PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+26-6Lines changed: 26 additions & 6 deletions b/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+26-6Lines changed: 26 additions & 6 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/utils.cc
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/utils.cc
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/framework/pir/CMakeLists.txt
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/cinn/hlir/framework/pir/CMakeLists.txt
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_cache.cc
Copy file name to clipboard
+10-33Lines changed: 10 additions & 33 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_cache.cc
Copy file name to clipboard
+10-33Lines changed: 10 additions & 33 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_cache.h
Copy file name to clipboard
+33-29Lines changed: 33 additions & 29 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_cache.h
Copy file name to clipboard
+33-29Lines changed: 33 additions & 29 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_task.cc
Copy file name to clipboard
+15-26Lines changed: 15 additions & 26 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_task.cc
Copy file name to clipboard
+15-26Lines changed: 15 additions & 26 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_task.h
Copy file name to clipboard
+4-5Lines changed: 4 additions & 5 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_task.h
Copy file name to clipboard
+4-5Lines changed: 4 additions & 5 deletions
@@ -82,7 +82,10 @@ pir::Block* GroupOp::block() {
 
 pir::Block* GroupOp::block() const {
   pir::Region& region = (*this)->region(0);
-  CHECK(!region.empty());
+  PADDLE_ENFORCE_EQ(region.empty(),
+                    false,
+                    ::common::errors::Unavailable(
+                        "Required GroupOp's region must not be emptpy."));
   return &region.front();
 }
 
@@ -155,7 +158,16 @@ pir::Block* FusionOp::block() {
   return &region.front();
 }
 
-std::vector<pir::Operation*> FusionOp::GetOperators() {
+pir::Block* FusionOp::block() const {
+  pir::Region& region = (*this)->region(0);
+  PADDLE_ENFORCE_EQ(region.empty(),
+                    false,
+                    ::common::errors::Unavailable(
+                        "Required FusionOp's region must not be emptpy."));
+  return &region.front();
+}
+
+std::vector<pir::Operation*> FusionOp::GetOperators() const {
   std::vector<pir::Operation*> rt_ops;
   for (auto& op : *block()) {
     rt_ops.push_back(&op);
@@ -304,7 +316,9 @@ void GenerateShapeOp::Build(
   if (inputs.empty()) {
     VLOG(3) << "GenerateShapeOp inputs is empty";
     for (const auto& attr : output_dim_exprs) {
-      CHECK(attr.isa<pir::Int64Attribute>());
+      PADDLE_ENFORCE(attr.isa<pir::Int64Attribute>(),
+                     ::common::errors::PreconditionNotMet(
+                         "Reqiured attr must be Int64Attribute."));
     }
   }
   argument.AddInputs(inputs);
@@ -466,11 +480,15 @@ bool GenerateShapeOp::InferSymbolicShape(
   const auto attr_dim_exprs = [&] {
     std::vector<symbol::DimExpr> dim_exprs{};
     pir::Attribute dim_expr_attr = this->attributes().at("output_dim_exprs");
-    CHECK(dim_expr_attr.isa<pir::ArrayAttribute>());
+    PADDLE_ENFORCE(dim_expr_attr.isa<pir::ArrayAttribute>(),
+                   ::common::errors::PreconditionNotMet(
+                       "Required dim_expr_attr is ArrayAttribute."));
     auto array = dim_expr_attr.dyn_cast<pir::ArrayAttribute>();
     for (int i = 0; i < array.size(); ++i) {
       const auto& dim_expr = ConvertAttributeToDimExpr(array.at(i));
-      CHECK(dim_expr.has_value());
+      PADDLE_ENFORCE(dim_expr.has_value(),
+                     ::common::errors::PreconditionNotMet(
+                         "Required dim_expr.has_value()==true."));
       dim_exprs.push_back(dim_expr.value());
     }
     return dim_exprs;
@@ -480,7 +498,9 @@ bool GenerateShapeOp::InferSymbolicShape(
         this->attributes().at("symbol_bindings");
     auto symbol_bindings =
         ConvertAttributeToSymbolBindings(symbol_bindings_attr);
-    CHECK(symbol_bindings.has_value());
+    PADDLE_ENFORCE(symbol_bindings.has_value(),
+                   ::common::errors::PreconditionNotMet(
+                       "Required symbol_bindings.has_value()==true."));
     return symbol_bindings.value();
   }();
   auto DimExprs4InputDim =
 
@@ -77,7 +77,8 @@ class IR_API FusionOp : public pir::Op<FusionOp> {
                     const cinn::dialect::GroupInfo &group_info);
 
   pir::Block *block();
-  std::vector<pir::Operation *> GetOperators();
+  pir::Block *block() const;
+
   std::vector<pir::Operation *> GetOperators() const;
 
   void VerifySig();
 
@@ -16,4 +16,4 @@ cinn_cc_library(cinn_transforms SRCS ${cinn_transforms_srcs} DEPS
 cc_library(
   add_cinn_pass
   SRCS add_cinn_pass.cc
-  DEPS op_dialect pir cinn_op_dialect cinnapi pir_transforms cinn_transforms)
+  DEPS pir_transforms cinn_transforms)
@@ -78,7 +78,8 @@ CompileGroupAsOpAttribute(const std::vector<OpLoweringGroupPtr>& group_list) {
 
 std::unordered_map<std::string, ::pir::Attribute> GetJitKernelAttr(
     const OpLoweringGroupPtr& group) {
-  auto kernel_info = CompilationCache::Instance().GetKernelInfo(group);
+  hlir::framework::pir::FusionInfo fusion_info(*group);
+  auto kernel_info = CompilationCache::Instance().GetKernelInfo(fusion_info);
   std::unordered_map<std::string, ::pir::Attribute> attrs{
       {cinn::dialect::JitKernelOp::kAttrName,
        cinn::dialect::CINNKernelInfoAttribute::get(pir::IrContext::Instance(),
 
@@ -11,4 +11,5 @@ gather_srcs(
   trivial_op_impl.cc
   trivial_op_util.cc
   compilation_task.cc
-  compilation_cache.cc)
+  compilation_cache.cc
+  fusion_info.cc)
@@ -39,38 +39,20 @@ void* BackendResource::GetInferFuncPtr() const {
   return ptr;
 }
 
-std::shared_ptr<backends::Compiler>& BackendResource::GetBackendCompiler() {
-  return backend_compiler_;
-}
-
-const std::shared_ptr<backends::Compiler>& BackendResource::GetBackendCompiler()
-    const {
-  return backend_compiler_;
-}
-
-void BackendResource::SetHostFnName(const std::string& name) {
-  host_fn_name_ = name;
-}
-
-void BackendResource::SetInferFnName(const std::string& name) {
-  infer_fn_name_ = name;
-}
-
-pir::CINNKernelInfo BackendResource::GernerateKernelInfo(
-    const std::shared_ptr<pir::OpLoweringGroup>& group) const {
+pir::CINNKernelInfo BackendResource::GenerateKernelInfo() const {
   pir::CINNKernelInfo kernel_info;
   kernel_info.fn_name = host_fn_name_;
   kernel_info.fn_ptr = GetHostFuncPtr();
   kernel_info.infer_shape_fn_ptr = GetInferFuncPtr();
-  kernel_info.int_args_map = group->int_args_map();
+  kernel_info.int_args_map = GetIntArgsMap();
   return kernel_info;
 }
 }  // namespace pir
 
 bool CompilationCache::Has(const CacheKey& key) const {
-  const bool has_existed = cache_.find(KeyHash(key)) != cache_.end();
-  VLOG(6) << "Check IsExisted in CompilationCache: " << key->FuncName() << " "
-          << has_existed;
+  const bool has_existed = cache_.find(key) != cache_.end();
+  VLOG(6) << "Check IsExisted in CompilationCache: " << has_existed << " - "
+          << key;
   return has_existed;
 }
 
@@ -79,24 +61,19 @@ const CompilationCache::CacheValue& CompilationCache::Get(
   PADDLE_ENFORCE_EQ(
       Has(key),
       true,
-      phi::errors::NotFound("%s is not in CompliatonCache.", key->FuncName()));
-  return cache_.at(KeyHash(key));
+      phi::errors::NotFound("%s is not in CompliatonCache.", key));
+  return cache_.at(key);
 }
 
 pir::CINNKernelInfo CompilationCache::GetKernelInfo(const CacheKey& key) const {
-  return Get(key)->GetKernelInfo(key);
+  return Get(key)->GetKernelInfo();
 }
 
 void CompilationCache::Insert(const CacheKey& key, const CacheValue& value) {
-  VLOG(6) << "Insert CompilationCache for: " << key->FuncName();
-  cache_.insert({KeyHash(key), value});
+  VLOG(6) << "Insert CompilationCache for: " << key;
+  cache_.insert({key, value});
 }
 
 void CompilationCache::Clear() { cache_.clear(); }
 
-size_t CompilationCache::KeyHash(const CacheKey& key) const {
-  // TODO(Aurelius84): use a better hash function in next pr.
-  return std::hash<std::string>{}(key->FuncName());
-}
-
 }  // namespace cinn::hlir::framework
@@ -19,6 +19,7 @@
 #include "paddle/cinn/backends/compiler.h"
 #include "paddle/cinn/common/macros.h"
 #include "paddle/cinn/common/target.h"
+#include "paddle/cinn/hlir/framework/pir/fusion_info.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 
 namespace cinn::hlir::framework {
@@ -27,76 +28,79 @@ namespace pir {
 class OpLoweringGroup;
 class BackendResource final {
  public:
-  BackendResource(const Target& target) {
-    backend_compiler_ = backends::Compiler::Create(target);
-  }
-
   BackendResource(const Target& target,
                   const std::string& host_fn_name,
-                  const std::string& infer_fn_name)
-      : host_fn_name_(host_fn_name), infer_fn_name_(infer_fn_name) {
+                  const std::string& infer_fn_name,
+                  const std::map<int, CINNKernelInfo::ArgDimIdx>& int_args_map)
+      : host_fn_name_(host_fn_name),
+        infer_fn_name_(infer_fn_name),
+        int_args_map_(int_args_map) {
     backend_compiler_ = backends::Compiler::Create(target);
   }
 
   void* GetHostFuncPtr() const;
   void* GetInferFuncPtr() const;
-  pir::CINNKernelInfo GernerateKernelInfo(
-      const std::shared_ptr<pir::OpLoweringGroup>& group) const;
-  std::shared_ptr<backends::Compiler>& GetBackendCompiler();
-  const std::shared_ptr<backends::Compiler>& GetBackendCompiler() const;
-  void SetHostFnName(const std::string& name);
-  void SetInferFnName(const std::string& name);
+  const std::map<int, CINNKernelInfo::ArgDimIdx>& GetIntArgsMap() const {
+    return int_args_map_;
+  }
+  const std::shared_ptr<backends::Compiler>& GetBackendCompiler() const {
+    return backend_compiler_;
+  }
+  pir::CINNKernelInfo GenerateKernelInfo() const;
 
  private:
   std::string host_fn_name_;
   std::string infer_fn_name_;
-  // std::string host_code_;
-  // std::vector<std::string> device_code_;
-  std::shared_ptr<backends::Compiler> backend_compiler_;
+  std::map<int, CINNKernelInfo::ArgDimIdx> int_args_map_;
+
+  std::shared_ptr<backends::Compiler> backend_compiler_{nullptr};
 };
 
 class CompilationResult final {
  public:
-  explicit CompilationResult(const Target& target)
-      : target_(target), backend_resource_(target) {}
-
-  BackendResource& MutableBackendResource() { return backend_resource_; }
-  const BackendResource& GetBackendResource() const {
+  explicit CompilationResult(const Target& target) : target_(target) {}
+  const std::shared_ptr<BackendResource>& GetBackendResource() const {
     return backend_resource_;
   }
-  pir::CINNKernelInfo GetKernelInfo(
-      const std::shared_ptr<pir::OpLoweringGroup>& group) {
-    return backend_resource_.GernerateKernelInfo(group);
+
+  void SetBackendResource(const std::shared_ptr<BackendResource>& other) {
+    backend_resource_ = other;
+  }
+
+  pir::CINNKernelInfo GetKernelInfo() {
+    // TODO(Aurelius84): add ENFORCE_NOT_NULL
+    return backend_resource_->GenerateKernelInfo();
   }
 
  private:
   Target target_;
-  BackendResource backend_resource_;
+  std::shared_ptr<BackendResource> backend_resource_{nullptr};
 };
+
 }  // namespace pir
 
 class CompilationCache {
  public:
-  using CacheKey = std::shared_ptr<pir::OpLoweringGroup>;
+  using CacheKey = pir::FusionInfo;
   using CacheValue = std::shared_ptr<pir::CompilationResult>;
 
   static CompilationCache& Instance() {
-    static CompilationCache instance;
+    thread_local static CompilationCache instance;
     return instance;
   }
 
   bool Has(const CacheKey& key) const;
   const CacheValue& Get(const CacheKey& key) const;
-  pir::CINNKernelInfo GetKernelInfo(const CacheKey& key) const;
   void Insert(const CacheKey& key, const CacheValue& value);
   void Clear();
-  size_t KeyHash(const CacheKey& key) const;
+
+  pir::CINNKernelInfo GetKernelInfo(const CacheKey& key) const;
 
  private:
   CompilationCache() = default;
   CINN_DISALLOW_COPY_AND_ASSIGN(CompilationCache);
 
-  std::unordered_map<size_t, CacheValue> cache_;
+  std::unordered_map<CacheKey, CacheValue> cache_;
 };
 
 }  // namespace cinn::hlir::framework
@@ -42,15 +42,9 @@ std::string GroupCompilationContext::PrintPredicate2Funcs() const {
   return ss.str();
 }
 
-void CompilationTask::operator()() {
-  VLOG(4) << "Run Compilation Task for : " << context_->group_.get();
-  if (CompilationCache::Instance().Has(context_->group_)) {
-    VLOG(4) << "Found cached kernel info for group: "
-            << context_->group_->FuncName();
-    return;
-  }
+std::shared_ptr<pir::CompilationResult> CompilationTask::operator()() {
   Lowering();
-  CodegenAndJit();
+  return CodegenAndJit();
 }
 
 void CompilationTask::Lowering() {
@@ -62,7 +56,7 @@ void CompilationTask::Lowering() {
                              /* apply pass = */ true));
 }
 
-void CompilationTask::CodegenAndJit() {
+std::shared_ptr<pir::CompilationResult> CompilationTask::CodegenAndJit() {
   ir::Module::Builder builder(cinn::common::UniqName("module"),
                               context_->target_);
   CHECK_EQ(context_->predicates_.size(), context_->lowered_funcs_.size());
@@ -74,27 +68,22 @@ void CompilationTask::CodegenAndJit() {
   }
   builder.SetInferShapeFunc(context_->infer_shape_lowered_func_);
   ir::Module ir_module = builder.Build();
-  BuildPirCINNKernelInfo(ir_module);
-}
-
-pir::CINNKernelInfo CompilationTask::GetCINNKernelInfo() {
-  if (!CompilationCache::Instance().Has(context_->group_)) {
-    PADDLE_THROW(phi::errors::NotFound(
-        "Kernel info has been cached for current group."));
-  }
-  return CompilationCache::Instance().GetKernelInfo(context_->group_);
+  return BuildPirCINNKernelInfo(ir_module);
 }
 
-void CompilationTask::BuildPirCINNKernelInfo(const ir::Module& module) {
+std::shared_ptr<pir::CompilationResult> CompilationTask::BuildPirCINNKernelInfo(
+    const ir::Module& module) {
   auto compilation_result =
       std::make_shared<pir::CompilationResult>(context_->target_);
-  pir::BackendResource& backend_resource =
-      compilation_result->MutableBackendResource();
-  backend_resource.GetBackendCompiler()->Build(module, "");
-  backend_resource.SetHostFnName(context_->group_->FuncName());
-  backend_resource.SetInferFnName(context_->group_->FuncName() +
-                                  "_infer_shape");
-  CompilationCache::Instance().Insert(context_->group_, compilation_result);
+  auto backend_resource = std::make_shared<pir::BackendResource>(
+      context_->target_,
+      context_->group_->FuncName(),
+      context_->group_->FuncName() + "_infer_shape",
+      context_->group_->int_args_map());
+  VLOG(5) << "Start to compile module into cuda kernel...";
+  backend_resource->GetBackendCompiler()->Build(module, "");
+  compilation_result->SetBackendResource(backend_resource);
+  return compilation_result;
 }
 
 }  // namespace framework
 
@@ -50,14 +50,13 @@ class CompilationTask {
   explicit CompilationTask(GroupCompilationContext* context)
       : context_(context) {}
 
-  void operator()();
-  pir::CINNKernelInfo GetCINNKernelInfo();
+  std::shared_ptr<pir::CompilationResult> operator()();
 
  private:
   void Lowering();
-  void CodegenAndJit();
-  std::unique_ptr<Instruction> BuildInstruction();
-  void BuildPirCINNKernelInfo(const ir::Module& module);
+  std::shared_ptr<pir::CompilationResult> CodegenAndJit();
+  std::shared_ptr<pir::CompilationResult> BuildPirCINNKernelInfo(
+      const ir::Module& module);
 
   GroupCompilationContext* context_;
 };