[NewIR]Support Instruction.Run in CINN for Runtime::Program (#55680)

Aurelius84 · web-flow · commit f9e1b2d2ad1d · 2023-07-25T15:05:40.000+08:00
diff --git a/paddle/cinn/hlir/framework/new_ir_compiler.h b/paddle/cinn/hlir/framework/new_ir_compiler.h
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+#include <absl/types/variant.h>
 #include <memory>
 #include <unordered_map>
 #include "paddle/cinn/common/context.h"
@@ -30,9 +31,15 @@ namespace cinn {
 namespace hlir {
 namespace framework {
 
-// TODO(Aurelius): Need add name mapping logic in REGISTER_CINN_OP
-// macros or attempt to unify Op name with Paddle and CINN.
-static const std::unordered_map<std::string, std::string> OP_NAMES = {
+struct CompatibleInfo {
+  static constexpr char* kInputPrefix = "input_";
+  static constexpr char* kOutputPrefix = "output_";
+  // TODO(Aurelius): Need add name mapping logic in REGISTER_CINN_OP
+  // macros or attempt to unify Op name with Paddle and CINN.
+  static const std::unordered_map<std::string, std::string> OP_NAMES;
+};
+
+const std::unordered_map<std::string, std::string> CompatibleInfo::OP_NAMES = {
     {"pd.full", "fill_constant"}, {"pd.matmul", "matmul"}};
 
 // TODO(Aurelius84): Need abstract this logic to implement Proxy for
@@ -70,18 +77,32 @@ class NewIRCompiler final {
     compiler_->Build(build_module, "");
 
     auto instructions = BuildInstructions(groups);
+
+    // TODO(Aurelius84): Instantiate all tensors on compile-time, which is
+    // controlled by 'options.with_instantiate_variables' in GraphCompiler.
+    // Moreover, it's better to implement InsertBufferHandlers() logic
+    // to automatically insert Malloc and Free instructions.
+    for (auto& name : scope_->var_names()) {
+      std::string var_name({name.data(), name.size()});
+      VLOG(4) << "Instantiate " << var_name << " on compile-time";
+      auto* var = scope_->Var<Tensor>(var_name);
+      auto& tensor = absl::get<Tensor>(*var);
+      tensor->mutable_data(target_, tensor->type());
+    }
     return std::make_unique<Program>(scope_, std::move(instructions));
   }
 
   std::vector<ir::LoweredFunc> GetOpFunc(const ::ir::Operation& op, int idx) {
     std::vector<ir::Tensor> inputs;
     std::vector<common::CINNValue> cinn_inputs;
-    VLOG(4) << "GetOpFunc for op: " << op.name();
+    auto op_name = op.name();
+    VLOG(4) << "GetOpFunc for op: " << op_name;
     // step 1: Deal with Oprands
     for (int i = 0; i < op.num_operands(); ++i) {
       auto in_value = op.operand(i);
       // TODO(Aurelius84): For now, use addr as name but it's not wise.
-      std::string input_id = std::to_string(std::hash<::ir::Value>()(in_value));
+      std::string input_id = CompatibleInfo::kInputPrefix +
+                             std::to_string(std::hash<::ir::Value>()(in_value));
       // NOTE(Aurelius84): whether need to support other Type?
       auto type_info =
           in_value.type().dyn_cast<paddle::dialect::DenseTensorType>();
@@ -100,8 +121,7 @@ class NewIRCompiler final {
       cinn_inputs.push_back(common::CINNValue(temp));
     }
     for (auto out_name : OpGetOutputNames(op)) {
-      cinn_inputs.push_back(
-          common::CINNValue(op.name().substr(3) + "_" + out_name));
+      cinn_inputs.push_back(common::CINNValue(out_name));
     }
 
     VLOG(4) << "inputs.size(): " << inputs.size();
@@ -124,14 +144,14 @@ class NewIRCompiler final {
     {
       VLOG(4) << "op.attributes():" << op.attributes().size();
       auto attrs = utils::ConvertAttributes(op.attributes());
-      node_attrs.node_name = OP_NAMES.at(op.name());
+      node_attrs.node_name = CompatibleInfo::OP_NAMES.at(op_name);
       node_attrs.attr_store = std::move(attrs);
     }
     auto& strategy = Operator::GetAttrs<StrategyFunction>("CINNStrategy");
     // NOTE(Aurelius84): Do we need replace all hlir::framework Operator with
     // ::ir::Program ？
     const hlir::framework::Operator* cinn_op =
-        Operator::Get(OP_NAMES.at(op.name()));
+        Operator::Get(CompatibleInfo::OP_NAMES.at(op_name));
     auto impl = OpStrategy::SelectImpl(
         strategy[cinn_op](node_attrs, inputs, out_types, out_shapes, target_));
     common::CINNValuePack C =
@@ -223,7 +243,8 @@ class NewIRCompiler final {
     std::unordered_set<std::string> repeat;
     for (int i = 0; i < op.num_operands(); ++i) {
       auto value = op.operand(i);
-      std::string name = std::to_string(std::hash<::ir::Value>()(value));
+      std::string name = CompatibleInfo::kInputPrefix +
+                         std::to_string(std::hash<::ir::Value>()(value));
       if (repeat.count(name)) {
         continue;
       }
@@ -237,7 +258,8 @@ class NewIRCompiler final {
     std::vector<std::string> names;
     for (int i = 0; i < op.num_results(); ++i) {
       auto value = op.result(i);
-      std::string name = std::to_string(std::hash<::ir::Value>()(value));
+      std::string name = CompatibleInfo::kOutputPrefix +
+                         std::to_string(std::hash<::ir::Value>()(value));
       names.push_back(std::move(name));
     }
     return names;
@@ -257,11 +279,12 @@ std::shared_ptr<Scope> BuildScope(const Target& target,
   std::unordered_set<::ir::Value> visited;
   auto scope = std::make_shared<Scope>();
 
-  auto create_var = [&](::ir::Value value) {
+  auto create_var = [&](const std::string& name_prefix, ::ir::Value value) {
     if (visited.count(value) > 0) return;
     visited.emplace(value);
 
-    std::string name = std::to_string(std::hash<::ir::Value>()(value));
+    std::string name =
+        name_prefix + std::to_string(std::hash<::ir::Value>()(value));
     auto type_info = value.type().dyn_cast<paddle::dialect::DenseTensorType>();
     auto* var = scope->Var<Tensor>(name);
     auto& tensor = absl::get<Tensor>(*var);
@@ -279,12 +302,12 @@ std::shared_ptr<Scope> BuildScope(const Target& target,
     // visit OpOprands
     for (auto i = 0; i < (*it)->num_operands(); ++i) {
       auto in_value = (*it)->operand(i);
-      create_var(in_value);
+      create_var(CompatibleInfo::kInputPrefix, in_value);
     }
 
     for (auto i = 0; i < (*it)->num_results(); ++i) {
       auto out_value = (*it)->result(i);
-      create_var(out_value);
+      create_var(CompatibleInfo::kOutputPrefix, out_value);
     }
   }
   return scope;
diff --git a/test/cpp/ir/cinn/graph_compiler_new_ir_test.cc b/test/cpp/ir/cinn/graph_compiler_new_ir_test.cc
@@ -24,6 +24,7 @@
 #include "paddle/cinn/frontend/net_builder.h"
 #include "paddle/cinn/frontend/optimize.h"
 #include "paddle/cinn/hlir/framework/graph_compiler.h"
+#include "paddle/cinn/utils/data_util.h"
 
 #include "paddle/cinn/hlir/framework/new_ir_compiler.h"
 
@@ -33,15 +34,16 @@ TEST(GraphCompier, TestNewIR) {
   ::ir::Program program(ctx);
   ::ir::Builder builder = ::ir::Builder(ctx, program.block());
 
+  const float value = 2.0;
   auto full_op_x =
       builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 128},
-                                             1.0,
+                                             value,
                                              phi::DataType::FLOAT32,
                                              phi::CPUPlace());
 
   auto full_op_y =
       builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{128, 64},
-                                             2.0,
+                                             value,
                                              phi::DataType::FLOAT32,
                                              phi::CPUPlace());
   // TODO(Aurelius84): test more op
@@ -61,7 +63,15 @@ TEST(GraphCompier, TestNewIR) {
   cinn::hlir::framework::NewIRCompiler ir_compiler(program, target, scope);
   auto runtime_program = ir_compiler.Build();
 
-  // FIXME(Aurelius84): It raised illegal memory access while deconstructor
-  // after running all instruction, but it's ok under GLOG_v=10.
-  // ASSERT_NO_THROW(runtime_program->Execute());
+  ASSERT_NO_THROW(runtime_program->Execute());
+
+  for (auto& var_name : scope->var_names()) {
+    std::string name = {var_name.begin(), var_name.end()};
+    std::vector<float> data =
+        cinn::GetTensorData<float>(scope->GetTensor(name), target);
+    for (int i = 0; i < data.size(); ++i) {
+      LOG_FIRST_N(INFO, 3) << "data: " << data[i];
+      ASSERT_NEAR(data[i], value, 1e-5);
+    }
+  }
 }
