predictor supports phi, test=develop (PaddlePaddle#40856)

Author: Shixiaowei02

paddle/infrt/api/.gitignore

@@ -0,0 +1 @@
+infrt_api_test.cc

paddle/infrt/api/CMakeLists.txt

@@ -3,6 +3,7 @@ core_gather_headers()
 gather_srcs(infrt_src SRCS
     infrt_api.cc
     )
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/infrt_api_test.cc.in ${CMAKE_CURRENT_SOURCE_DIR}/infrt_api_test.cc)
 
 # Disable temporarily for the external-kernel's mkldnn is outdate
-# cc_test(test_infrt_api SRCS infrt_api_test.cc DEPS infrt ${MLIR_IR_LIBS})
+cc_test_tiny(test_infrt_api SRCS infrt_api_test.cc DEPS infrt ${MLIR_IR_LIBS})

paddle/infrt/api/infrt_api.cc

@@ -22,18 +22,27 @@
 #include <unordered_map>
 #include <vector>
 
+#include "mlir/Pass/PassManager.h"
+#include "paddle/infrt/backends/host/phi_allocator.h"
 #include "paddle/infrt/common/global.h"
 #include "paddle/infrt/dialect/dense_tensor.h"
 #include "paddle/infrt/dialect/infrt/ir/infrt_dialect.h"
+#include "paddle/infrt/dialect/infrt/pass/infrt_op_fuse_pass.h"
 #include "paddle/infrt/dialect/mlir_loader.h"
+#include "paddle/infrt/dialect/phi/ir/phi_base.h"
+#include "paddle/infrt/dialect/phi/pass/phi_op_convert_pass.h"
 #include "paddle/infrt/host_context/core_runtime.h"
 #include "paddle/infrt/host_context/kernel_registry.h"
 #include "paddle/infrt/host_context/mlir_function_executable.h"
 #include "paddle/infrt/host_context/mlir_to_runtime_translate.h"
 #include "paddle/infrt/host_context/op_executable.h"
+#include "paddle/infrt/host_context/paddle_mlir.h"
 #include "paddle/infrt/host_context/value.h"
 #include "paddle/infrt/kernel/basic_kernels.h"
 #include "paddle/infrt/kernel/control_flow_kernels.h"
+#include "paddle/infrt/kernel/phi/dense_tensor_kernels.h"
+#include "paddle/infrt/kernel/phi/infershaped/infershaped_kernel_launchers.h"
+#include "paddle/infrt/kernel/phi/registry.h"
 #include "paddle/infrt/kernel/tensor_kernels.h"
 #include "paddle/infrt/kernel/tensor_shape_kernels.h"
 #include "paddle/infrt/kernel/test_kernels.h"
@@ -84,12 +93,12 @@ class PredictExecutor : public MlirToRuntimeTranslator {
 
   PredictExecutor(mlir::ModuleOp module,
                   KernelRegistry* registry,
-                  TensorMap* map)
+                  ::infrt::phi::DenseTensorMap&& map)
       : MlirToRuntimeTranslator(module, &core_runtime),
         core_runtime(registry),
         registry_(registry) {
     CHECK(registry_);
-    Init(map);
+    Init(std::move(map));
   }
 
   void Run() {
@@ -100,18 +109,18 @@ class PredictExecutor : public MlirToRuntimeTranslator {
 
   int GetInputNum() { return inputs_.size(); }
 
-  DenseHostTensor* GetInput(int i) { return inputs_[i]; }
+  ::phi::DenseTensor* GetInput(int i) { return inputs_[i]; }
 
   int GetOutputNum() { return outputs_.size(); }
 
-  DenseHostTensor* GetOutput(int i) { return outputs_[i]; }
+  ::phi::DenseTensor* GetOutput(int i) { return outputs_[i]; }
 
  private:
-  void Init(TensorMap* map) {
+  void Init(::infrt::phi::DenseTensorMap&& map) {
     EmitFunctions();
     llvm::Optional<mlir::FuncOp> predict_func_ = llvm::None;
     for (auto func_op : impl_->module.getOps<mlir::FuncOp>()) {
-      if (func_op.getName().str() != "predict") continue;
+      if (func_op.getName().str() != "main_graph") continue;
       predict_func_ = func_op;
       break;
     }
@@ -125,30 +134,43 @@ class PredictExecutor : public MlirToRuntimeTranslator {
         new MlirFunctionExecutable(predict_func, registry_, impl_->func_defs);
 
     // process parammeters
+    VLOG(3) << "Arguments num of predict func: "
+            << predict_func.getNumArguments();
     for (size_t i = 0; i < predict_func.getNumArguments(); ++i) {
       auto arg = predict_func.getArgument(i);
       auto type = arg.getType();
       // this param is TensorMap
-      if (type.isa<infrt::DenseHostTensorMapType>()) {
-        auto* value = new host_context::Value(std::move(*map));
+      if (type.isa<::infrt::phi::DenseTensorMapType>()) {
+        auto* value = new host_context::Value(std::move(map));
         arguments_.push_back(value);
         AddValue(predict_func.getArgument(i), value);
-      } else {
+      } else if (type.isa<::infrt::DenseTensorType>()) {
         // this param is an input Tensor
-        auto dht = DenseHostTensor();
+        auto dht = ::phi::DenseTensor();
         auto* value = new host_context::Value(std::move(dht));
         arguments_.push_back(value);
-        inputs_.push_back(&(value->get<DenseHostTensor>()));
+        inputs_.push_back(&(value->get<::phi::DenseTensor>()));
+      } else {
+        llvm_unreachable("The input type has not been supported by predictor.");
       }
     }
 
     // process results
     auto& last_op = predict_func.front().back();
     if (last_op.getName().getStringRef() == "infrt.return") {
       for (size_t i = 0; i < last_op.getNumOperands(); ++i) {
-        auto* value = AddValue(mlir::Value(last_op.getOperand(i)));
-        results_.push_back(ValueRef(value));
-        outputs_.push_back(&(value->get<DenseHostTensor>()));
+        auto operand = last_op.getOperand(i);
+        if (operand.getType().isa<::infrt::DenseTensorType>()) {
+          auto r = impl_->value_map.try_emplace(
+              operand, ValueRef(new host_context::Value(::phi::DenseTensor())));
+          CHECK(r.second) << "Duplicate add mlir value ["
+                          << DumpToString(operand) << "]";
+          auto* value = r.first->second.get();
+          results_.push_back(ValueRef(value));
+          outputs_.push_back(&(value->get<::phi::DenseTensor>()));
+        } else {
+          llvm_unreachable("infrt.return only supports DenseTensor now.");
+        }
       }
     }
   }
@@ -166,22 +188,22 @@ class PredictExecutor : public MlirToRuntimeTranslator {
  private:
   KernelRegistry* registry_{};
   MlirFunctionExecutable* function_executable_;
-  llvm::SmallVector<DenseHostTensor*, 1> inputs_;
+  llvm::SmallVector<::phi::DenseTensor*, 1> inputs_;
   llvm::SmallVector<host_context::Value*, 2> arguments_;
-  llvm::SmallVector<DenseHostTensor*, 1> outputs_;
+  llvm::SmallVector<::phi::DenseTensor*, 1> outputs_;
   llvm::SmallVector<ValueRef, 1> results_;
 };
 
-std::shared_ptr<InfRtPredictor> CreateInfRtPredictor(
+std::unique_ptr<InfRtPredictor> CreateInfRtPredictor(
     const InfRtConfig& config) {
-  auto x = std::make_shared<InfRtPredictor>();
+  auto x = std::make_unique<InfRtPredictor>();
   x->Init(config);
   return x;
 }
 
 struct InfRtPredictor::Impl {
-  mlir::OwningModuleRef module_ref;
   std::unique_ptr<PredictExecutor> executor;
+  MLIRModelGenImpl module_gen_;
 };
 
 InfRtPredictor::InfRtPredictor() : impl_(new Impl) {}
@@ -190,8 +212,7 @@ InfRtPredictor::~InfRtPredictor() {}
 void InfRtPredictor::Run() { impl_->executor->Run(); }
 
 int InfRtPredictor::Init(const InfRtConfig& config) {
-  mlir::MLIRContext* context = infrt::Global::getMLIRContext();
-  auto module_ref = dialect::LoadMlirFile(config.mlir_path(), context);
+  mlir::MLIRContext* context = ::infrt::Global::getMLIRContext();
 
   KernelRegistry* registry = new KernelRegistry();
 
@@ -200,8 +221,32 @@ int InfRtPredictor::Init(const InfRtConfig& config) {
   kernel::RegisterTensorShapeKernels(registry);
   kernel::RegisterTensorKernels(registry);
   kernel::RegisterControlFlowKernels(registry);
-
-  impl_->module_ref = std::move(module_ref);
+#ifdef INFRT_WITH_PHI
+  kernel::RegisterPhiKernels(registry);
+  kernel::RegisterInferShapeLaunchers(registry);
+#if defined(INFRT_WITH_GPU) && defined(INFRT_WITH_TRT)
+  kernel::RegisterTrtKernels(registry);
+#endif  // INFRT_WITH_GPU && INFRT_WITH_TRT
+#endif
+
+  auto module_op = impl_->module_gen_.ImportPaddleModel(config.model_dir(),
+                                                        config.param_dir());
+
+  context->loadAllAvailableDialects();
+  ::mlir::PassManager pm(context);
+  ::mlir::OpPassManager& phi_pass_manager = pm.nest<::mlir::FuncOp>();
+  std::vector<::infrt::Place> valid_places = {{::infrt::TargetType::CPU,
+                                               ::infrt::PrecisionType::FLOAT32,
+                                               ::infrt::LayoutType::NCHW}};
+  phi_pass_manager.addPass(::infrt::createPhiOpCvtPass(valid_places));
+  phi_pass_manager.addPass(::infrt::createInfrtOpFusePass());
+  if (mlir::failed(pm.run(module_op))) {
+    std::cout << "\npass failed!\n" << std::endl;
+    return 4;
+  }
+#ifndef NDEBUG
+  module_op.dump();
+#endif  // NDEBUG
 
   // load extra shared library
   for (const std::string& lib_path : config.shared_libs()) {
@@ -222,23 +267,24 @@ int InfRtPredictor::Init(const InfRtConfig& config) {
   }
 
   // Load params
-  TensorMap* tensor_map = LoadParams(config.model_dir());
+  auto tensor_map = ::infrt::kernel::phi::LoadCombinedParameters(
+      config.model_dir(), config.param_dir());
 
   // Create PredictExecutor
   impl_->executor.reset(
-      new PredictExecutor(impl_->module_ref.get(), registry, tensor_map));
+      new PredictExecutor(module_op, registry, std::move(tensor_map)));
   return 0;
 }
 
 int InfRtPredictor::GetInputNum() { return impl_->executor->GetInputNum(); }
 
-DenseHostTensor* InfRtPredictor::GetInput(int i) {
+::phi::DenseTensor* InfRtPredictor::GetInput(int i) {
   return impl_->executor->GetInput(i);
 }
 
 int InfRtPredictor::GetOutputNum() { return impl_->executor->GetOutputNum(); }
 
-DenseHostTensor* InfRtPredictor::GetOutput(int i) {
+::phi::DenseTensor* InfRtPredictor::GetOutput(int i) {
   return impl_->executor->GetOutput(i);
 }
 

paddle/infrt/api/infrt_api.h

@@ -17,22 +17,22 @@
 #include <string>
 #include <vector>
 
-#include "paddle/infrt/tensor/dense_host_tensor.h"
+#include "paddle/phi/core/dense_tensor.h"
 
 namespace infrt {
 
 class InfRtConfig {
   std::string model_dir_;
-  std::string mlir_path_;
+  std::string param_dir_;
   std::vector<std::string> shared_libs_;
 
  public:
   InfRtConfig() = default;
   void set_model_dir(const std::string& model_dir) { model_dir_ = model_dir; }
   const std::string& model_dir() const { return model_dir_; }
 
-  void set_mlir_path(const std::string& mlir_path) { mlir_path_ = mlir_path; }
-  const std::string& mlir_path() const { return mlir_path_; }
+  void set_param_dir(const std::string& param_dir) { param_dir_ = param_dir; }
+  const std::string& param_dir() const { return param_dir_; }
 
   void set_shared_libs(const std::vector<std::string>& shared_libs) {
     shared_libs_ = shared_libs;
@@ -49,15 +49,15 @@ class InfRtPredictor {
   void Run();
   int Init(const InfRtConfig& config);
   int GetInputNum();
-  tensor::DenseHostTensor* GetInput(int i);
+  ::phi::DenseTensor* GetInput(int i);
   int GetOutputNum();
-  tensor::DenseHostTensor* GetOutput(int i);
+  ::phi::DenseTensor* GetOutput(int i);
 
  protected:
   struct Impl;
   std::unique_ptr<Impl> impl_;
 };
 
-std::shared_ptr<InfRtPredictor> CreateInfRtPredictor(const InfRtConfig& config);
+std::unique_ptr<InfRtPredictor> CreateInfRtPredictor(const InfRtConfig& config);
 
 }  // namespace infrt