PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.cc
Copy file name to clipboard
+450-276Lines changed: 450 additions & 276 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.cc
Copy file name to clipboard
+450-276Lines changed: 450 additions & 276 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/CMakeLists.txt
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions b/‎paddle/cinn/hlir/framework/pir/CMakeLists.txt
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_cache.cc
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_cache.cc
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_cache.h
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_cache.h
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_task.cc
Copy file name to clipboard
+25-26Lines changed: 25 additions & 26 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_task.cc
Copy file name to clipboard
+25-26Lines changed: 25 additions & 26 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_task.h
Copy file name to clipboard
+6-11Lines changed: 6 additions & 11 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_task.h
Copy file name to clipboard
+6-11Lines changed: 6 additions & 11 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/op_lowering_group.h
Copy file name to clipboard
+16-1Lines changed: 16 additions & 1 deletion b/‎paddle/cinn/hlir/framework/pir/op_lowering_group.h
Copy file name to clipboard
+16-1Lines changed: 16 additions & 1 deletion
@@ -8,6 +8,7 @@ gather_srcs(
   op_lowering_impl.cc
   op_mapper.cc
   op_lowering_util.cc
+  compilation_task.cc
+  compilation_cache.cc
   trivial_op_impl.cc
-  trivial_op_util.cc
-  compilation_task.cc)
+  trivial_op_util.cc)
@@ -0,0 +1,102 @@
+// 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/framework/pir/compilation_cache.h"
+#include "paddle/cinn/hlir/framework/pir/op_lowering_group.h"
+
+#include "paddle/common/enforce.h"
+
+namespace cinn::hlir::framework {
+
+namespace pir {
+void* BackendResource::GetHostFuncPtr() const {
+  VLOG(4) << "Lookup kernel name: " << host_fn_name_;
+  void* ptr = backend_compiler_->Lookup(host_fn_name_);
+  PADDLE_ENFORCE_NOT_NULL(ptr,
+                          phi::errors::InvalidArgument(
+                              "Can't find kernel function %s", host_fn_name_));
+  return ptr;
+}
+
+void* BackendResource::GetInferFuncPtr() const {
+  VLOG(4) << "Lookup infer shape fn name: " << infer_fn_name_;
+  void* ptr = backend_compiler_->Lookup(infer_fn_name_);
+  PADDLE_ENFORCE_NOT_NULL(
+      ptr,
+      phi::errors::InvalidArgument("Can't find infer shape function %s",
+                                   infer_fn_name_));
+  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 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();
+  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;
+  return has_existed;
+}
+
+const CompilationCache::CacheValue& CompilationCache::Get(
+    const CacheKey& key) const {
+  PADDLE_ENFORCE_EQ(
+      Has(key),
+      true,
+      phi::errors::NotFound("%s is not in CompliatonCache.", key->FuncName()));
+  return cache_.at(KeyHash(key));
+}
+
+pir::CINNKernelInfo CompilationCache::GetKernelInfo(const CacheKey& key) const {
+  return Get(key)->GetKernelInfo(key);
+}
+
+void CompilationCache::Insert(const CacheKey& key, const CacheValue& value) {
+  VLOG(6) << "Insert CompilationCache for: " << key->FuncName();
+  cache_.insert({KeyHash(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
@@ -0,0 +1,102 @@
+// 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>
+#include <unordered_map>
+#include "paddle/cinn/backends/compiler.h"
+#include "paddle/cinn/common/macros.h"
+#include "paddle/cinn/common/target.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
+
+namespace cinn::hlir::framework {
+
+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) {
+    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);
+
+ 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_;
+};
+
+class CompilationResult final {
+ public:
+  explicit CompilationResult(const Target& target)
+      : target_(target), backend_resource_(target) {}
+
+  BackendResource& MutableBackendResource() { return backend_resource_; }
+  const BackendResource& GetBackendResource() const {
+    return backend_resource_;
+  }
+  pir::CINNKernelInfo GetKernelInfo(
+      const std::shared_ptr<pir::OpLoweringGroup>& group) {
+    return backend_resource_.GernerateKernelInfo(group);
+  }
+
+ private:
+  Target target_;
+  BackendResource backend_resource_;
+};
+}  // namespace pir
+
+class CompilationCache {
+ public:
+  using CacheKey = std::shared_ptr<pir::OpLoweringGroup>;
+  using CacheValue = std::shared_ptr<pir::CompilationResult>;
+
+  static CompilationCache& Instance() {
+    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;
+
+ private:
+  CompilationCache() = default;
+  CINN_DISALLOW_COPY_AND_ASSIGN(CompilationCache);
+
+  std::unordered_map<size_t, CacheValue> cache_;
+};
+
+}  // namespace cinn::hlir::framework
@@ -17,7 +17,7 @@
 #include "paddle/cinn/hlir/framework/pir/compilation_task.h"
 #include "paddle/cinn/common/target.h"
 #include "paddle/cinn/hlir/framework/op_lowering.h"
-#include "paddle/cinn/ir/module.h"
+#include "paddle/common/enforce.h"
 
 namespace cinn {
 namespace hlir {
@@ -29,7 +29,6 @@ void GroupCompilationContext::SetLoweredFuncs(
        funcs.predicate2funcs) {
     predicates_.push_back(std::move(predicate2func.first));
     lowered_funcs_.push_back(std::move(predicate2func.second));
-    ++func_size_;
   }
   infer_shape_lowered_func_ = std::move(funcs.infer_shape_func);
 }
@@ -43,15 +42,13 @@ std::string GroupCompilationContext::PrintPredicate2Funcs() const {
   return ss.str();
 }
 
-void* GroupCompilationContext::FuncPtr() {
-  return backend_compiler_->Lookup(host_func_name_);
-}
-
-std::shared_ptr<backends::Compiler> GroupCompilationContext::BackendCompiler() {
-  return backend_compiler_;
-}
-
 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;
+  }
   Lowering();
   CodegenAndJit();
 }
@@ -77,25 +74,27 @@ void CompilationTask::CodegenAndJit() {
   }
   builder.SetInferShapeFunc(context_->infer_shape_lowered_func_);
   ir::Module ir_module = builder.Build();
+  BuildPirCINNKernelInfo(ir_module);
+}
 
-  context_->backend_compiler_ = backends::Compiler::Create(context_->target_);
-  context_->backend_compiler_->Build(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_);
 }
 
-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);
-  auto* infer_shape_fn_ptr =
-      context_->backend_compiler_->Lookup(fn_name + "_infer_shape");
-  CHECK(infer_shape_fn_ptr);
-  pir::CINNKernelInfo cinn_kernel_info;
-  cinn_kernel_info.fn_name = fn_name;
-  cinn_kernel_info.fn_ptr = fn_ptr;
-  cinn_kernel_info.infer_shape_fn_ptr = infer_shape_fn_ptr;
-  cinn_kernel_info.int_args_map = context_->group_->int_args_map();
-  return cinn_kernel_info;
+void 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);
 }
 
 }  // namespace framework
 
@@ -16,13 +16,16 @@
 #include "paddle/cinn/backends/compiler.h"
 #include "paddle/cinn/common/target.h"
 #include "paddle/cinn/hlir/framework/instruction.h"
+#include "paddle/cinn/hlir/framework/pir/compilation_cache.h"
 #include "paddle/cinn/hlir/framework/pir/op_lowering_impl.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/cinn/ir/group_schedule/base_group_scheduler.h"
+#include "paddle/cinn/ir/module.h"
 
 namespace cinn {
 namespace hlir {
 namespace framework {
+class CompilationTask;
 
 class GroupCompilationContext {
  public:
@@ -32,23 +35,14 @@ class GroupCompilationContext {
 
   void SetLoweredFuncs(BucketLoweredFuncsWrapper&& funcs);
   std::string PrintPredicate2Funcs() const;
-  void* FuncPtr();
-  std::shared_ptr<backends::Compiler> BackendCompiler();
 
  private:
   friend class CompilationTask;
-
   const Target& target_;
   const pir::OpLoweringGroupPtr& group_;
-
-  size_t func_size_ = 0;
   std::vector<ir::SymbolicPredicate> predicates_;
   std::vector<ir::LoweredFunc> lowered_funcs_;
   ir::LoweredFunc infer_shape_lowered_func_;
-  std::string host_func_name_;
-  std::string host_code_;
-  std::vector<std::string> device_code_;
-  std::shared_ptr<backends::Compiler> backend_compiler_;
 };
 
 class CompilationTask {
@@ -57,13 +51,14 @@ class CompilationTask {
       : context_(context) {}
 
   void operator()();
+  pir::CINNKernelInfo GetCINNKernelInfo();
 
+ private:
   void Lowering();
   void CodegenAndJit();
   std::unique_ptr<Instruction> BuildInstruction();
-  pir::CINNKernelInfo BuildPirCINNKernelInfo();
+  void BuildPirCINNKernelInfo(const ir::Module& module);
 
- private:
   GroupCompilationContext* context_;
 };
 
 
@@ -19,6 +19,7 @@
 #include <vector>
 #include "glog/logging.h"
 
+#include "paddle/cinn/common/context.h"
 #include "paddle/cinn/hlir/framework/op.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/pir/include/core/builtin_type_interfaces.h"
@@ -47,6 +48,20 @@ class OpLoweringGroup {
   explicit OpLoweringGroup(std::initializer_list<::pir::Operation*> group_ops)
       : ops_(group_ops) {}
 
+  struct SharedGroupHasher {
+    size_t operator()(
+        const std::shared_ptr<OpLoweringGroup>& group) const noexcept {
+      return std::hash<std::string>()(group->group_id());
+    }
+  };
+  struct SharedGroupComparator {
+    bool operator()(
+        const std::shared_ptr<OpLoweringGroup>& first,
+        const std::shared_ptr<OpLoweringGroup>& second) const noexcept {
+      return first->group_id() == second->group_id();
+    }
+  };
+
   std::vector<::pir::Value> GetGroupOutputValues() const {
     std::unordered_set<::pir::Operation*> group_ops_set(this->ops_.begin(),
                                                         this->ops_.end());
@@ -265,7 +280,7 @@ class OpLoweringGroup {
 
  private:
   // group id, consisted of op's id.
-  std::string group_id_{""};
+  std::string group_id_{common::UniqName("group_")};
   // op in this group
   std::vector<::pir::Operation*> ops_;
   // output ops of the group.