未验证 提交 c7302f96 编写于 作者: X xiongkun 提交者: GitHub

Add C++ EinsumOp which support 2 operands einsum. (#42105)

* full api fix

* when out is None, go old dygraph mode

* by static check

* first version: support 2-inputs forwards. TODO: 1. backward  2. BroadCast  3. MultiVariable

* time out -> 120
上级 e95838dd
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <string>
#include <vector>
#include "paddle/fluid/framework/infershape_utils.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/phi/core/ddim.h"
#include "paddle/phi/kernels/impl/einsum_impl.h"
namespace paddle {
namespace operators {
class EinsumOp : public framework::OperatorWithKernel {
public:
using framework::OperatorWithKernel::OperatorWithKernel;
};
class EinsumOpMaker : public framework::OpProtoAndCheckerMaker {
public:
void Make() override {
AddInput("Operands", "(TensorList), The input tensor of einsum op.")
.AsDuplicable();
AddOutput("Out", "(Tensor), The output tensor of einsum op.");
AddAttr<std::string>("equation",
"(string) A einsum equation. such as `ij,jk->ik`"
"There must have `->` and the number of operands in "
"equation must equals the `Operands` length.");
AddComment(R"DOC(
Einsum Operator.
This operator is used to perform einsum operation for given operands and equation.
)DOC");
}
};
class EinsumGradOp : public framework::OperatorWithKernel {
public:
using framework::OperatorWithKernel::OperatorWithKernel;
void InferShape(framework::InferShapeContext* ctx) const override {
auto x_name = "Operands";
auto x_grad_name = framework::GradVarName(x_name);
ctx->SetOutputsDim(x_grad_name, ctx->GetInputsDim(x_name));
ctx->ShareAllLoD(x_name, x_grad_name);
}
protected:
framework::OpKernelType GetExpectedKernelType(
const framework::ExecutionContext& ctx) const override {
auto dtype = OperatorWithKernel::IndicateVarDataType(
ctx, framework::GradVarName("Out"));
return framework::OpKernelType(dtype, ctx.GetPlace());
}
};
template <typename T>
class EinsumGradMaker : public framework::SingleGradOpMaker<T> {
public:
using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
void Apply(GradOpPtr<T> retv) const override {
retv->SetType("einsum_grad");
retv->SetInput("Operands", this->Input("Operands"));
retv->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
retv->SetAttrMap(this->Attrs());
retv->SetOutput(framework::GradVarName("Operands"),
this->InputGrad("Operands", false));
}
};
} // namespace operators
} // namespace paddle
namespace ops = paddle::operators;
DECLARE_INFER_SHAPE_FUNCTOR(einsum, EinsumInferShapeFunctor,
PD_INFER_META(phi::EinsumInferShape));
REGISTER_OPERATOR(einsum, ops::EinsumOp, ops::EinsumOpMaker,
EinsumInferShapeFunctor,
ops::EinsumGradMaker<paddle::framework::OpDesc>,
ops::EinsumGradMaker<paddle::imperative::OpBase>);
REGISTER_OPERATOR(einsum_grad, ops::EinsumGradOp);
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "paddle/phi/kernels/einsum_grad_kernel.h"
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/impl/einsum_grad_impl.h"
PD_REGISTER_KERNEL(
einsum_grad, CPU, ALL_LAYOUT, phi::EinsumGradKernel, float, double) {}
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "paddle/phi/kernels/einsum_kernel.h"
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/impl/einsum_impl.h"
PD_REGISTER_KERNEL(einsum, CPU, ALL_LAYOUT, phi::EinsumKernel, float, double) {}
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "paddle/phi/core/dense_tensor.h"
namespace phi {
template <typename T, typename Context>
void EinsumGradKernel(const Context& dev_ctx,
const std::vector<const DenseTensor*>& x,
const DenseTensor& out_grad,
const std::string& equation,
std::vector<DenseTensor*> x_grad);
} // namespace phi
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "paddle/phi/core/dense_tensor.h"
namespace phi {
template <typename T, typename Context>
void EinsumKernel(const Context& dev_ctx,
const std::vector<const DenseTensor*>& inputs,
const std::string& equation,
DenseTensor* out);
} // namespace phi
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "paddle/phi/kernels/einsum_kernel.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/impl/einsum_grad_impl.h"
PD_REGISTER_KERNEL(
einsum_grad, GPU, ALL_LAYOUT, phi::EinsumGradKernel, float, double) {}
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "paddle/phi/kernels/einsum_kernel.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/impl/einsum_impl.h"
PD_REGISTER_KERNEL(einsum, GPU, ALL_LAYOUT, phi::EinsumKernel, float, double) {}
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/impl/einsum_impl.h"
#include "paddle/phi/kernels/tile_kernel.h"
#include "paddle/utils/string/string_helper.h"
namespace phi {
template <typename T, typename Context>
DenseTensor PerformTileAndReduction(const Context& dev_ctx,
const LabelMap& label2type,
const LabelMap& label2shape,
const std::vector<int>& broadcast_dims,
const std::vector<int>& ellipsis_dims,
std::string op_label, // value pass
DenseTensor& t) { // NOLINT
ReplaceEllipsis(op_label);
DenseTensor ret;
std::vector<int> repeat_times;
std::vector<int> resize_dims;
std::vector<int> recover_shape;
for (int c : op_label) {
if (label2type[c] == LabelType::Reduction) {
// '.' can't be Reduction, so we don't deal '.' here.
repeat_times.push_back(label2shape[c]);
resize_dims.push_back(1);
recover_shape.push_back(label2shape[c]);
} else {
if (c != '.') {
resize_dims.push_back(label2shape[c]);
repeat_times.push_back(1);
recover_shape.push_back(label2shape[c]);
} else {
int n_dims = broadcast_dims.size();
resize_dims.insert(
resize_dims.end(), broadcast_dims.begin(), broadcast_dims.end());
recover_shape.insert(
recover_shape.end(), ellipsis_dims.begin(), ellipsis_dims.end());
while (n_dims--) repeat_times.push_back(1);
}
}
}
t.Resize(make_ddim(resize_dims));
DenseTensor after_tile;
TileKernel<T, Context>(dev_ctx, t, repeat_times, &after_tile);
size_t n_ellipsis_idx = op_label.find(".", 0);
if (n_ellipsis_idx != std::string::npos) {
// may be we need reduce. broadcast_dims is not equal to ellipsis dims.
std::vector<int64_t> to_reduce;
for (size_t i = 0; i < broadcast_dims.size() - ellipsis_dims.size(); ++i)
to_reduce.push_back(i + n_ellipsis_idx);
int new_offset =
n_ellipsis_idx + broadcast_dims.size() - ellipsis_dims.size();
for (size_t i = 0; i < ellipsis_dims.size(); ++i)
if (ellipsis_dims[i] == 1) to_reduce.push_back(i + new_offset);
VLOG(5) << "PermformTileAndReduction: reduce sum axis: "
<< paddle::string::join_strings(to_reduce, ",");
if (to_reduce.size() != 0) {
ret = Sum<T, Context>(dev_ctx,
after_tile,
to_reduce,
after_tile.dtype(),
false); // not keep dim.
} else {
ret = after_tile;
}
} else {
ret = after_tile;
}
VLOG(5) << "PermformTileAndReduction: recover shape: "
<< paddle::string::join_strings(recover_shape, ",");
ret.Resize(make_ddim(recover_shape));
return ret;
}
template <typename T, typename Context>
void EinsumGradKernel(const Context& dev_ctx,
const std::vector<const DenseTensor*>& x,
const DenseTensor& out_grad,
const std::string& equation,
std::vector<DenseTensor*> x_grad) {
VLOG(5) << "Start EisumGradKernel:";
LabelMap labelshape(0);
LabelMap labeltype(LabelType::Reduction);
std::vector<LabelMap> label2perms(x.size(), LabelMap(-1));
std::vector<char> all_labels; // order: ABO, AO, BO, AB, Reduce
std::vector<std::vector<int>> ellipsis_dims(2);
std::vector<int> broadcast_dims;
std::vector<int> output_dims;
std::vector<DDim> input_dims;
for (auto& i : x) {
input_dims.push_back(i->dims());
}
std::string right;
ParseEinsumEquation(equation,
input_dims,
&labelshape,
&labeltype,
&all_labels,
&label2perms,
&ellipsis_dims,
&broadcast_dims,
&output_dims,
&right);
auto gather_labels_except_reduction = [&labeltype](std::string all) {
std::string res("");
for (auto c : all)
if (labeltype[static_cast<int>(c)] != LabelType::Reduction) res += c;
return res;
};
if (x.size() == 1) { // Unary
auto splits = paddle::string::split_string(equation, "->");
auto left = splits[0];
right = splits[1].substr(1);
auto new_equation = right + "->" + gather_labels_except_reduction(left);
auto new_operands = std::vector<const DenseTensor*>();
new_operands.push_back(&out_grad);
DenseTensor before_tile;
EinsumKernel<T, Context>(dev_ctx, new_operands, new_equation, &before_tile);
*(x_grad[0]) = PerformTileAndReduction<T, Context>(dev_ctx,
labeltype,
labelshape,
broadcast_dims,
ellipsis_dims[0],
left,
before_tile);
} else {
auto splits = paddle::string::split_string(equation, "->");
auto left = splits[0];
auto ops = paddle::string::split_string(left, ",");
right = splits[1].substr(1);
auto equation_for_A =
right + "," + ops[1] + "->" + gather_labels_except_reduction(ops[0]);
auto equation_for_B =
right + "," + ops[0] + "->" + gather_labels_except_reduction(ops[1]);
auto operands_for_A = std::vector<const DenseTensor*>();
auto operands_for_B = std::vector<const DenseTensor*>();
DenseTensor dA, dB;
operands_for_A.push_back(&out_grad);
operands_for_A.push_back(x[1]);
operands_for_B.push_back(&out_grad);
operands_for_B.push_back(x[0]);
DenseTensor before_tile;
EinsumKernel<T, Context>(dev_ctx, operands_for_A, equation_for_A, &dA);
EinsumKernel<T, Context>(dev_ctx, operands_for_B, equation_for_B, &dB);
*(x_grad[0]) = PerformTileAndReduction<T, Context>(dev_ctx,
labeltype,
labelshape,
broadcast_dims,
ellipsis_dims[0],
ops[0],
dA);
*(x_grad[1]) = PerformTileAndReduction<T, Context>(dev_ctx,
labeltype,
labelshape,
broadcast_dims,
ellipsis_dims[1],
ops[1],
dB);
}
}
} // namespace phi
此差异已折叠。
/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
#include "paddle/phi/core/compat/op_utils.h"
namespace phi {
KernelSignature EinsumOpArgumentMapping(const ArgumentMappingContext& ctx) {
return KernelSignature("einsum", {"Operands"}, {"equation"}, {"Out"});
}
KernelSignature EinsumGradOpArgumentMapping(const ArgumentMappingContext& ctx) {
return KernelSignature("einsum_grad",
{"Operands", {"Out@GRAD"}},
{"equation"},
{{"Operands@GRAD"}});
}
} // namespace phi
PD_REGISTER_ARG_MAPPING_FN(einsum, phi::EinsumOpArgumentMapping);
PD_REGISTER_ARG_MAPPING_FN(einsum_grad, phi::EinsumGradOpArgumentMapping);
......@@ -1065,6 +1065,7 @@ set_tests_properties(test_lstm_cudnn_op PROPERTIES TIMEOUT 120)
set_tests_properties(test_stack_op PROPERTIES TIMEOUT 120)
set_tests_properties(test_bilinear_interp_v2_op PROPERTIES TIMEOUT 120)
set_tests_properties(test_svd_op PROPERTIES TIMEOUT 80)
set_tests_properties(test_einsum_op PROPERTIES TIMEOUT 120)
set_tests_properties(test_qr_op PROPERTIES TIMEOUT 60)
set_tests_properties(test_deformable_psroi_pooling PROPERTIES TIMEOUT 120)
set_tests_properties(test_trilinear_interp_v2_op PROPERTIES TIMEOUT 120)
......
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import print_function
import unittest
import numpy as np
import paddle
from op_test import OpTest
class TestEinsumBinary(OpTest):
def setUp(self):
paddle.enable_static()
self.op_type = "einsum"
self.disable = False
self.set_mandatory()
self.init_input()
np.random.seed(123)
out = np.einsum(self.equation, *self.inputs)
self.operands = []
for idx, inp in enumerate(self.inputs):
self.operands.append(("x" + str(idx), inp))
self.inputs = {"Operands": self.operands}
self.attrs = {"equation": self.equation}
self.outputs = {'Out': out}
def init_input(self):
self.inputs = []
for t, s in zip(self.types, self.shapes):
self.inputs.append(np.random.random(s).astype(t))
def set_mandatory(self):
self.disable = False
self.shapes = [(10, 10, 20), (20, 6)]
self.types = [np.float64, np.float64]
self.equation = "mij,jk->ki"
def test_check_output(self):
if not self.disable:
self.check_output()
def test_grad(self):
if not self.disable:
self.check_grad([op[0] for op in self.operands], ["Out"])
class TestEinsum1(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(20, 3, 3), (20, 3, 3)]
self.types = [np.float64, np.float64]
self.equation = "mij,mjk->mik"
class TestEinsum2(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(20, 3, 3), (20, 3, 3)]
self.types = [np.float64, np.float64]
self.equation = "mij,mjk->ikm"
class TestEinsum3(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(10, 10), (10, 10)]
self.types = [np.float64, np.float64]
self.equation = "ij,jk->ik" # }}}
class TestEinsumWithReduction(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(10, 3, 5), (5, 30)]
self.types = [np.float64, np.float64]
self.equation = "ijk,kl->jl"
class TestEinsumWithReduction1(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(10, 3, 3, 5), (10, 5, 10, 10)]
self.types = [np.float64, np.float64]
self.equation = "mijk,mklh->ljm"
class TestEinsumWithUnary(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(10, 10, 3, 5)]
self.types = [np.float64]
self.equation = "mijk->mi"
class TestEinsumWithUnary1(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(5, 10, 3, 3), (3, 6, 3, 10)]
self.types = [np.float64, np.float64]
self.equation = "imjl,jklm->imk"
class TestEinsumWithBroadcast1(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(5, 10, 3, 3)]
self.types = [np.float64]
self.equation = "i...->..."
class TestEinsumWithBroadcast2(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(10, 11), (3, 4, 5, 10)]
self.types = [np.float64, np.float64]
self.equation = "...ij,...i->j..."
class TestEinsumWithBroadcast3(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(10, 3, 2, 3, 4), (12, 10)]
self.types = [np.float64, np.float64]
self.equation = "k...,...jk->...k"
class TestEinsumWithBroadcast4(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(10, 3, 2, 3, 4), (12, 10)]
self.types = [np.float64, np.float64]
self.equation = "a...d,...cb->...abcd"
class TestEinsumWithBroadcast5(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(3, 2, 2, 10), (10, 3, 2, 2)]
self.types = [np.float64, np.float64]
self.equation = "...a,a...->..."
class TestEinsumWithBroadcast6(TestEinsumBinary):
def set_mandatory(self):
self.shapes = [(100), (100)]
self.types = [np.float64, np.float64]
self.equation = "i,i->"
if __name__ == "__main__":
unittest.main()
......@@ -20,6 +20,10 @@ from .linalg import dot, matmul, transpose
from .manipulation import squeeze, unsqueeze, reshape
from .math import multiply
from .math import sum as paddle_sum
from ..fluid.framework import _in_legacy_dygraph
from paddle import _C_ops
from ..fluid.data_feeder import check_variable_and_dtype, check_type, check_dtype
from ..fluid.layer_helper import LayerHelper
from paddle.common_ops_import import dygraph_only
......@@ -660,6 +664,26 @@ def plan_einsum(operands, g_view, g_shape, g_supports, g_count, n_bcast):
return plan
def einsum_v2(equation, *operands):
if _in_legacy_dygraph():
# dygraph
return _C_ops.einsum(operands, 'equation', equation)
# static graph
for inp in operands:
check_variable_and_dtype(inp, 'dtype', ['float32', 'float64'], 'einsum')
check_type(equation, 'equation', str, 'einsum')
helper = LayerHelper('einsum', **locals())
out = helper.create_variable_for_type_inference(dtype=operands[0].dtype)
attrs = dict()
attrs['equation'] = equation
helper.append_op(
type='einsum',
inputs={'Operands': operands},
outputs={'Out': out},
attrs=attrs, )
return out
def einsum(equation, *operands):
r"""
einsum(equation, *operands)
......@@ -817,6 +841,9 @@ def einsum(equation, *operands):
# [0.50226176, 0.24512935, 0.39881429],
# [0.51476848, 0.23367381, 0.39229113]]])
"""
import os
if int(os.environ.get('FLAGS_new_einsum', "0")):
return einsum_v2(equation, *operands)
nop = len(operands)
assert nop > 0, "At least one operand is expected."
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册