runtime(cinn): update cinn jit instruction to support dynamic shape (#59470)

6clc · web-flow · commit ed58b869bc71 · 2023-12-01T10:51:16.000+08:00
* runtime(cinn): update cinn jit instruction to support dynamic shape

* runtime(cinn): update cinn jit instruction to solve conflict
diff --git a/paddle/cinn/hlir/dialect/operator/ir/attribute_storage.h b/paddle/cinn/hlir/dialect/operator/ir/attribute_storage.h
@@ -77,12 +77,12 @@ struct GroupInfoAttributeStorage : public pir::AttributeStorage {
   ParamKey data_;
 };
 
-struct JITInfoAttributeStorage : public pir::AttributeStorage {
-  using ParamKey = cinn::hlir::framework::pir::CUDAJITInfo;
-  explicit JITInfoAttributeStorage(const ParamKey& key) : data_(key) {}
+struct CINNKernelInfoAttributeStorage : public pir::AttributeStorage {
+  using ParamKey = cinn::hlir::framework::pir::CINNKernelInfo;
+  explicit CINNKernelInfoAttributeStorage(const ParamKey& key) : data_(key) {}
 
-  static JITInfoAttributeStorage* Construct(const ParamKey& key) {
-    return new JITInfoAttributeStorage(key);
+  static CINNKernelInfoAttributeStorage* Construct(const ParamKey& key) {
+    return new CINNKernelInfoAttributeStorage(key);
   }
 
   static std::size_t HashValue(const ParamKey& key) {
diff --git a/paddle/cinn/hlir/dialect/operator/ir/op_attribute.cc b/paddle/cinn/hlir/dialect/operator/ir/op_attribute.cc
@@ -20,12 +20,12 @@ const GroupInfo &GroupInfoAttribute::data() const {
   return storage()->GetAsKey();
 }
 
-const cinn::hlir::framework::pir::CUDAJITInfo &CUDAJITInfoAttribute::data()
-    const {
+const cinn::hlir::framework::pir::CINNKernelInfo &
+CINNKernelInfoAttribute::data() const {
   return storage()->GetAsKey();
 }
 }  // namespace dialect
 }  // namespace cinn
 
 IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::GroupInfoAttribute)
-IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::CUDAJITInfoAttribute)
+IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::CINNKernelInfoAttribute)
diff --git a/paddle/cinn/hlir/dialect/operator/ir/op_attribute.h b/paddle/cinn/hlir/dialect/operator/ir/op_attribute.h
@@ -33,22 +33,22 @@ class GroupInfoAttribute : public pir::Attribute {
   const GroupInfo& data() const;
 };
 
-class CUDAJITInfoAttribute : public pir::Attribute {
+class CINNKernelInfoAttribute : public pir::Attribute {
  public:
   using Attribute::Attribute;
 
-  DECLARE_ATTRIBUTE_UTILITY_FUNCTOR(CUDAJITInfoAttribute,
-                                    JITInfoAttributeStorage);
+  DECLARE_ATTRIBUTE_UTILITY_FUNCTOR(CINNKernelInfoAttribute,
+                                    CINNKernelInfoAttributeStorage);
 
-  bool operator<(const CUDAJITInfoAttribute& right) const {
+  bool operator<(const CINNKernelInfoAttribute& right) const {
     return storage() < right.storage();
   }
 
-  const cinn::hlir::framework::pir::CUDAJITInfo& data() const;
+  const cinn::hlir::framework::pir::CINNKernelInfo& data() const;
 };
 
 }  // namespace dialect
 }  // namespace cinn
 
 IR_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::GroupInfoAttribute)
-IR_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::CUDAJITInfoAttribute)
+IR_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::CINNKernelInfoAttribute)
diff --git a/paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc b/paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc
@@ -41,7 +41,7 @@ void OperatorDialect::initialize() {
   RegisterOp<ConcatOp>();
   RegisterOp<SplitOp>();
   RegisterAttribute<GroupInfoAttribute>();
-  RegisterAttribute<CUDAJITInfoAttribute>();
+  RegisterAttribute<CINNKernelInfoAttribute>();
 }
 
 void OperatorDialect::PrintType(pir::Type type, std::ostream &os) const {}
@@ -57,14 +57,14 @@ void OperatorDialect::PrintAttribute(pir::Attribute attr,
          << "[" << data.fn_name << "]";
     }
     { os << "<#AttrNotImplemented>"; }
-  } else if (attr.isa<CUDAJITInfoAttribute>()) {
-    auto cuda_jit_info = attr.dyn_cast<CUDAJITInfoAttribute>();
+  } else if (attr.isa<CINNKernelInfoAttribute>()) {
+    auto cinn_kernel_info = attr.dyn_cast<CINNKernelInfoAttribute>();
 
-    os << "(" << cuda_jit_info.data().fn_ptr;
+    os << "(" << cinn_kernel_info.data().fn_ptr;
     os << ')';
   } else {
     PADDLE_THROW(phi::errors::Unimplemented(
-        "cinn dialect only support GrupInfo and CUDAJITInfo"));
+        "cinn dialect only support GroupInfo and CINNKernelInfo"));
   }
 }
 
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.cc
@@ -177,7 +177,7 @@ class GroupOpPattern : public pir::OpRewritePattern<cinn::dialect::GroupOp> {
       auto fn_ptr_res = ir_compiler->BuildCUDAJITInfo({group});
       std::unordered_map<std::string, ::pir::Attribute> op_attrs{
           {cinn::dialect::JitKernelOp::kAttrName,
-           cinn::dialect::CUDAJITInfoAttribute::get(ctx, fn_ptr_res[0])},
+           cinn::dialect::CINNKernelInfoAttribute::get(ctx, fn_ptr_res[0])},
       };
 
       // Generate jit kernel op input and output
diff --git a/paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.cc b/paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.cc
@@ -46,16 +46,16 @@ void JitKernelOp::VerifySig() {
 
   auto& attributes = this->attributes();
 
-  IR_ENFORCE(
-      attributes.count(kAttrName) > 0 &&
-          attributes.at(kAttrName).isa<cinn::dialect::CUDAJITInfoAttribute>(),
-      "Type of attribute: instruction is not right.");
+  IR_ENFORCE(attributes.count(kAttrName) > 0 &&
+                 attributes.at(kAttrName)
+                     .isa<cinn::dialect::CINNKernelInfoAttribute>(),
+             "Type of attribute: instruction is not right.");
 }
 
-const hlir::framework::pir::CUDAJITInfo& JitKernelOp::cuda_jit_info() {
+const hlir::framework::pir::CINNKernelInfo& JitKernelOp::cinn_kernel_info() {
   return attributes()
       .at(kAttrName)
-      .dyn_cast<cinn::dialect::CUDAJITInfoAttribute>()
+      .dyn_cast<cinn::dialect::CINNKernelInfoAttribute>()
       .data();
 }
 
diff --git a/paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.h b/paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.h
@@ -27,7 +27,7 @@ class JitKernelOp : public ::pir::Op<JitKernelOp> {
   static const char* name() { return "cinn_runtime.jit_kernel"; }
   // TODO(Aurelius84): Think deeply what should contains
   static constexpr uint32_t attributes_num = 1;
-  static constexpr char* kAttrName = "jit_info";
+  static constexpr char* kAttrName = "kernel_info";
   static const char* attributes_name[attributes_num];
 
   static void Build(::pir::Builder& builder,             // NOLINT
@@ -36,7 +36,7 @@ class JitKernelOp : public ::pir::Op<JitKernelOp> {
                     const ::pir::AttributeMap& attributes,
                     const std::vector<::pir::Type>& out_types);
 
-  const hlir::framework::pir::CUDAJITInfo& cuda_jit_info();
+  const hlir::framework::pir::CINNKernelInfo& cinn_kernel_info();
 
   void VerifySig();
 };
diff --git a/paddle/cinn/hlir/framework/pir/compilation_task.cc b/paddle/cinn/hlir/framework/pir/compilation_task.cc
@@ -90,6 +90,17 @@ std::unique_ptr<Instruction> CompilationTask::BuildInstruction() {
   return instr;
 }
 
+pir::CINNKernelInfo CompilationTask::BuildPirCINNKernelInfo() {
+  std::string fn_name = context_->group_->FuncName();
+  VLOG(4) << "Lookup kernel name: " << fn_name;
+  auto* fn_ptr = context_->backend_compiler_->Lookup(fn_name);
+  CHECK(fn_ptr);
+  pir::CINNKernelInfo cinn_kernel_info;
+  cinn_kernel_info.fn_ptr = fn_ptr;
+  cinn_kernel_info.int_args_map = context_->group_->int_args_map;
+  return cinn_kernel_info;
+}
+
 }  // namespace framework
 }  // namespace hlir
 }  // namespace cinn
diff --git a/paddle/cinn/hlir/framework/pir/compilation_task.h b/paddle/cinn/hlir/framework/pir/compilation_task.h
@@ -63,6 +63,7 @@ class CompilationTask {
   void Lowering();
   void CodegenAndJit();
   std::unique_ptr<Instruction> BuildInstruction();
+  pir::CINNKernelInfo BuildPirCINNKernelInfo();
 
  private:
   GroupCompilationContext* context_;
diff --git a/paddle/cinn/hlir/framework/pir/group.h b/paddle/cinn/hlir/framework/pir/group.h
@@ -81,6 +81,7 @@ struct Group {
   std::vector<std::string> output_names;
   std::vector<::pir::Value> output_values;
   std::string fn_name{""};
+  std::map<int, CINNKernelInfo::ArgDimIdx> int_args_map;
 
   struct SharedGroupHasher {
     size_t operator()(const std::shared_ptr<Group>& group) const noexcept {
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -232,7 +232,7 @@ OpLowererImpl::BucketLower(const GroupPtr& group,
     group_scheduler->Schedule();
     cond2func_bodies = group_scheduler->GetIRs();
   } else {
-    cond2func_bodies.emplace_back(ir::Expr(1),
+    cond2func_bodies.emplace_back(ir::Expr(true),
                                   ir_sch.GetModule().GetExprs()[0]);
   }
 
@@ -506,17 +506,40 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
     for (auto arg : group_func_args) {
       args_set.insert(arg.name());
     }
+    for (auto& op : group->ops) {
+      // collect all output tensor.
+      for (auto opresult : op->results()) {
+        if (tensor_map.count(opresult) == 0) {
+          continue;
+        }
+        auto tensor = tensor_map.at(opresult);
+        if (args_set.count("_" + tensor->name) != 0) {
+          continue;
+        }
+        group->output_values.push_back(opresult);
+        group_func_arg_tensors->push_back(tensor);
+        group->output_names.push_back(tensor->name);
+        group_func_args.emplace_back(tensor->buffer, ir::Argument::IO::kOutput);
+      }
+    }
+  }
 
-    for (auto& tensor_pair : tensor_map) {
-      if (args_set.count("_" + tensor_pair.second->name)) {
-        continue;
+  std::map<int, CINNKernelInfo::ArgDimIdx> mps;
+  // update args for dynamic dim
+  int num_tensor_args = static_cast<int>(group_func_args.size());
+  int non_tensor_arg_idx = group_func_args.size();
+  for (int tensor_arg_idx = 0; tensor_arg_idx < num_tensor_args;
+       tensor_arg_idx++) {
+    auto tensor_dim = (*group_func_arg_tensors)[tensor_arg_idx]->sym_shape;
+    int tensor_dim_size = tensor_dim.size();
+    for (int tensor_arg_dim_idx = 0; tensor_arg_dim_idx < tensor_dim_size;
+         tensor_arg_dim_idx++) {
+      if (tensor_dim[tensor_arg_dim_idx]->IsDynamic()) {
+        group_func_args.emplace_back(ir::_Var_::Make(
+            tensor_dim[tensor_arg_dim_idx]->GetSymbolName(), common::Int(32)));
+        group->int_args_map[non_tensor_arg_idx++] = {tensor_arg_idx,
+                                                     tensor_arg_dim_idx};
       }
-      group_func_arg_tensors->push_back(tensor_pair.second);
-      // use the underlying tensor name to be consistent with the argument name
-      // in the lowered function
-      group->output_names.push_back(tensor_pair.second->name);
-      group_func_args.emplace_back(tensor_pair.second->buffer,
-                                   ir::Argument::IO::kOutput);
     }
   }
 
diff --git a/paddle/cinn/hlir/framework/pir/utils.h b/paddle/cinn/hlir/framework/pir/utils.h
@@ -29,11 +29,26 @@ namespace framework {
 
 namespace pir {
 
-struct CUDAJITInfo {
+struct CINNKernelInfo {
   void* fn_ptr;
-  std::vector<int> block_dims;
-  std::vector<int> grid_dims;
-  void* compiler;
+
+  struct ArgDimIdx {
+    int arg_idx;
+    int dim_idx;
+  };
+  // int_args_map records the int_args_map.key argument (dtype is Int) in the
+  // kernel parameter taken from the dim_idx dimension of the shape of the
+  // ArgDimIdx.arg_idx argument.
+  // Examples:
+  //   a func like: foo(tensor A, tensor B, int S1, int S2)
+  //   S1 = A.shape[3]
+  //   S2 = B.shape[2]
+  //   int_args_map will be like
+  //   {
+  //     2: {0, 3},
+  //     3: {1, 2}
+  //   }
+  std::map<int, ArgDimIdx> int_args_map;
 };
 
 struct CompatibleInfo {
diff --git a/paddle/cinn/hlir/framework/pir_compiler.cc b/paddle/cinn/hlir/framework/pir_compiler.cc
@@ -42,44 +42,49 @@ std::unique_ptr<Program> PirCompiler::Build() {
   return std::move(Build(groups));
 }
 
-std::vector<pir::CUDAJITInfo> PirCompiler::BuildCUDAJITInfo(
+std::vector<pir::CINNKernelInfo> PirCompiler::BuildCUDAJITInfo(
     const std::vector<pir::GroupPtr>& groups) {
-  std::vector<pir::CUDAJITInfo> vec_res;
+  std::vector<pir::CINNKernelInfo> cinn_kernel_info_vecs(groups.size());
 
-  auto op_lowerer = CreateOpLowerer<pir::GroupPtr>(target_);
-
-  std::vector<std::vector<ir::LoweredFunc>> lowered_funcs;
-  for (int i = 0; i < groups.size(); ++i) {
-    lowered_funcs.emplace_back(op_lowerer.Lower(groups[i]));
-  }
-
-  for (auto&& lowered_func : lowered_funcs) {
-    ProcessFunction(lowered_func);
-  }
-
-  compiler_ = backends::Compiler::Create(target_);
-  auto build_module = m_builder_.Build();
-  compiler_->Build(build_module, "");
-
-  auto instructions = BuildInstructions(groups);
+  if (FLAGS_cinn_bucket_compile) {
+    for (int i = 0; i < groups.size(); ++i) {
+      group_compilation_contexts_.emplace_back(target_, groups[i], scope_);
+    }
+    auto worker_fn = [&](int index) {
+      CompilationTask task(&group_compilation_contexts_[index]);
+      task();
+      cinn_kernel_info_vecs[index] = task.BuildPirCINNKernelInfo();
+    };
+    utils::parallel_run(
+        worker_fn, utils::SequenceDispatcher(0, groups.size()), -1);
+  } else {
+    auto op_lowerer = CreateOpLowerer<pir::GroupPtr>(target_);
 
-  auto fn_ptrs = compiler_->GetFnPtr();
+    std::vector<std::vector<ir::LoweredFunc>> lowered_funcs;
+    for (int i = 0; i < groups.size(); ++i) {
+      lowered_funcs.emplace_back(op_lowerer.Lower(groups[i]));
+    }
 
-  auto* compilter_ptr = compiler_.release();
-  for (int idx = 0; idx < groups.size(); ++idx) {
-    pir::CUDAJITInfo jit_info;
-    jit_info.fn_ptr = fn_ptrs[idx];
-    jit_info.compiler = reinterpret_cast<void*>(compilter_ptr);
+    for (auto&& lowered_func : lowered_funcs) {
+      ProcessFunction(lowered_func);
+    }
+    compiler_ = backends::Compiler::Create(target_);
+    auto build_module = m_builder_.Build();
+    compiler_->Build(build_module, "");
 
-    lowered_funcs[idx][0]->cuda_axis_info.CopyBlockDimsTo(
-        &(jit_info.block_dims));
+    auto fn_ptrs = compiler_->GetFnPtr();
 
-    lowered_funcs[idx][0]->cuda_axis_info.CopyGridDimsTo(&(jit_info.grid_dims));
+    for (int idx = 0; idx < groups.size(); ++idx) {
+      pir::CINNKernelInfo cinn_kernel_info;
+      auto fn_name = groups[idx]->FuncName();
+      auto fn_ptr = compiler_->Lookup(fn_name);
+      cinn_kernel_info.fn_ptr = fn_ptr;
+      cinn_kernel_info.int_args_map = groups[idx]->int_args_map;
 
-    vec_res.push_back(jit_info);
+      cinn_kernel_info_vecs[idx] = cinn_kernel_info;
+    }
   }
-
-  return vec_res;
+  return cinn_kernel_info_vecs;
 }
 
 std::unique_ptr<Program> PirCompiler::Build(
diff --git a/paddle/cinn/hlir/framework/pir_compiler.h b/paddle/cinn/hlir/framework/pir_compiler.h
@@ -41,7 +41,7 @@ class PirCompiler final {
 
   std::unique_ptr<Program> Build();
 
-  std::vector<pir::CUDAJITInfo> BuildCUDAJITInfo(
+  std::vector<pir::CINNKernelInfo> BuildCUDAJITInfo(
       const std::vector<pir::GroupPtr>& groups);
 
   std::unique_ptr<Program> Build(const std::vector<pir::GroupPtr>& groups);
diff --git a/paddle/fluid/framework/new_executor/instruction/cinn_jit_instruction.cc b/paddle/fluid/framework/new_executor/instruction/cinn_jit_instruction.cc
diff --git a/test/cpp/pir/cinn/jit_instruction_test.cc b/test/cpp/pir/cinn/jit_instruction_test.cc
diff --git a/test/cpp/pir/cinn/pir_compiler_test.cc b/test/cpp/pir/cinn/pir_compiler_test.cc

Original file line number	Diff line number	Diff line change
`@@ -20,12 +20,12 @@ const GroupInfo &GroupInfoAttribute::data() const {`
`20`	`20`	`return storage()->GetAsKey();`
`21`	`21`	`}`
`22`	`22`
`23`		`-const cinn::hlir::framework::pir::CUDAJITInfo &CUDAJITInfoAttribute::data()`
`24`		`- const {`
	`23`	`+const cinn::hlir::framework::pir::CINNKernelInfo &`
	`24`	`+CINNKernelInfoAttribute::data() const {`
`25`	`25`	`return storage()->GetAsKey();`
`26`	`26`	`}`
`27`	`27`	`} // namespace dialect`
`28`	`28`	`} // namespace cinn`
`29`	`29`
`30`	`30`	`IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::GroupInfoAttribute)`
`31`		`-IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::CUDAJITInfoAttribute)`
	`31`	`+IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::CINNKernelInfoAttribute)`
Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ void OperatorDialect::initialize() {`
`41`	`41`	`RegisterOp<ConcatOp>();`
`42`	`42`	`RegisterOp<SplitOp>();`
`43`	`43`	`RegisterAttribute<GroupInfoAttribute>();`
`44`		`- RegisterAttribute<CUDAJITInfoAttribute>();`
	`44`	`+ RegisterAttribute<CINNKernelInfoAttribute>();`
`45`	`45`	`}`
`46`	`46`
`47`	`47`	`void OperatorDialect::PrintType(pir::Type type, std::ostream &os) const {}`
`@@ -57,14 +57,14 @@ void OperatorDialect::PrintAttribute(pir::Attribute attr,`
`57`	`57`	`<< "[" << data.fn_name << "]";`
`58`	`58`	`}`
`59`	`59`	`{ os << "<#AttrNotImplemented>"; }`
`60`		`- } else if (attr.isa<CUDAJITInfoAttribute>()) {`
`61`		`- auto cuda_jit_info = attr.dyn_cast<CUDAJITInfoAttribute>();`
	`60`	`+ } else if (attr.isa<CINNKernelInfoAttribute>()) {`
	`61`	`+ auto cinn_kernel_info = attr.dyn_cast<CINNKernelInfoAttribute>();`
`62`	`62`
`63`		`- os << "(" << cuda_jit_info.data().fn_ptr;`
	`63`	`+ os << "(" << cinn_kernel_info.data().fn_ptr;`
`64`	`64`	`os << ')';`
`65`	`65`	`} else {`
`66`	`66`	`PADDLE_THROW(phi::errors::Unimplemented(`
`67`		`- "cinn dialect only support GrupInfo and CUDAJITInfo"));`
	`67`	`+ "cinn dialect only support GroupInfo and CINNKernelInfo"));`
`68`	`68`	`}`
`69`	`69`	`}`
`70`	`70`