[PRIM][IR]Support prim in new ir

paddle/fluid/framework/type_info.cc

@@ -17,7 +17,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/raw_tensor.h"
 #include "paddle/fluid/framework/string_array.h"
 #include "paddle/fluid/prim/utils/static/desc_tensor.h"
-#include "paddle/fluid/primitive/type/static_tensor.h"
+#include "paddle/fluid/primitive/type/lazy_tensor.h"
 
 namespace phi {
 
@@ -41,8 +41,7 @@ template class TypeInfoTraits<phi::TensorBase, paddle::framework::Strings>;
 template class TypeInfoTraits<phi::TensorBase, paddle::framework::FeedList>;
 template class TypeInfoTraits<phi::TensorBase, egr::VariableCompatTensor>;
 template class TypeInfoTraits<phi::TensorBase, paddle::prim::DescTensor>;
-template class TypeInfoTraits<phi::TensorBase,
-                              paddle::primitive::experimental::StaticTensor>;
+template class TypeInfoTraits<phi::TensorBase, paddle::primitive::LazyTensor>;
 template class TypeInfoTraits<phi::TensorBase,
                               paddle::framework::VariableRefArray>;
 

paddle/fluid/ir/dialect/op_generator/api_gen.py

@@ -95,6 +95,8 @@
     'expand',
     'tile',
     'add_grad',
+    'divide_grad',
+    'sum_grad',
 ]
 OP_RESULT = 'ir::OpResult'
 VECTOR_TYPE = 'ir::VectorType'

paddle/fluid/ir/dialect/op_generator/op_gen.py

@@ -21,10 +21,13 @@
     gen_exclusive_interface_str,
     gen_op_infer_meta_str,
     gen_op_vjp_str,
-    vjp_interface_gen_op_list,
 )
 from op_member_func_gen import gen_op_get_inputs_outputs_str
 from op_verify_gen import gen_verify_func_str
+from vjp_interface_gen_op_list import (
+    vjp_interface_declare_gen_op_list,
+    vjp_interface_implementation_gen_op_list,
+)
 
 # =====================================
 # String Template for h file code gen
@@ -112,7 +115,7 @@ class {op_name} : public ir::Op<{op_name}{interfaces}{traits}> {{
 #include "paddle/phi/infermeta/backward.h"
 #include "paddle/phi/api/lib/utils/allocator.h"
 #include "paddle/fluid/primitive/rule/vjp/vjp.h"
-#include "paddle/fluid/primitive/type/static_tensor.h"
+#include "paddle/fluid/primitive/type/lazy_tensor.h"
 #include "paddle/ir/core/op_base.h"
 
 {input}
@@ -756,7 +759,7 @@ def OpGenerator(
 
         if (
             op_info.backward_name
-            and op_info.op_phi_name[0] in vjp_interface_gen_op_list
+            and op_info.op_phi_name[0] in vjp_interface_declare_gen_op_list
         ):
             op_interfaces += ["VjpInterface"]
         exclusive_interface_str = gen_exclusive_interface_str(op_info)
@@ -1055,7 +1058,8 @@ def OpGenerator(
             # TODO(chenzhiyang) add vjp gen code
             if (
                 op_info.backward_name
-                and op_info.op_phi_name[0] in vjp_interface_gen_op_list
+                and op_info.op_phi_name[0]
+                in vjp_interface_implementation_gen_op_list
             ):
                 op_vjp_str = gen_op_vjp_str(
                     op_class_name,

paddle/fluid/ir/dialect/op_generator/op_interface_gen.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # generator interfaces
-from vjp_interface_gen_op_list import vjp_interface_gen_op_list
+from vjp_interface_gen_op_list import vjp_interface_declare_gen_op_list
 
 OP_INFER_SHAPE_TEMPLATE = """
 void {op_name}::InferMeta( phi::InferMetaContext *infer_meta ) {{
@@ -23,13 +23,13 @@
 """
 
 OP_VJP_FORWARD_INPUT_OR_OUTPUT_TEMPLATE = """
-    {input_type} {input_name}(std::make_shared<primitive::experimental::StaticTensor>(op_obj.{input_name}()));"""
+    {input_type} {input_name}(std::make_shared<primitive::LazyTensor>(op_obj.{input_name}()));"""
 
 OP_VJP_FORWARD_OUTPUT_GRAD_TEMPLATE = """
-    Tensor {output_grad_name}(std::make_shared<primitive::experimental::StaticTensor>(out_grads[{idx1}][{idx2}]));"""
+    Tensor {output_grad_name}(std::make_shared<primitive::LazyTensor>(out_grads[{idx1}][{idx2}]));"""
 
 OP_VJP_FORWARD_OUTPUT_GRAD_LIST_TEMPLATE = """
-    std::vector<Tensor> {output_grad_name}(std::make_shared<primitive::experimental::StaticTensor>(out_grads[{idx1}]));"""
+    std::vector<Tensor> {output_grad_name}(std::make_shared<primitive::LazyTensor>(out_grads[{idx1}]));"""
 
 OP_VJP_ATTRIBUTE_TEMPLATE = """
     {attr_type} {attr_name} = op->attribute("{attr_name}").dyn_cast<{attr_parse_type}>().data();"""
@@ -39,7 +39,7 @@
 
 
 OP_VJP_CALL_VJP_TEMPLATE = """  std::vector<std::vector<Tensor>> tensor_res =
-    primitive::experimental::{op_phi_name}_vjp(
+    primitive::{op_phi_name}_vjp(
     {inputs_list}stop_gradients);"""
 
 OP_VJP_STOPGRADIENT_TEMPLATE = """
@@ -48,7 +48,7 @@
         res[i].resize(tensor_res[i].size());
         for (size_t j = 0; j < tensor_res[i].size(); ++j) {{
             if(tensor_res[i][j].defined()){{
-                res[i][j] = std::static_pointer_cast<primitive::experimental::StaticTensor>(tensor_res[i][j].impl())->getValue().dyn_cast<ir::OpResult>();
+                res[i][j] = std::static_pointer_cast<primitive::LazyTensor>(tensor_res[i][j].impl())->getValue().dyn_cast<ir::OpResult>();
             }}
         }}
     }}"""
@@ -166,6 +166,6 @@ def gen_exclusive_interface_str(op_info):
         exclusive_interface_str += (
             "  static void InferMeta( phi::InferMetaContext *infer_meta );"
         )
-    if op_info.op_phi_name[0] in vjp_interface_gen_op_list:
+    if op_info.op_phi_name[0] in vjp_interface_declare_gen_op_list:
         exclusive_interface_str += "\n  static std::vector<std::vector<ir::OpResult>> Vjp(ir::Operation* op, const std::vector<std::vector<ir::OpResult>>& out_grads, const std::vector<std::vector<bool>>& stop_gradients);"
     return exclusive_interface_str

paddle/fluid/ir/dialect/op_generator/vjp_interface_gen_op_list.py

@@ -21,4 +21,5 @@
 # TODO(wanghao107)
 # remove this file and support Vjp methods
 # code gen.
-vjp_interface_gen_op_list = ["tanh", "mean", "add"]
+vjp_interface_declare_gen_op_list = ["tanh", "mean", "divide", "sum", "add"]
+vjp_interface_implementation_gen_op_list = ["tanh", "mean", "divide", "add"]

paddle/fluid/ir/dialect/pd_op_vjp_manual.cc

@@ -15,13 +15,40 @@
 #include "paddle/fluid/ir/dialect/pd_attribute.h"
 #include "paddle/fluid/ir/dialect/pd_op.h"
 #include "paddle/fluid/primitive/rule/vjp/vjp.h"
-#include "paddle/fluid/primitive/type/static_tensor.h"
+#include "paddle/fluid/primitive/type/lazy_tensor.h"
 #include "paddle/ir/core/op_base.h"
 #include "paddle/phi/common/int_array.h"
 
 // TODO(wanghao107)
 // this file will be generated in pd_op.cc
 
 namespace paddle {
-namespace dialect {}  // namespace dialect
+namespace dialect {
+std::vector<std::vector<ir::OpResult>> SumOp::Vjp(
+    ir::Operation* op,
+    const std::vector<std::vector<ir::OpResult>>& out_grads,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  SumOp op_obj = op->dyn_cast<SumOp>();
+  Tensor x(std::make_shared<primitive::LazyTensor>(op_obj.x()));
+  Tensor out_grad(std::make_shared<primitive::LazyTensor>(out_grads[0][0]));
+
+  IntArray axis = op_obj.axis()
+                      .GetDefiningOp()
+                      ->attribute("value")
+                      .dyn_cast<paddle::dialect::IntArrayAttribute>()
+                      .data();
+  bool keepdim = op->attribute("keepdim").dyn_cast<ir::BoolAttribute>().data();
+  bool reduce_all = false;
+  std::vector<std::vector<Tensor>> tensor_res = primitive::sum_vjp(
+      x, out_grad, axis, keepdim, reduce_all, stop_gradients);
+  std::vector<std::vector<ir::OpResult>> res(1, std::vector<ir::OpResult>(1));
+  if (tensor_res[0][0].defined()) {
+    res[0][0] =
+        std::static_pointer_cast<primitive::LazyTensor>(tensor_res[0][0].impl())
+            ->getValue()
+            .dyn_cast<ir::OpResult>();
+  }
+  return res;
+}
+}  // namespace dialect
 }  // namespace paddle

paddle/fluid/primitive/CMakeLists.txt

@@ -1,2 +1,3 @@
+add_subdirectory(utils)
 add_subdirectory(backend)
 add_subdirectory(rule)

paddle/fluid/primitive/backend/eager_backend.cc

@@ -19,8 +19,6 @@
 
 namespace paddle {
 namespace primitive {
-namespace backend {
-namespace experimental {}  // namespace experimental
-}  // namespace backend
+namespace backend {}  // namespace backend
 }  // namespace primitive
 }  // namespace paddle

paddle/fluid/primitive/backend/eager_backend.h

@@ -21,8 +21,6 @@
 
 namespace paddle {
 namespace primitive {
-namespace backend {
-namespace experimental {}  // namespace experimental
-}  // namespace backend
+namespace backend {}  // namespace backend
 }  // namespace primitive
 }  // namespace paddle