[CINN] CINN plug-in Inference Executor (#61949)

jiahy0825 · web-flow · commit de1f8ab5dafd · 2024-02-27T19:50:53.000+08:00
* [CINN] CINN plug-in Inference Executor

* Fix headerfile

* Fix compile dependency

* Change cinn_pass_entry position

* cinn_cc_library -&gt; cc_library

* Fix compile error

* Resolve dependency

* Fix sub_graph_checker

* Change cmake

* Fix compile error

* Refine inference-related codes

* Refine inference-related codes

* Refine inference-related codes

* Refine inference-related codes

* Fix compile error

* Add doc

* remove preprocess macro

* Remove doc
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/CMakeLists.txt b/paddle/cinn/hlir/dialect/operator/transforms/CMakeLists.txt
@@ -14,4 +14,8 @@ if(NOT CINN_ONLY)
   cinn_cc_library(cinn_transforms SRCS ${cinn_transforms_srcs} DEPS
                   ${cinn_transforms_deps})
 
+  cc_library(
+    add_cinn_pass
+    SRCS add_cinn_pass.cc
+    DEPS op_dialect pir cinn_op_dialect cinnapi pir_transforms cinn_transforms)
 endif()
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
@@ -0,0 +1,137 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h"
+
+#include "paddle/common/errors.h"
+#include "paddle/common/flags.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/pir/include/core/ir_context.h"
+#include "paddle/pir/include/core/program.h"
+#include "paddle/pir/include/dialect/shape/ir/shape_dialect.h"
+#include "paddle/pir/include/pass/pass_manager.h"
+
+#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/dynamic_reshape_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/check_infer_symbolic_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_0d_to_1d_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_dynamic_to_static_dim_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_static_dim_to_dynamic_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/divide_group_op_to_fusion_op_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/move_generate_shape_ops_to_prologue_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/simplify_dim_expr_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/substitute_dim_expr_based_on_constraints_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/insert_broadcast_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/remove_unchanged_reshape_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/replace_dynamic_expand_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.h"
+#include "paddle/fluid/pir/transforms/build_cinn_pass.h"
+#include "paddle/fluid/pir/transforms/dead_code_elimination_pass.h"
+#include "paddle/fluid/pir/transforms/shape_optimization_pass.h"
+
+COMMON_DECLARE_bool(print_ir);
+COMMON_DECLARE_bool(check_infer_symbolic);
+
+namespace cinn::dialect::ir {
+
+namespace {
+bool HasDynamicShape(const pir::Program &program) {
+  for (const auto &op : *program.block()) {
+    if (op.isa<pir::CombineOp>()) {
+      continue;
+    }
+    for (uint32_t i = 0; i < op.num_results(); ++i) {
+      if (op.result(i) && op.result(i).type()) {
+        auto shape_type =
+            op.result(i).type().dyn_cast<pir::ShapedTypeInterface>();
+        if (shape_type && shape_type.IsDynamicShape()) {
+          return true;
+        }
+      }
+    }
+  }
+  return false;
+}
+}  // namespace
+
+void AddCinnPass(std::shared_ptr<pir::PassManager> &pass_manager,  // NOLINT
+                 pir::Program &program) {                          // NOLINT
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
+
+  bool has_dynamic_shape = HasDynamicShape(program);
+
+  if (FLAGS_print_ir) {
+    pass_manager->EnableIRPrinting();
+  }
+
+  pass_manager->AddPass(cinn::dialect::ir::CreateConvert0DTo1DPass());
+  if (!has_dynamic_shape && FLAGS_check_infer_symbolic) {
+    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
+    pass_manager->AddPass(cinn::dialect::ir::CreateCheckInferSymbolicPass());
+  }
+  pass_manager->AddPass(cinn::dialect::ir::CreatePdOpToCinnOpPass());
+  pass_manager->AddPass(cinn::dialect::ir::CreateRemoveUnchangedReshapePass());
+  pass_manager->AddPass(
+      cinn::dialect::ir::CreateAddBroadcastToElementwisePass());
+  pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
+
+  if (has_dynamic_shape) {
+    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
+    pass_manager->AddPass(cinn::dialect::ir::CreateSimplifyDimExprPass());
+    pass_manager->AddPass(
+        cinn::dialect::ir::CreateSubstituteDimExprBasedOnConstraintsPass());
+    pass_manager->AddPass(cinn::dialect::ir::CreateInsertBroadcastPass());
+    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
+    pass_manager->AddPass(
+        cinn::dialect::ir::CreateFuseShapeOpsIntoGenerateShapeOpPass());
+    pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
+    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
+  }
+
+  pass_manager->AddPass(pir::CreateBuildCinnPass());
+
+  pass_manager->AddPass(
+      cinn::dialect::ir::CreateMoveGenerateShapeOpsToProloguePass());
+  pass_manager->AddPass(cinn::dialect::ir::CreateDivideGroupOpToFusionOpPass());
+  pass_manager->AddPass(cinn::dialect::ir::CreateDynamicReshapeOpPass());
+  pass_manager->AddPass(cinn::dialect::ir::CreateReplaceDynamicExpandOpPass());
+  pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
+
+  bool force_static_shape = false;
+  if (auto pass = cinn::dialect::ir::CreateConvertDynamicToStaticDimPass()) {
+    pass_manager->AddPass(std::move(pass.value()));
+    force_static_shape = true;
+  }
+  if (auto pass = cinn::dialect::ir::CreateConvertStaticDimToDynamicPass()) {
+    pass_manager->AddPass(std::move(pass.value()));
+  }
+
+  if (has_dynamic_shape && !force_static_shape) {
+    pass_manager->AddPass(
+        cinn::dialect::ir::CreateLowerCinnDyShapeFusionOpPass());
+  }
+  pass_manager->AddPass(cinn::dialect::ir::CreateLowerCinnFusionOpPass());
+  pass_manager->AddPass(
+      cinn::dialect::ir::CreateSplitGenerateShapeIntoShapeOpsPass());
+}
+
+}  // namespace cinn::dialect::ir
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h b/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h
@@ -0,0 +1,31 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+
+namespace pir {
+
+class PassManager;
+class Program;
+
+}  // namespace pir
+
+namespace cinn::dialect::ir {
+
+void AddCinnPass(std::shared_ptr<pir::PassManager> &pass_manager,  // NOLINT
+                 pir::Program &program);                           // NOLINT
+
+}  // namespace cinn::dialect::ir
diff --git a/paddle/fluid/inference/api/CMakeLists.txt b/paddle/fluid/inference/api/CMakeLists.txt
@@ -77,33 +77,29 @@ if(WITH_GPU AND TENSORRT_FOUND)
   set(inference_deps ${inference_deps} tensorrt_engine tensorrt_converter)
 endif()
 
+set(ANALYSIS_PREDICTOR_SRCS analysis_predictor.cc resource_manager.cc
+                            infer_context.cc ${mkldnn_quantizer_src})
+set(ANALYSIS_PREDICTOR_DEPS
+    ${inference_deps}
+    zero_copy_tensor
+    ir_pass_manager
+    op_compatible_info
+    infer_io_utils
+    model_utils
+    fleet_executor)
+
 if(WITH_ONNXRUNTIME)
-  cc_library(
-    analysis_predictor
-    SRCS analysis_predictor.cc onnxruntime_predictor.cc resource_manager.cc
-         infer_context.cc ${mkldnn_quantizer_src}
-    DEPS ${inference_deps}
-         zero_copy_tensor
-         ir_pass_manager
-         op_compatible_info
-         infer_io_utils
-         model_utils
-         onnxruntime
-         paddle2onnx
-         fleet_executor)
-else()
-  cc_library(
-    analysis_predictor
-    SRCS analysis_predictor.cc resource_manager.cc infer_context.cc
-         ${mkldnn_quantizer_src}
-    DEPS ${inference_deps}
-         zero_copy_tensor
-         ir_pass_manager
-         op_compatible_info
-         infer_io_utils
-         model_utils
-         fleet_executor)
+  set(ANALYSIS_PREDICTOR_SRCS ${ANALYSIS_PREDICTOR_SRCS}
+                              onnxruntime_predictor.cc)
+  set(ANALYSIS_PREDICTOR_DEPS ${ANALYSIS_PREDICTOR_DEPS} onnxruntime
+                              paddle2onnx)
+elseif(WITH_CINN)
+  set(ANALYSIS_PREDICTOR_DEPS ${ANALYSIS_PREDICTOR_DEPS} add_cinn_pass)
 endif()
+cc_library(
+  analysis_predictor
+  SRCS ${ANALYSIS_PREDICTOR_SRCS}
+  DEPS ${ANALYSIS_PREDICTOR_DEPS})
 
 if(WITH_ONNXRUNTIME AND WIN32)
   # Copy onnxruntime for some c++ test in Windows, since the test will
diff --git a/paddle/fluid/inference/api/analysis_config.cc b/paddle/fluid/inference/api/analysis_config.cc
@@ -535,7 +535,7 @@ AnalysisConfig::AnalysisConfig(const AnalysisConfig &other) {
   CP_MEMBER(with_profile_);
 
   // cinn compiler related.
-  CP_MEMBER(use_cinn_compiler_);
+  CP_MEMBER(use_cinn_);
 
   // glog related.
   CP_MEMBER(with_glog_info_);
@@ -606,7 +606,7 @@ AnalysisConfig::AnalysisConfig(const AnalysisConfig &other) {
 #undef CP_MEMBER
 
   Update();
-  if (use_tensorrt_ || use_cinn_compiler_) {
+  if (use_tensorrt_ || use_cinn_) {
     // Update() will reset all the passes, when some tensorRT pass is deleted in
     // other.pass_builder(), it will set again, so we just remove the
     // deleted_pass.
@@ -987,7 +987,7 @@ void AnalysisConfig::Update() {
   }
 
   // TODO(wilber): An ugly method to update pass, need to be fixed.
-  if (use_cinn_compiler_) {
+  if (use_cinn_) {
     pass_builder()->ClearPasses();
     for (const auto &pass : kCINNCompilerPasses) {
       pass_builder()->AppendPass(pass);
@@ -1479,7 +1479,7 @@ std::string AnalysisConfig::Summary() {
   }
 
   // cinn compiler
-  os.InsertRow({"use_cinn_compiler", use_cinn_compiler_ ? "true" : "false"});
+  os.InsertRow({"use_cinn_compiler", use_cinn_ ? "true" : "false"});
 
   // ir info
   os.InsertRow(
@@ -1603,9 +1603,9 @@ void AnalysisConfig::Exp_EnableMixedPrecisionOps(
   mixed_white_list_ = white_list;
 }
 
-void AnalysisConfig::Exp_EnableCINNCompiler() {
+void AnalysisConfig::EnableCINN() {
 #ifdef PADDLE_WITH_CINN
-  use_cinn_compiler_ = true;
+  use_cinn_ = true;
   Update();
 #else
   PADDLE_THROW(platform::errors::Unavailable(
@@ -1614,8 +1614,6 @@ void AnalysisConfig::Exp_EnableCINNCompiler() {
 #endif
 }
 
-bool AnalysisConfig::cinn_compiler_enabled() const {
-  return use_cinn_compiler_;
-}
+bool AnalysisConfig::cinn_enabled() const { return use_cinn_; }
 
 }  // namespace paddle
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -105,6 +105,10 @@
 #include "paddle/fluid/platform/device/gpu/cuda/cuda_profiler.h"
 #endif
 
+#ifdef PADDLE_WITH_CINN
+#include "paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h"
+#endif
+
 #include "paddle/common/flags.h"
 #include "paddle/fluid/ir_adaptor/translator/translate.h"
 #include "paddle/fluid/pir/transforms/constant_folding_pass.h"
@@ -874,6 +878,21 @@ bool AnalysisPredictor::PrepareExecutor() {
         DecompProgram decomp_object(pir_program_.get());
         decomp_object.decomp_program();
       }
+#ifdef PADDLE_WITH_CINN
+      if (config_.cinn_enabled()) {
+        VLOG(4) << "[CINN] Begin AddCinnPass";
+        auto cinn_pm = std::make_shared<::pir::PassManager>(
+            ::pir::IrContext::Instance(), 2);
+        cinn::dialect::ir::AddCinnPass(cinn_pm, *pir_program_.get());
+        if (!config_.glog_info_disabled()) {
+          cinn_pm->EnablePrintStatistics();
+        }
+        if (config_.ir_debug_) {
+          cinn_pm->EnableIRPrinting();
+        }
+        cinn_pm->Run(pir_program_.get());
+      }
+#endif
 
       if (config_.use_gpu()) {
         ::pir::PassManager gpu_pm(::pir::IrContext::Instance(), 2);
@@ -1793,8 +1812,8 @@ void AnalysisPredictor::PrepareArgument() {
     argument_->SetMKLDNNEnabledOpTypes(config_.mkldnn_enabled_op_types_);
   }
 
-  if (config_.use_cinn_compiler_) {
-    argument_->SetUseCinnCompiler(config_.use_cinn_compiler_);
+  if (config_.cinn_enabled()) {
+    argument_->SetUseCinnCompiler(true);
   }
 
 #ifdef PADDLE_WITH_DNNL
@@ -1869,7 +1888,7 @@ void AnalysisPredictor::PrepareArgument() {
     LOG(INFO) << "Model is mixed precision type with " << model_precision_
               << ", we will use a new PassStrategy. Note that only GPU/XPU "
                  "backend is supported for now.";
-    if (!config_.use_cinn_compiler_) {
+    if (!config_.cinn_enabled()) {
       const auto &deleted_passes = pass_builder->GetAllDeletedPasses();
       if (config_.tensorrt_engine_enabled()) {
         pass_builder->ClearPasses();
diff --git a/paddle/fluid/inference/api/paddle_analysis_config.h b/paddle/fluid/inference/api/paddle_analysis_config.h
@@ -1196,15 +1196,15 @@ struct PD_INFER_DECL AnalysisConfig {
   ///
   /// \brief Enable use cinn compiler optimization.
   ///
-  void Exp_EnableCINNCompiler();
+  void EnableCINN();
 
   ///
   /// \brief A boolean state telling whether the CINN compiler optimization is
   /// turned on.
   ///
   /// \return bool Whether the CINN compiler optimization is turned on.
   ///
-  bool cinn_compiler_enabled() const;
+  bool cinn_enabled() const;
 
  protected:
   // Update the config.
@@ -1350,7 +1350,7 @@ struct PD_INFER_DECL AnalysisConfig {
   bool lite_zero_copy_;
 
   // CINN compiler related.
-  bool use_cinn_compiler_{false};
+  bool use_cinn_{false};
 
   // XPU related.
   bool use_xpu_{false};
diff --git a/paddle/fluid/pybind/CMakeLists.txt b/paddle/fluid/pybind/CMakeLists.txt
@@ -51,7 +51,7 @@ set(PYBIND_DEPS
 
 if(WITH_CINN)
   set(PYBIND_DEPS ${PYBIND_DEPS} pir_transforms cinn_transforms
-                  sub_graph_checker)
+                  sub_graph_checker add_cinn_pass)
 endif()
 
 if(WITH_PSCORE)
diff --git a/paddle/fluid/pybind/inference_api.cc b/paddle/fluid/pybind/inference_api.cc
@@ -978,6 +978,7 @@ void BindAnalysisConfig(py::module *m) {
       .def("enable_mkldnn", &AnalysisConfig::EnableMKLDNN)
       .def("disable_mkldnn", &AnalysisConfig::DisableMKLDNN)
       .def("mkldnn_enabled", &AnalysisConfig::mkldnn_enabled)
+      .def("enable_cinn", &AnalysisConfig::EnableCINN)
       .def("set_cpu_math_library_num_threads",
            &AnalysisConfig::SetCpuMathLibraryNumThreads)
       .def("cpu_math_library_num_threads",
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
