From ffeac6d5e809d86a1761b00781a3ac68a38554f3 Mon Sep 17 00:00:00 2001 From: cyberslack_lee Date: Mon, 26 Jun 2023 10:51:48 +0800 Subject: [PATCH] support auto generation for gather (#54084) --- paddle/fluid/operators/gather_op.cc | 191 ------------------ paddle/fluid/operators/unity_build_rule.cmake | 1 - paddle/phi/api/yaml/backward.yaml | 13 ++ paddle/phi/api/yaml/legacy_backward.yaml | 13 -- paddle/phi/api/yaml/legacy_ops.yaml | 10 - paddle/phi/api/yaml/op_compat.yaml | 8 + paddle/phi/api/yaml/op_version.yaml | 7 + paddle/phi/api/yaml/ops.yaml | 10 + paddle/phi/ops/compat/gather_sig.cc | 40 ---- 9 files changed, 38 insertions(+), 255 deletions(-) delete mode 100644 paddle/fluid/operators/gather_op.cc delete mode 100644 paddle/phi/ops/compat/gather_sig.cc diff --git a/paddle/fluid/operators/gather_op.cc b/paddle/fluid/operators/gather_op.cc deleted file mode 100644 index c01f2d2d528..00000000000 --- a/paddle/fluid/operators/gather_op.cc +++ /dev/null @@ -1,191 +0,0 @@ -/* Copyright (c) 2016 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 -#include -#include - -#include "paddle/fluid/framework/infershape_utils.h" -#include "paddle/fluid/framework/op_registry.h" -#include "paddle/fluid/framework/op_version_registry.h" -#include "paddle/fluid/prim/api/composite_backward/composite_backward_api.h" -#include "paddle/fluid/prim/utils/static/composite_grad_desc_maker.h" -#include "paddle/phi/core/ddim.h" -#include "paddle/phi/core/infermeta_utils.h" -#include "paddle/phi/infermeta/backward.h" -#include "paddle/phi/infermeta/binary.h" - -namespace paddle { -namespace operators { - -class GatherOp : public framework::OperatorWithKernel { - public: - using framework::OperatorWithKernel::OperatorWithKernel; - - protected: - phi::KernelKey GetExpectedKernelType( - const framework::ExecutionContext& ctx) const override { - return phi::KernelKey(OperatorWithKernel::IndicateVarDataType(ctx, "X"), - ctx.device_context().GetPlace()); - } - phi::KernelKey GetKernelTypeForVar( - const std::string& var_name, - const phi::DenseTensor& tensor, - const phi::KernelKey& expected_kernel_type) const override { - if (var_name == "Axis") { - return phi::KernelKey(phi::Backend::ALL_BACKEND, - expected_kernel_type.layout(), - expected_kernel_type.dtype()); - } - return phi::KernelKey( - tensor.place(), tensor.layout(), expected_kernel_type.dtype()); - } -}; - -class GatherGradOp : public framework::OperatorWithKernel { - public: - using framework::OperatorWithKernel::OperatorWithKernel; - - protected: - phi::KernelKey GetExpectedKernelType( - const framework::ExecutionContext& ctx) const override { - return phi::KernelKey(OperatorWithKernel::IndicateVarDataType( - ctx, framework::GradVarName("Out")), - ctx.device_context().GetPlace()); - } - phi::KernelKey GetKernelTypeForVar( - const std::string& var_name, - const phi::DenseTensor& tensor, - const phi::KernelKey& expected_kernel_type) const override { - if (var_name == "Axis") { - return phi::KernelKey(phi::Backend::ALL_BACKEND, - expected_kernel_type.layout(), - expected_kernel_type.dtype()); - } - return phi::KernelKey( - tensor.place(), tensor.layout(), expected_kernel_type.dtype()); - } -}; - -class GatherOpMaker : public framework::OpProtoAndCheckerMaker { - public: - void Make() override { - AddInput("X", "The source input of gather op"); - AddInput("Index", "The index input of gather op"); - AddInput("Axis", - "The Tensor which contains the axis that we do gather operation.") - .AsDispensable(); - AddOutput("Out", "The output of gather op"); - AddAttr( - "axis", - "The Tensor which contains the axis that we do gather operation.") - .SetDefault(0); - AddComment(R"DOC( -Gather Operator. - -$Out = X[Index]$ - -Out is obtained by gathering entries of the outer-most dimension -of X indexed by Index and concatenate them together. - -Example: - -X = [[1, 2], - [3, 4], - [5, 6]] - -Index = [[1, 2]] - -Then: - -Out = [[3, 4], - [5, 6]] - -)DOC"); - } -}; - -template -class GatherGradOpMaker : public framework::SingleGradOpMaker { - public: - using framework::SingleGradOpMaker::SingleGradOpMaker; - - protected: - void Apply(GradOpPtr op) const override { - op->SetType("gather_grad"); - op->SetInput("Index", this->Input("Index")); - op->SetInput("Axis", this->Input("Axis")); - - op->SetInput("X", this->Input("X")); - op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out")); - op->SetOutput(framework::GradVarName("X"), this->InputGrad("X")); - op->SetAttrMap(this->Attrs()); - } -}; - -class GatherCompositeGradOpMaker : public prim::CompositeGradOpMakerBase { - public: - using prim::CompositeGradOpMakerBase::CompositeGradOpMakerBase; - - protected: - void Apply() override { - paddle::Tensor index = this->GetSingleForwardInput("Index"); - paddle::optional tensor_axis = - this->GetOptionalSingleForwardInput("Axis"); - paddle::Tensor x = this->GetSingleForwardInput("X"); - paddle::Tensor dout = this->GetSingleOutputGrad("Out"); - paddle::Tensor dx = this->GetSingleInputGrad("X"); - auto* dx_ptr = this->GetOutputPtr(&dx); - std::string dx_name = this->GetOutputName(*dx_ptr); - int axis = static_cast(this->Attr("axis")); - VLOG(3) << "Runing gather_grad composite func"; - if (tensor_axis.is_initialized()) { - PADDLE_THROW(platform::errors::Unimplemented( - "We don't support dynamic index from tensor for gather composite " - "grad for now. ")); - } else { - prim::gather_grad(x, index, dout, axis, dx_ptr); - } - this->RecoverOutputName(dx, dx_name); - } -}; - -DECLARE_NO_NEED_BUFFER_VARS_INFERER(GatherGradNoNeedBufferVarInferer, "X"); - -} // namespace operators -} // namespace paddle - -namespace ops = paddle::operators; -DECLARE_INFER_SHAPE_FUNCTOR(gather, - GatherInferShapeFunctor, - PD_INFER_META(phi::GatherInferMeta)); -REGISTER_OPERATOR(gather, - ops::GatherOp, - ops::GatherOpMaker, - ops::GatherGradOpMaker, - ops::GatherGradOpMaker, - ops::GatherCompositeGradOpMaker, - GatherInferShapeFunctor); -DECLARE_INFER_SHAPE_FUNCTOR(gather_grad, - GatherGradInferShapeFunctor, - PD_INFER_META(phi::GeneralUnaryGradInferMeta)); -REGISTER_OPERATOR(gather_grad, - ops::GatherGradOp, - ops::GatherGradNoNeedBufferVarInferer, - GatherGradInferShapeFunctor); - -REGISTER_OP_VERSION(gather).AddCheckpoint( - R"ROC(upgrad gather, add a new input [Axis])ROC", - paddle::framework::compatible::OpVersionDesc().NewInput( - "Axis", "Specify the axis of gather operation.")); diff --git a/paddle/fluid/operators/unity_build_rule.cmake b/paddle/fluid/operators/unity_build_rule.cmake index 4de9ed86fd0..3cfb53ea14b 100644 --- a/paddle/fluid/operators/unity_build_rule.cmake +++ b/paddle/fluid/operators/unity_build_rule.cmake @@ -104,7 +104,6 @@ register_unity_group( flatten_op.cc fsp_op.cc gather_nd_op.cc - gather_op.cc gather_tree_op.cc gaussian_random_batch_size_like_op.cc mkldnn/gaussian_random_mkldnn_op.cc diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml index 43dc112ca79..18fc6cbe07d 100644 --- a/paddle/phi/api/yaml/backward.yaml +++ b/paddle/phi/api/yaml/backward.yaml @@ -840,6 +840,19 @@ kernel : func : frame_grad +- backward_op : gather_grad + forward : gather(Tensor x, Tensor index, Scalar axis=0) -> Tensor(out) + args : (Tensor x, Tensor index, Tensor out_grad, Scalar axis=0) + output : Tensor(x_grad) + infer_meta : + func : GeneralUnaryGradInferMeta + param: [x] + kernel : + data_type: out_grad + func : gather_grad + composite : gather_grad(x, index, out_grad, axis, x_grad) + no_need_buffer : x + - backward_op : gather_nd_grad forward : gather_nd (Tensor x, Tensor index) -> Tensor(out) args : (Tensor x, Tensor index, Tensor out_grad) diff --git a/paddle/phi/api/yaml/legacy_backward.yaml b/paddle/phi/api/yaml/legacy_backward.yaml index 307ead7b599..df644df098a 100755 --- a/paddle/phi/api/yaml/legacy_backward.yaml +++ b/paddle/phi/api/yaml/legacy_backward.yaml @@ -302,19 +302,6 @@ kernel : func : frobenius_norm_grad -- backward_op : gather_grad - forward : gather(Tensor x, Tensor index, Scalar axis=0) -> Tensor(out) - args : (Tensor x, Tensor index, Tensor out_grad, Scalar axis=0) - output : Tensor(x_grad) - infer_meta : - func : UnchangedInferMeta - param: [x] - kernel : - data_type: x - func : gather_grad - composite : gather_grad(x, index, out_grad, axis, x_grad) - no_need_buffer : x - - backward_op : hardswish_grad forward : hardswish (Tensor x) -> Tensor(out) args : (Tensor x, Tensor out_grad) diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml index e5d1351758b..fccf4fdfb65 100755 --- a/paddle/phi/api/yaml/legacy_ops.yaml +++ b/paddle/phi/api/yaml/legacy_ops.yaml @@ -422,16 +422,6 @@ optional : skip_update, master_params inplace : (params -> params_out), (moments1 -> moments1_out), (moments2 -> moments2_out), (beta1_pows -> beta1_pows_out), (beta2_pows -> beta2_pows_out), (master_params -> master_params_out) -- op : gather - args : (Tensor x, Tensor index, Scalar(int) axis=0) - output : Tensor(out) - infer_meta : - func : GatherInferMeta - kernel : - func : gather - data_type: x - backward : gather_grad - - op : gaussian args : (IntArray shape, float mean, float std, int seed, DataType dtype, Place place={}) output: Tensor(out) diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml index f2894e1d967..4a554ad336a 100755 --- a/paddle/phi/api/yaml/op_compat.yaml +++ b/paddle/phi/api/yaml/op_compat.yaml @@ -1132,6 +1132,14 @@ - op : gather backward : gather_grad + inputs : + {x : X, index : Index} + outputs : + out : Out + scalar : + axis : + data_type : int + tensor_name : Axis - op : gather_nd backward : gather_nd_grad diff --git a/paddle/phi/api/yaml/op_version.yaml b/paddle/phi/api/yaml/op_version.yaml index 8894a7bac14..48e13add28e 100644 --- a/paddle/phi/api/yaml/op_version.yaml +++ b/paddle/phi/api/yaml/op_version.yaml @@ -197,6 +197,13 @@ - delete_attr : dims comment : The attr 'dims' is deleted. +- op : gather + version : + - checkpoint : Upgrade gather, add a new input [Axis] + action : + - add_input : Axis + comment : Specify the axis of gather operation. + - op : gaussian_random version : - checkpoint : Upgrade gaussian_random add new inputs [ShapeTensor] and [ShapeTensorList] diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml index a0dd85dc3f5..e433f9e6e3e 100644 --- a/paddle/phi/api/yaml/ops.yaml +++ b/paddle/phi/api/yaml/ops.yaml @@ -940,6 +940,16 @@ data_type : dtype backend : place +- op : gather + args : (Tensor x, Tensor index, Scalar axis=0) + output : Tensor(out) + infer_meta : + func : GatherInferMeta + kernel : + func : gather + data_type: x + backward : gather_grad + - op : gather_nd args : (Tensor x, Tensor index) output : Tensor diff --git a/paddle/phi/ops/compat/gather_sig.cc b/paddle/phi/ops/compat/gather_sig.cc deleted file mode 100644 index 8618b5a226f..00000000000 --- a/paddle/phi/ops/compat/gather_sig.cc +++ /dev/null @@ -1,40 +0,0 @@ -/* 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 GatherOpArgumentMapping(const ArgumentMappingContext& ctx) { - if (ctx.HasInput("Axis")) { - return KernelSignature("gather", {"X", "Index"}, {"Axis"}, {"Out"}); - } else { - return KernelSignature("gather", {"X", "Index"}, {"axis"}, {"Out"}); - } -} - -KernelSignature GatherGradOpArgumentMapping(const ArgumentMappingContext& ctx) { - if (ctx.HasInput("Axis")) { - return KernelSignature( - "gather_grad", {"X", "Index", "Out@GRAD"}, {"Axis"}, {"X@GRAD"}); - } else { - return KernelSignature( - "gather_grad", {"X", "Index", "Out@GRAD"}, {"axis"}, {"X@GRAD"}); - } -} - -} // namespace phi - -PD_REGISTER_ARG_MAPPING_FN(gather, phi::GatherOpArgumentMapping); -PD_REGISTER_ARG_MAPPING_FN(gather_grad, phi::GatherGradOpArgumentMapping); -- GitLab