make mechanism of caching grad op symbol shape to work (#65591)

zyfncg · web-flow · commit 90b2cc2fa6d2 · 2024-07-01T22:12:10.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/cache_grad_op_symbol_shape_gen.py b/paddle/fluid/pir/dialect/op_generator/cache_grad_op_symbol_shape_gen.py
@@ -111,10 +111,10 @@ def gen_cpp_file_code(self, cpp_file_path):
                 continue
             if op_info_item.backward_name not in self.op_info_maps:
                 continue
+            if op_info_item.kernel_map is None:
+                continue
 
             grad_op_item = self.op_info_maps[op_info_item.backward_name]
-            if grad_op_item.infer_meta_map is None:
-                continue
 
             for op_phi_name in op_info_item.op_phi_name:
                 create_grad_op_shape_info_code = ""
@@ -143,20 +143,6 @@ def gen_cpp_file_code(self, cpp_file_path):
                                 index=index,
                             )
                         )
-                    elif input_name in op_info_item.mutable_attribute_name_list:
-                        # mutable attribute
-                        index = len(
-                            op_info_item.input_name_list
-                        ) + op_info_item.mutable_attribute_name_list.index(
-                            input_name
-                        )
-                        create_grad_op_shape_info_code += (
-                            GET_INPUT_SHAPE_CODE_TEMPLATE.format(
-                                input_name=input_name,
-                                name_suffix=SHAPE_VAR_NAME_SUFFIX,
-                                index=index,
-                            )
-                        )
                     elif input_name.endswith("_grad"):
                         # output grad
                         origin_out_name = input_name[:-5]
@@ -171,13 +157,35 @@ def gen_cpp_file_code(self, cpp_file_path):
                             )
                         )
                     else:
-                        raise (
+                        raise ValueError(
                             f"Not found input name {input_name} for backward op {op_info_item.backward_name}."
                         )
+                # mutable attribute
+                for (
+                    mutable_attribute_name
+                ) in grad_op_item.mutable_attribute_name_list:
+                    assert (
+                        mutable_attribute_name
+                        in op_info_item.mutable_attribute_name_list
+                    ), f"{mutable_attribute_name} is not found in {op_info_item.backward_name}'s mutable_attribute name list."
+                    index = len(
+                        op_info_item.input_name_list
+                    ) + op_info_item.mutable_attribute_name_list.index(
+                        mutable_attribute_name
+                    )
+                    create_grad_op_shape_info_code += (
+                        GET_INPUT_SHAPE_CODE_TEMPLATE.format(
+                            input_name=mutable_attribute_name,
+                            name_suffix=SHAPE_VAR_NAME_SUFFIX,
+                            index=index,
+                        )
+                    )
 
                 create_grad_op_output_shape_code = ""
                 for output_name in grad_op_item.output_name_list:
-                    assert output_name.endswith("_grad")
+                    if not output_name.endswith("_grad"):
+                        create_grad_op_output_shape_code = ""
+                        break
                     origin_input_name = output_name[:-5]
                     if (
                         origin_input_name
@@ -216,7 +224,10 @@ def gen_cpp_file_code(self, cpp_file_path):
                     input_shape_list=", ".join(
                         [
                             input_name + SHAPE_VAR_NAME_SUFFIX
-                            for input_name in grad_op_item.input_name_list
+                            for input_name in (
+                                grad_op_item.input_name_list
+                                + grad_op_item.mutable_attribute_name_list
+                            )
                         ]
                     ),
                     create_grad_op_output_shape_code=create_grad_op_output_shape_code,
@@ -228,6 +239,21 @@ def gen_cpp_file_code(self, cpp_file_path):
                     ),
                 )
 
+                if len(op_info_item.kernel_map['func']) == 1:
+                    continue
+                for kernel_func_name in op_info_item.kernel_map['func']:
+                    is_inplace_version = op_phi_name.endswith('_')
+                    op_origin_name = (
+                        op_phi_name[:-1] if is_inplace_version else op_phi_name
+                    )
+                    if kernel_func_name == op_origin_name:
+                        continue
+                    inplace_suffix = '_' if is_inplace_version else ''
+                    body_code += UNIMPLEMENTED_CODE_TEMPLATE.format(
+                        op_name=to_pascal_case(kernel_func_name)
+                        + inplace_suffix
+                    )
+
         directory_path = os.path.dirname(cpp_file_path)
         if not os.path.exists(directory_path):
             os.makedirs(directory_path, exist_ok=True)
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -91,6 +91,7 @@
 #include "paddle/fluid/pir/dialect/operator/interface/decomp.h"
 #include "paddle/fluid/pir/dialect/operator/interface/decomp_vjp.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_symbolic_shape.h"
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infermeta.h"
 #include "paddle/fluid/pir/dialect/operator/interface/layout_transformation.h"
 #include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
@@ -1301,6 +1302,14 @@ def AutoCodeGen(
         op_traits = op_info.traits_list
         op_interfaces = op_info.interfaces_list
         op_interfaces += ["paddle::dialect::OpYamlInfoInterface"]
+        if (
+            dialect_name == "pd_op"
+            and op_info.backward_name
+            and not op_info.is_sparse_op
+        ):
+            op_interfaces += [
+                "paddle::dialect::CacheGradOpSymbolicShapeInterface"
+            ]
         exclusive_interface_str = gen_exclusive_interface_str(
             op_info, op_info_items
         )
diff --git a/paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h b/paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h
@@ -28,16 +28,16 @@ class CacheGradOpSymbolicShapeInterface
  public:
   /// Defined these methods with the interface.
   struct Concept {
-    explicit Concept(bool (*cache_grad_op_symbolic_shape)(
+    explicit Concept(void (*cache_grad_op_symbolic_shape)(
         pir::Operation *op, pir::InferSymbolicShapeContext *infer_context))
         : cache_grad_op_symbolic_shape(cache_grad_op_symbolic_shape) {}
-    bool (*cache_grad_op_symbolic_shape)(
+    void (*cache_grad_op_symbolic_shape)(
         pir::Operation *op, pir::InferSymbolicShapeContext *infer_context);
   };
 
   template <class ConcreteOp>
   struct Model : public Concept {
-    static inline bool CacheGradOpSymbolicShape(
+    static inline void CacheGradOpSymbolicShape(
         pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
       return op->dyn_cast<ConcreteOp>().CacheGradOpSymbolicShape(infer_context);
     }
@@ -50,7 +50,7 @@ class CacheGradOpSymbolicShapeInterface
       : pir::OpInterfaceBase<CacheGradOpSymbolicShapeInterface>(op),
         impl_(impl) {}
 
-  bool CacheGradOpSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
+  void CacheGradOpSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
 
  private:
   Concept *impl_;
diff --git a/paddle/pir/src/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.cc b/paddle/pir/src/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.cc
@@ -22,9 +22,9 @@
 
 namespace pir {
 
-bool CacheGradOpSymbolicShapeInterface::CacheGradOpSymbolicShape(
+void CacheGradOpSymbolicShapeInterface::CacheGradOpSymbolicShape(
     pir::InferSymbolicShapeContext *infer_context) {
-  return impl_->cache_grad_op_symbolic_shape(operation(), infer_context);
+  impl_->cache_grad_op_symbolic_shape(operation(), infer_context);
 }
 
 }  // namespace pir
diff --git a/paddle/pir/src/dialect/shape/transforms/shape_optimization_pass.cc b/paddle/pir/src/dialect/shape/transforms/shape_optimization_pass.cc
@@ -295,6 +295,8 @@ void InferSymExprForOp(Operation* op,
         }
       } else {
         // risk set
+        LOG(WARNING) << "Not found symbolic shape cache for " << op->name()
+                     << "[id:" << op->id() << "]";
         for (uint32_t i = 0; i < op->num_results(); ++i) {
           infer_context->SetSymbolForValueByStaticShape(op->result(i));
         }
@@ -349,12 +351,8 @@ void CacheBackwardOpSymbolicShape(Operation* op,
       op->dyn_cast<pir::CacheGradOpSymbolicShapeInterface>();
   if (cache_grad_op_symbolic_shape_interface) {
     VLOG(3) << "CacheBackwardOpSymbolicShape for: " << op->name();
-    PADDLE_ENFORCE_EQ(
-        cache_grad_op_symbolic_shape_interface.CacheGradOpSymbolicShape(
-            infer_context),
-        true,
-        common::errors::Fatal("CacheBackwardOpSymbolicShape for %s failed.",
-                              op->name()));
+    cache_grad_op_symbolic_shape_interface.CacheGradOpSymbolicShape(
+        infer_context);
   }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -22,9 +22,9 @@`
`22`	`22`
`23`	`23`	`namespace pir {`
`24`	`24`
`25`		`-bool CacheGradOpSymbolicShapeInterface::CacheGradOpSymbolicShape(`
	`25`	`+void CacheGradOpSymbolicShapeInterface::CacheGradOpSymbolicShape(`
`26`	`26`	`pir::InferSymbolicShapeContext *infer_context) {`
`27`		`- return impl_->cache_grad_op_symbolic_shape(operation(), infer_context);`
	`27`	`+ impl_->cache_grad_op_symbolic_shape(operation(), infer_context);`
`28`	`28`	`}`
`29`	`29`
`30`	`30`	`} // namespace pir`
Original file line number	Diff line number	Diff line change
`@@ -295,6 +295,8 @@ void InferSymExprForOp(Operation* op,`
`295`	`295`	`}`
`296`	`296`	`} else {`
`297`	`297`	`// risk set`
	`298`	`+ LOG(WARNING) << "Not found symbolic shape cache for " << op->name()`
	`299`	`+ << "[id:" << op->id() << "]";`
`298`	`300`	`for (uint32_t i = 0; i < op->num_results(); ++i) {`
`299`	`301`	`infer_context->SetSymbolForValueByStaticShape(op->result(i));`
`300`	`302`	`}`
`@@ -349,12 +351,8 @@ void CacheBackwardOpSymbolicShape(Operation* op,`
`349`	`351`	`op->dyn_cast<pir::CacheGradOpSymbolicShapeInterface>();`
`350`	`352`	`if (cache_grad_op_symbolic_shape_interface) {`
`351`	`353`	`VLOG(3) << "CacheBackwardOpSymbolicShape for: " << op->name();`
`352`		`- PADDLE_ENFORCE_EQ(`
`353`		`- cache_grad_op_symbolic_shape_interface.CacheGradOpSymbolicShape(`
`354`		`- infer_context),`
`355`		`- true,`
`356`		`- common::errors::Fatal("CacheBackwardOpSymbolicShape for %s failed.",`
`357`		`- op->name()));`
	`354`	`+ cache_grad_op_symbolic_shape_interface.CacheGradOpSymbolicShape(`
	`355`	`+ infer_context);`
`358`	`356`	`}`
`359`	`357`	`}`
`360`	`358`