[PIR & Inference] constant_folding_pass support place (#59430)

yuanlehome · web-flow · commit 363795bd2d34 · 2023-11-30T00:51:47.000+08:00
* constant_folding_pass support place

* update

* add ut

* refine pattern_rewrite_test

* fix win ci
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -740,6 +740,9 @@ static void DisablePrepareDataOpt(
 }
 
 bool AnalysisPredictor::PrepareExecutor() {
+  PADDLE_ENFORCE_NOT_NULL(sub_scope_,
+                          platform::errors::PreconditionNotMet(
+                              "The sub_scope should not be nullptr."));
   if (config_.dist_config().use_dist_model()) {
 #if defined(PADDLE_WITH_DISTRIBUTE) && defined(PADDLE_WITH_PSCORE)
     VLOG(3) << "use_dist_model is enabled, will init FleetExecutor.";
@@ -770,14 +773,26 @@ bool AnalysisPredictor::PrepareExecutor() {
           paddle::TranslateLegacyProgramToProgram(*inference_program_));
 
       ::pir::PassManager pm_for_op_program(::pir::IrContext::Instance(), 2);
+      //----------------------------------------------------------------------------------------------//
+      // Operator fusion pass
       pm_for_op_program.AddPass(::pir::CreateConv2dFusePass());
+      //----------------------------------------------------------------------------------------------//
+
+      //----------------------------------------------------------------------------------------------//
+      // Functional pass
+      //----------------------------------------------------------------------------------------------//
 
-      pm_for_op_program.AddPass(::pir::CreateConstantFoldingPass(sub_scope_));
+      //----------------------------------------------------------------------------------------------//
+      // Basic pass required by the framework
+      pm_for_op_program.AddPass(
+          ::pir::CreateParamsSyncAmongDevicesPass(place_, sub_scope_));
+      pm_for_op_program.AddPass(
+          ::pir::CreateConstantFoldingPass(place_, sub_scope_));
       pm_for_op_program.AddPass(::pir::CreateDeadCodeEliminationPass());
       pm_for_op_program.AddPass(
           ::pir::CreateReplaceFetchWithShadowOutputPass());
-      pm_for_op_program.AddPass(
-          ::pir::CreateParamsSyncAmongDevicesPass(place_, sub_scope_));
+      //----------------------------------------------------------------------------------------------//
+
       // pm_for_op_program.EnableIRPrinting();
       pm_for_op_program.Run(pir_program_.get());
 
@@ -807,10 +822,6 @@ bool AnalysisPredictor::PrepareExecutor() {
     executor_->MakeReusePlan(reuse_table);
   }
 
-  PADDLE_ENFORCE_NOT_NULL(sub_scope_,
-                          platform::errors::PreconditionNotMet(
-                              "The sub_scope should not be nullptr."));
-
   return true;
 }
 
diff --git a/paddle/fluid/pir/transforms/constant_folding_pass.cc b/paddle/fluid/pir/transforms/constant_folding_pass.cc
@@ -21,7 +21,7 @@
 
 #include "paddle/fluid/framework/new_executor/interpretercore.h"
 #include "paddle/fluid/framework/scope.h"
-// #include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
+#include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
@@ -51,28 +51,47 @@ class ConstantFoldingPattern : public pir::RewritePattern {
  public:
   ConstantFoldingPattern(
       pir::IrContext* context,
-      size_t* suffix,
+      size_t* counter,
+      const phi::Place& place,
       paddle::framework::Scope* scope,
       paddle::framework::interpreter::ExecutionConfig* exe_config,
       std::vector<std::string>* deleted_vars)
       : RewritePattern(MatchAnyOpTypeTag(),
                        1 /*benefit*/,
                        context,
                        {} /*generated_names*/),
-        counter_(suffix),
+        counter_(counter),
+        place_(place),
         scope_(scope),
         exe_config_(exe_config),
         deleted_vars_(deleted_vars) {
     exe_config_->create_local_scope = false;
   }
 
   bool Match(pir::Operation* op) const override {
+    // 1. Some ops do not need to be processed
     if (op->HasTrait<pir::SideEffectTrait>() ||
         op->isa<pir::ConstantTensorOp>() || op->isa<pir::ParameterOp>() ||
         op->isa<paddle::dialect::FeedOp>())
       return false;
-    if (!ValidOp(op)) {
-      return false;
+
+    for (uint32_t i = 0; i < op->num_operands(); i++) {
+      // 2. inputs must come from parameter op or constant op
+      // 3. inputs must be a dense tensor type
+      if (!(pir::GetDefiningOpForInput(op, i)->isa<pir::ParameterOp>() ||
+            pir::GetDefiningOpForInput(op, i)->isa<pir::ConstantTensorOp>()) ||
+          !op->operand_source(i)
+               .type()
+               .isa<paddle::dialect::DenseTensorType>()) {
+        return false;
+      }
+    }
+
+    for (uint32_t i = 0; i < op->num_results(); i++) {
+      // 4. outputs must be a dense tensor type
+      if (!op->result(i).type().isa<paddle::dialect::DenseTensorType>()) {
+        return false;
+      }
     }
 
     return true;
@@ -87,9 +106,9 @@ class ConstantFoldingPattern : public pir::RewritePattern {
     // execute program
     exe_config_->skip_gc_vars.insert(output_var_name);
     auto kernel_program =
-        paddle::dialect::PdOpLowerToKernelPass(&new_program, phi::CPUPlace{});
+        paddle::dialect::PdOpLowerToKernelPass(&new_program, place_);
     paddle::framework::InterpreterCore core(
-        phi::CPUPlace{}, {}, kernel_program->block(), scope_, *exe_config_);
+        place_, {}, kernel_program->block(), scope_, *exe_config_);
 
     core.Run({});
 
@@ -144,50 +163,6 @@ class ConstantFoldingPattern : public pir::RewritePattern {
     return true;
   }
 
-  bool ValidOp(pir::Operation* op) const {
-    for (uint32_t i = 0; i < op->num_operands(); i++) {
-      // 1. inputs must come from parameter op or constant op
-      // 2. inputs must be a dense tensor type
-      if (!(pir::GetDefiningOpForInput(op, i)->isa<pir::ParameterOp>() ||
-            pir::GetDefiningOpForInput(op, i)->isa<pir::ConstantTensorOp>()) ||
-          !op->operand_source(i)
-               .type()
-               .isa<paddle::dialect::DenseTensorType>()) {
-        return false;
-      }
-    }
-
-    for (uint32_t i = 0; i < op->num_results(); i++) {
-      // 3. outputs must be a dense tensor type
-      if (!op->result(i).type().isa<paddle::dialect::DenseTensorType>()) {
-        return false;
-      }
-      // TODO(liuyuanle): Uncomment follows doce after have a solution to slove
-      // the performance degradation
-      // auto use_ops = pir::GetUseOpsForOutput(op, i);
-      // for (auto* use_op : use_ops) {
-      //   if (use_op->isa<pir::CombineOp>()) {
-      //     return false;
-      //   }
-      //   if (use_op->HasInterface<paddle::dialect::OpYamlInfoInterface>()) {
-      //     auto [input_infos, _1, _2, _3, _4] =
-      //         use_op->dyn_cast<paddle::dialect::OpYamlInfoInterface>()
-      //             .GetOpInfo();
-      //     for (const auto& input_info : input_infos) {
-      //       if (input_info.type_name.find("IntArrayAttribute") !=
-      //               std::string::npos ||
-      //           input_info.type_name.find("ScalarAttribute") !=
-      //               std::string::npos) {
-      //         return false;
-      //       }
-      //     }
-      //   }
-      // }
-    }
-
-    return true;
-  }
-
   std::string BuildProgramFromOperation(pir::Operation* op,
                                         pir::Program* new_program) const {
     pir::Builder builder = pir::Builder(ir_context(), new_program->block());
@@ -233,16 +208,18 @@ class ConstantFoldingPattern : public pir::RewritePattern {
   }
 
  private:
-  size_t* counter_{nullptr};
-  paddle::framework::Scope* scope_{nullptr};
-  paddle::framework::interpreter::ExecutionConfig* exe_config_{nullptr};
-  std::vector<std::string>* deleted_vars_{nullptr};
+  size_t* counter_;
+  phi::Place place_;
+  paddle::framework::Scope* scope_;
+  paddle::framework::interpreter::ExecutionConfig* exe_config_;
+  std::vector<std::string>* deleted_vars_;
 };
 
 class ConstantFoldingPass : public pir::Pass {
  public:
-  explicit ConstantFoldingPass(paddle::framework::Scope* scope)
-      : pir::Pass("constant_folding_pass", 1), scope_(scope) {
+  explicit ConstantFoldingPass(const phi::Place& place,
+                               paddle::framework::Scope* scope)
+      : pir::Pass("constant_folding_pass", 1), place_(place), scope_(scope) {
     PADDLE_ENFORCE_NOT_NULL(
         scope_, phi::errors::InvalidArgument("scope can not be nullptr"));
   }
@@ -251,7 +228,7 @@ class ConstantFoldingPass : public pir::Pass {
   bool Initialize(pir::IrContext* context) override {
     pir::RewritePatternSet ps(context);
     ps.Add<ConstantFoldingPattern>(
-        context, &counter_, scope_, &exe_config_, &deleted_vars_);
+        context, &counter_, place_, scope_, &exe_config_, &deleted_vars_);
     patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
     return true;
   }
@@ -276,6 +253,7 @@ class ConstantFoldingPass : public pir::Pass {
 
  private:
   size_t counter_{0};
+  phi::Place place_;
   paddle::framework::Scope* scope_{nullptr};
   paddle::framework::interpreter::ExecutionConfig exe_config_{};
   std::vector<std::string> deleted_vars_;
@@ -288,8 +266,8 @@ class ConstantFoldingPass : public pir::Pass {
 namespace pir {
 
 std::unique_ptr<Pass> CreateConstantFoldingPass(
-    paddle::framework::Scope* scope) {
-  return std::make_unique<ConstantFoldingPass>(scope);
+    const phi::Place& place, paddle::framework::Scope* scope) {
+  return std::make_unique<ConstantFoldingPass>(place, scope);
 }
 
 }  // namespace pir
diff --git a/paddle/fluid/pir/transforms/constant_folding_pass.h b/paddle/fluid/pir/transforms/constant_folding_pass.h
@@ -15,6 +15,7 @@
 #pragma once
 
 #include <memory>
+#include "paddle/phi/common/place.h"
 #include "paddle/pir/core/dll_decl.h"
 
 namespace paddle {
@@ -28,6 +29,6 @@ namespace pir {
 class Pass;
 
 IR_API std::unique_ptr<Pass> CreateConstantFoldingPass(
-    paddle::framework::Scope* scope);
+    const phi::Place& place, paddle::framework::Scope* scope);
 
 }  // namespace pir
diff --git a/test/cpp/pir/pattern_rewrite/CMakeLists.txt b/test/cpp/pir/pattern_rewrite/CMakeLists.txt
@@ -1,12 +1,7 @@
-set(PATTERN_REWRITE_TEST_DEPS pir_transforms gtest op_dialect_vjp pir)
-
-if(WITH_DISTRIBUTE)
-  set(PATTERN_REWRITE_TEST_DEPS ${PATTERN_REWRITE_TEST_DEPS} fleet_executor
-                                parallel_executor)
-endif()
-
-cc_test_old(pattern_rewrite_test SRCS pattern_rewrite_test.cc DEPS
-            ${PATTERN_REWRITE_TEST_DEPS})
+cc_test(
+  pattern_rewrite_test
+  SRCS pattern_rewrite_test.cc
+  DEPS gtest op_dialect_vjp pir pir_transforms)
 
 cc_test(
   drr_test
@@ -32,6 +27,3 @@ cc_test(
   drr_attention_fuse_test
   SRCS drr_attention_fuse_test.cc
   DEPS pir_transforms drr gtest op_dialect_vjp pir)
-
-set_tests_properties(pattern_rewrite_test
-                     PROPERTIES ENVIRONMENT "FLAGS_enable_pir_in_executor=true")
diff --git a/test/cpp/pir/pattern_rewrite/drr_attention_fuse_test.cc b/test/cpp/pir/pattern_rewrite/drr_attention_fuse_test.cc
@@ -17,9 +17,12 @@
 #include <cstdint>
 #include <vector>
 
+#include "paddle/fluid/framework/scope.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/transforms/constant_folding_pass.h"
 #include "paddle/fluid/pir/transforms/fusion/attention_fuse_pass.h"
+
 #include "paddle/pir/core/builtin_dialect.h"
 #include "paddle/pir/pass/pass_manager.h"
 
@@ -142,6 +145,8 @@ TEST(DrrTest, AttentionFuse) {
 
   pir::PassManager pm(ctx);
   pm.AddPass(pir::CreateAttentionFusePass());
+  paddle::framework::Scope scope;
+  pm.AddPass(pir::CreateConstantFoldingPass(phi::CPUPlace{}, &scope));
   pm.EnableIRPrinting();
 
   CHECK_EQ(pm.Run(&program), true);
diff --git a/test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc b/test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc
