[PIR+CINN]Fix CompilationCache key not Found while sym_dim is changed (#63871)

Aurelius84 · web-flow · commit 418416db913f · 2024-04-26T11:00:20.000+08:00
* [PIR+CINN]Fix CompilationCache key not Found while sym_dim is changed

* fix typo

* fix typo
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.cc b/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.cc
@@ -64,13 +64,13 @@ void CompileGroupToJitKernelOp(
   auto op_attr_map = CompileGroupAsOpAttribute(group_list);
   VLOG(4) << "The size of group_map is : " << group_map->size();
   for (auto& [block, group] : *group_map) {
+    auto& yield_op = block->back();
+    CHECK(yield_op.isa<pir::YieldOp>()) << "Last op of block should be yield";
     std::vector<pir::Type> output_types;
-    const auto& group_output_values = group->output_values();
+    const auto& group_output_values = yield_op.operands_source();
     for (size_t i = 0; i < group_output_values.size(); ++i) {
       output_types.push_back(group_output_values[i].type());
     }
-    auto& yield_op = block->back();
-    CHECK(yield_op.isa<pir::YieldOp>()) << "Last op of block should be yield";
     rewriter.set_insertion_point(&yield_op);
     const auto& group_inputs = GetBlockOutsideInput(group->ops());
     auto jit_kernel_op = rewriter.Build<cinn::dialect::JitKernelOp>(
@@ -283,6 +283,7 @@ void SetLeafBlockByGroupView(
   }
 
   auto new_group = CloneGroup(origin_group, block, &ir_mapping);
+  new_group->SetIsBroadcastLeaf(true);
   PADDLE_ENFORCE_EQ(
       origin_group->ops().size(),
       new_group->ops().size(),
diff --git a/paddle/cinn/hlir/framework/pir/compilation_task.cc b/paddle/cinn/hlir/framework/pir/compilation_task.cc
@@ -48,12 +48,14 @@ std::shared_ptr<pir::CompilationResult> CompilationTask::operator()() {
 }
 
 void CompilationTask::Lowering() {
+  VLOG(5) << "Begin to lowering group: " << *context_->group_;
   auto op_lowerer = CreateOpLowerer<pir::OpLoweringGroupPtr>(context_->target_);
   context_->SetLoweredFuncs(
       op_lowerer.BucketLower(context_->group_,
                              /* apply op schedule = */ false,
                              /* apply group schedule = */ true,
                              /* apply pass = */ true));
+  VLOG(5) << "End to lowering: " << context_->PrintPredicate2Funcs();
 }
 
 std::shared_ptr<pir::CompilationResult> CompilationTask::CodegenAndJit() {
diff --git a/paddle/cinn/hlir/framework/pir/fusion_info.cc b/paddle/cinn/hlir/framework/pir/fusion_info.cc
@@ -174,7 +174,9 @@ void FusionInfo::ParseInputDimExprs(const OpLoweringGroup& group) {
   };
 
   for (const auto& value : group.GetInputOpValues()) {
-    if (!TryGetDimExprsFromGroup(value)) {
+    if (group.IsBroadcastLeaf()) {
+      TryGetDimExprsFromGroup(value);
+    } else {
       TryeGetDimExprsFromGlobal(value);
     }
   }
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_group.h b/paddle/cinn/hlir/framework/pir/op_lowering_group.h
@@ -74,6 +74,11 @@ class OpLoweringGroup {
     }
   }
 
+  bool IsBroadcastLeaf() const { return is_broadcast_leaf_; }
+  void SetIsBroadcastLeaf(bool is_broadcast_leaf) {
+    is_broadcast_leaf_ = is_broadcast_leaf;
+  }
+
   const std::vector<::pir::Operation*>& ops() const { return ops_; }
   std::vector<::pir::Operation*>& mut_ops() { return ops_; }
   void SetOps(const std::vector<::pir::Operation*>& new_ops) { ops_ = new_ops; }
@@ -198,6 +203,7 @@ class OpLoweringGroup {
   std::unordered_set<::pir::Operation*> output_ops_;
   // op pattern kind.
   OpPatternKind op_pattern_kind_{kElementWise};
+  bool is_broadcast_leaf_{false};
 
   std::vector<std::string> input_names_;
   std::vector<std::string> output_names_;
diff --git a/paddle/cinn/hlir/framework/pir_compiler.cc b/paddle/cinn/hlir/framework/pir_compiler.cc
@@ -20,6 +20,7 @@
 #include "paddle/common/flags.h"
 
 PD_DECLARE_bool(enable_cinn_compile_cache);
+PD_DECLARE_int64(cinn_compile_thread_num);
 
 namespace cinn::hlir::framework {
 
@@ -52,16 +53,26 @@ class CompilationContextMapper {
   bool is_finalized_{false};
 };
 
+static size_t GetThreadNum(size_t task_size) {
+  size_t thread_size = task_size;
+  if (!FLAGS_enable_cinn_compile_cache) {
+    thread_size = 1;
+  } else if (FLAGS_cinn_compile_thread_num > 0) {
+    thread_size = FLAGS_cinn_compile_thread_num;
+  }
+  return thread_size;
+}
+
 std::vector<pir::CINNKernelInfo> PirCompiler::Build(
     const std::vector<pir::OpLoweringGroupPtr>& groups) {
   CompilationContextMapper ctx_mapper(target_, groups);
   auto& group_compilation_contexts = ctx_mapper.UniqueCompilationContexts();
   auto& compilation_results = ctx_mapper.MutableCompilationResult();
 
   const size_t task_size = group_compilation_contexts.size();
-  const size_t thread_size = FLAGS_enable_cinn_compile_cache ? task_size : 1;
+  const size_t thread_size = GetThreadNum(task_size);
   VLOG(5) << "Found " << task_size << " new groups parsed from "
-          << groups.size();
+          << groups.size() << " and compiles with " << thread_size;
   if (task_size > 0) {
     auto worker_fn = [&](int index) {
       CompilationTask task(&group_compilation_contexts[index]);
@@ -89,6 +100,8 @@ void CompilationContextMapper::Construct(
 
   for (size_t i = 0; i < groups.size(); ++i) {
     fusion_infos_.emplace_back(*groups[i]);
+    VLOG(5) << "Construct FusionInfo: " << fusion_infos_[i]
+            << " for group: " << *groups[i];
     // If FLAGS_enable_cinn_compile_cache=False, Cache strategy will not take
     // effects.
     if (IsNewAndUnique(fusion_infos_[i]) || !FLAGS_enable_cinn_compile_cache) {
diff --git a/paddle/common/flags.cc b/paddle/common/flags.cc
@@ -1037,7 +1037,17 @@ PHI_DEFINE_EXPORTED_bool(
     enable_cinn_compile_cache,
     true,
     "It controls whether to enable cinn compilation cache.");
-
+/*
+ * CINN related FLAG
+ * Name: FLAGS_deny_cinn_ops
+ * Since Version: 3.0 Beta
+ * Value Range: bool, default=-1
+ * Example: FLAGS_cinn_compile_thread_nume=8
+ */
+PHI_DEFINE_EXPORTED_int64(
+    cinn_compile_thread_num,
+    -1,
+    "It controls how many thread numbers applying compilation cache.");
 /*
  * CINN related FLAG
  * Name: FLAGS_enable_interpretercore_launch_cinn

Original file line number	Diff line number	Diff line change
`@@ -48,12 +48,14 @@ std::shared_ptr<pir::CompilationResult> CompilationTask::operator()() {`
`48`	`48`	`}`
`49`	`49`
`50`	`50`	`void CompilationTask::Lowering() {`
	`51`	`+ VLOG(5) << "Begin to lowering group: " << *context_->group_;`
`51`	`52`	`auto op_lowerer = CreateOpLowerer<pir::OpLoweringGroupPtr>(context_->target_);`
`52`	`53`	`context_->SetLoweredFuncs(`
`53`	`54`	`op_lowerer.BucketLower(context_->group_,`
`54`	`55`	`/* apply op schedule = */ false,`
`55`	`56`	`/* apply group schedule = */ true,`
`56`	`57`	`/* apply pass = */ true));`
	`58`	`+ VLOG(5) << "End to lowering: " << context_->PrintPredicate2Funcs();`
`57`	`59`	`}`
`58`	`60`
`59`	`61`	`std::shared_ptr<pir::CompilationResult> CompilationTask::CodegenAndJit() {`
Original file line number	Diff line number	Diff line change
`@@ -174,7 +174,9 @@ void FusionInfo::ParseInputDimExprs(const OpLoweringGroup& group) {`
`174`	`174`	`};`
`175`	`175`
`176`	`176`	`for (const auto& value : group.GetInputOpValues()) {`
`177`		`- if (!TryGetDimExprsFromGroup(value)) {`
	`177`	`+ if (group.IsBroadcastLeaf()) {`
	`178`	`+ TryGetDimExprsFromGroup(value);`
	`179`	`+ } else {`
`178`	`180`	`TryeGetDimExprsFromGlobal(value);`
`179`	`181`	`}`
`180`	`182`	`}`