[IR] Fix inplace bug of add_grad cpu kernel (PaddlePaddle#57042)

* enable inplace in dy2st with newir * fix bug --------- Co-authored-by: kangguangli <kangguangli@hotmail.com>
BeingGod · Sep 9, 2023 · 69a4d30 · 69a4d30
1 parent 5cdfdad
commit 69a4d30
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 6 deletions.
diff --git a/paddle/fluid/framework/CMakeLists.txt b/paddle/fluid/framework/CMakeLists.txt
@@ -1034,7 +1034,7 @@ cc_library(
 cc_library(
   executor_cache
   SRCS executor_cache.cc
-  DEPS parallel_executor standalone_executor phi_kernel_adaptor
+  DEPS parallel_executor standalone_executor phi_kernel_adaptor pd_inplace_pass
        pd_op_to_kernel_pass ir)
 if(WITH_PSCORE)
   get_property(RPC_DEPS GLOBAL PROPERTY RPC_DEPS)

diff --git a/paddle/fluid/framework/executor_cache.cc b/paddle/fluid/framework/executor_cache.cc
@@ -13,12 +13,16 @@
 // limitations under the License.
 
 #include "paddle/fluid/framework/executor_cache.h"
+
 #include "paddle/fluid/framework/new_executor/interpretercore.h"
 #include "paddle/fluid/framework/op_info.h"
+#include "paddle/fluid/ir/transforms/inplace_pass.h"
 #include "paddle/fluid/ir/transforms/pd_op_to_kernel_pass.h"
 #include "paddle/fluid/ir_adaptor/translator/translate.h"
 #include "paddle/ir/core/program.h"
 #include "paddle/ir/core/value.h"
+#include "paddle/ir/pass/pass.h"
+#include "paddle/ir/pass/pass_manager.h"
 
 namespace paddle {
 namespace framework {
@@ -440,6 +444,10 @@ std::unique_ptr<::ir::Program> ConstructFowardIrProgram(
 
   auto ir_res = paddle::dialect::PdOpLowerToKernelPass(program.get());
 
+  ::ir::PassManager pm(::ir::IrContext::Instance(), 3);
+  pm.AddPass(::ir::CreateInplacePass());
+  pm.Run(ir_res.get());
+
   return ir_res;
 }
 
@@ -513,6 +521,10 @@ std::unique_ptr<::ir::Program> ConstructBackwardIrProgram(
 
   auto res = paddle::dialect::PdOpLowerToKernelPass(program.get());
 
+  ::ir::PassManager pm(::ir::IrContext::Instance(), 3);
+  pm.AddPass(::ir::CreateInplacePass());
+  pm.Run(res.get());
+
   return res;
 }
 

diff --git a/paddle/fluid/ir/transforms/inplace_pass.cc b/paddle/fluid/ir/transforms/inplace_pass.cc
@@ -18,6 +18,7 @@
 #include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_type.h"
 #include "paddle/fluid/ir/dialect/paddle_dialect/trait/inplace.h"
 #include "paddle/fluid/ir/dialect/paddle_dialect/utils/op_yaml_info_parser.h"
+#include "paddle/fluid/ir/dialect/paddle_kernel_dialect/ir/kernel_attribute.h"
 #include "paddle/fluid/ir/dialect/paddle_kernel_dialect/ir/kernel_dialect.h"
 #include "paddle/fluid/ir/dialect/paddle_kernel_dialect/ir/kernel_type.h"
 #include "paddle/ir/core/builtin_op.h"
@@ -199,6 +200,20 @@ static std::unordered_map<ir::Operation*, std::string> GetInplaceOps(
         upper_op_attrs.at("op_name").dyn_cast<::ir::StrAttribute>().AsString();
     VLOG(6) << "analyse op: " << upper_op_name;
 
+    // NOTE(zhangbo): add_grad cpu kernel can't do inplace, for the reason shown
+    // in the function: CommonElementwiseBroadcastBackward
+    // (paddle/phi/kernels/funcs/elementwise_grad_base.h)
+    if ((upper_op_name == "pd.add_grad") &&
+        (upper_op_attrs.at("kernel_key")
+             .dyn_cast<paddle::dialect::KernelAttribute>()
+             .data()
+             .backend() == phi::Backend::CPU)) {
+      for (size_t i = 0; i < op->num_results(); ++i) {
+        visited_values.insert(op->result(i));
+      }
+      continue;
+    }
+
     if (upper_op_attrs.count("is_inplace") != 0 &&
         upper_op_attrs.at("is_inplace").dyn_cast<ir::BoolAttribute>().data()) {
       VLOG(6) << upper_op_name << " is already an inplace op.";
@@ -291,7 +306,7 @@ class InplacePass : public ir::Pass {
 
   void Run(ir::Operation* op) override {
     auto module_op = op->dyn_cast<ir::ModuleOp>();
-    IR_ENFORCE(module_op, "DcePass should run on module op.");
+    IR_ENFORCE(module_op, "InplacePass should run on module op.");
     auto* block = module_op.block();
 
     auto inplace_ops = details::GetInplaceOps(block);
@@ -330,4 +345,4 @@ std::unique_ptr<ir::Pass> CreateInplacePass() {
 
 }  // namespace ir
 
-REGISTER_PASS(inplace, InplacePass);
+REGISTER_IR_PASS(inplace, InplacePass);
diff --git a/paddle/ir/pass/pass_registry.h b/paddle/ir/pass/pass_registry.h
@@ -80,10 +80,10 @@ class PassRegistrar {
                 msg)
 
 // Register a new pass that can be applied on the IR.
-#define REGISTER_PASS(pass_type, pass_class)                                   \
+#define REGISTER_IR_PASS(pass_type, pass_class)                                \
   STATIC_ASSERT_PASS_GLOBAL_NAMESPACE(                                         \
       __reg_pass__##pass_type,                                                 \
-      "REGISTER_PASS must be called in global namespace");                     \
+      "REGISTER_IR_PASS must be called in global namespace");                  \
   static ::ir::PassRegistrar<pass_class> __pass_registrar_##pass_type##__(     \
       #pass_type);                                                             \
   int TouchPassRegistrar_##pass_type() {                                       \

diff --git a/paddle/ir/transforms/dead_code_elimination_pass.cc b/paddle/ir/transforms/dead_code_elimination_pass.cc
@@ -77,4 +77,4 @@ std::unique_ptr<Pass> CreateDeadCodeEliminationPass() {
 
 }  // namespace ir
 
-REGISTER_PASS(dead_code_elimination, DeadCodeEliminationPass);
+REGISTER_IR_PASS(dead_code_elimination, DeadCodeEliminationPass);