【Prim】Add gather vjp (#50305)

* tmp gather vjp * support gather * remove useless code * fix compiling error * fix ut * add eager test * add eager test * add seed * fix cpu error * fix transpose op compat * remove tensor index case * fix prim_cinn * fix ut
PaddlePaddle · Feb 22, 2023 · 4db8e5c · 4db8e5c
1 parent 613a3ff
commit 4db8e5c
Show file tree

Hide file tree

Showing 11 changed files with 836 additions and 28 deletions.
diff --git a/.gitignore b/.gitignore
@@ -26,7 +26,8 @@ paddle/phi/api/lib/tensor_operants.cc
 paddle/phi/extension.h
 paddle/phi/include/*
 paddle/phi/infermeta/generated.*
-
+paddle/fluid/prim/api/generated_prim/*.cc
+paddle/fluid/prim/api/generated_prim/*.h
 *.DS_Store
 *.vs
 build/

diff --git a/paddle/fluid/operators/elementwise/elementwise_add_op.cc b/paddle/fluid/operators/elementwise/elementwise_add_op.cc
@@ -61,10 +61,10 @@ class ElementwiseAddCompositeGradOpMaker
     paddle::experimental::Tensor y = this->GetSingleForwardInput("Y");
     paddle::experimental::Tensor out_grad = this->GetSingleOutputGrad("Out");
     paddle::experimental::Tensor dx = this->GetSingleInputGrad("X");
-    auto dx_ptr = this->GetOutputPtr(&dx);
+    auto* dx_ptr = this->GetOutputPtr(&dx);
     std::string dx_name = this->GetOutputName(dx);
     paddle::experimental::Tensor dy = this->GetSingleInputGrad("Y");
-    auto dy_ptr = this->GetOutputPtr(&dy);
+    auto* dy_ptr = this->GetOutputPtr(&dy);
     std::string dy_name = this->GetOutputName(dy);
     int axis = static_cast<int>(this->Attr<int>("axis"));
     VLOG(6) << "Runing add_grad composite func";

diff --git a/paddle/fluid/operators/gather_op.cc b/paddle/fluid/operators/gather_op.cc
@@ -19,6 +19,8 @@ limitations under the License. */
 #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"
@@ -132,6 +134,33 @@ class GatherGradOpMaker : public framework::SingleGradOpMaker<T> {
   }
 };
 
+class GatherCompositeGradOpMaker : public prim::CompositeGradOpMakerBase {
+ public:
+  using prim::CompositeGradOpMakerBase::CompositeGradOpMakerBase;
+
+ protected:
+  void Apply() override {
+    paddle::experimental::Tensor index = this->GetSingleForwardInput("Index");
+    paddle::optional<paddle::experimental::Tensor> tensor_axis =
+        this->GetOptionalSingleForwardInput("Axis");
+    paddle::experimental::Tensor x = this->GetSingleForwardInput("X");
+    paddle::experimental::Tensor dout = this->GetSingleOutputGrad("Out");
+    paddle::experimental::Tensor dx = this->GetSingleInputGrad("X");
+    auto* dx_ptr = this->GetOutputPtr(&dx);
+    std::string dx_name = this->GetOutputName(*dx_ptr);
+    int axis = static_cast<int>(this->Attr<int>("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<prim::DescTensor>(x, index, dout, axis, false, dx_ptr);
+    }
+    this->RecoverOutputName(dx, dx_name);
+  }
+};
+
 DECLARE_NO_NEED_BUFFER_VARS_INFERER(GatherGradNoNeedBufferVarInferer, "X");
 
 }  // namespace operators
@@ -146,6 +175,7 @@ REGISTER_OPERATOR(gather,
                   ops::GatherOpMaker,
                   ops::GatherGradOpMaker<paddle::framework::OpDesc>,
                   ops::GatherGradOpMaker<paddle::imperative::OpBase>,
+                  ops::GatherCompositeGradOpMaker,
                   GatherInferShapeFunctor);
 DECLARE_INFER_SHAPE_FUNCTOR(gather_grad,
                             GatherGradInferShapeFunctor,

diff --git a/paddle/fluid/prim/api/api.yaml b/paddle/fluid/prim/api/api.yaml
@@ -23,4 +23,5 @@
 - scatter
 - scatter_nd_add
 - tile
+- transpose
 - subtract