PaddlePaddle
diff --git a/‎paddle/cinn/backends/codegen_device_util.cc
Copy file name to clipboard
+8-5Lines changed: 8 additions & 5 deletions b/‎paddle/cinn/backends/codegen_device_util.cc
Copy file name to clipboard
+8-5Lines changed: 8 additions & 5 deletions
diff --git a/‎paddle/cinn/backends/compiler.cc
Copy file name to clipboard
+19-22Lines changed: 19 additions & 22 deletions b/‎paddle/cinn/backends/compiler.cc
Copy file name to clipboard
+19-22Lines changed: 19 additions & 22 deletions
diff --git a/‎paddle/cinn/backends/compiler.h
Copy file name to clipboard
+7-4Lines changed: 7 additions & 4 deletions b/‎paddle/cinn/backends/compiler.h
Copy file name to clipboard
+7-4Lines changed: 7 additions & 4 deletions
diff --git a/‎paddle/cinn/backends/llvm/execution_engine.cc
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion b/‎paddle/cinn/backends/llvm/execution_engine.cc
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/framework/pir/compilation_task.cc
Copy file name to clipboard
+183-40Lines changed: 183 additions & 40 deletions b/‎paddle/cinn/hlir/framework/pir/compilation_task.cc
Copy file name to clipboard
+183-40Lines changed: 183 additions & 40 deletions
@@ -51,7 +51,7 @@ ir::Module CreateSwitchWithBroadcastConditionModule(
       ir::Argument(kernel_args_num, ir::Argument::IO::kInput),
       ir::Argument(tensor_shape_args, ir::Argument::IO::kOutput)};
 
-  const auto &CreateSymbolArgDefines = [&]() -> std::vector<ir::Expr> {
+  const auto &symbolic_arg_define = [&]() -> std::vector<ir::Expr> {
     std::vector<ir::Expr> arg_defs;
     for (const auto &item : symbolic_shape_var_index) {
       ir::Expr call_get_value_in_kernel_args =
@@ -68,13 +68,13 @@ ir::Module CreateSwitchWithBroadcastConditionModule(
       arg_defs.push_back(stmt);
     }
     return arg_defs;
-  };
+  }();
 
   const auto &CreateSwitchFunction =
       [&](std::vector<ir::Argument> func_arguments,
           const std::vector<ir::Expr> &read_args,
           std::string name_extend) -> ir::Expr {
-    std::vector<ir::Expr> body_stmts(CreateSymbolArgDefines());
+    std::vector<ir::Expr> body_stmts(symbolic_arg_define);
     for (int i = 0; i < broadcast_conditions.size(); ++i) {
       ir::Expr callee = ir::Call::Make(Void(),
                                        case_func_names[i] + name_extend,
@@ -113,8 +113,11 @@ ir::Module CreateSwitchWithBroadcastConditionModule(
   module_builder.AddFunctionWithoutOptim(
       infer_shape_func_caller.as_lowered_func_ref());
   // no need cx86 func
-  ir::Expr cx86_func_caller = ir::_LoweredFunc_::Make(
-      wrapper_func_name + "_CX86", host_func_arguments, ir::Expr(), {});
+  ir::Expr cx86_func_caller =
+      ir::_LoweredFunc_::Make(wrapper_func_name + "_CX86",
+                              host_func_arguments,
+                              ir::Block::Make({}),
+                              {});
   module_builder.AddFunctionWithoutOptim(
       cx86_func_caller.as_lowered_func_ref());
   return module_builder.Build();
 
@@ -227,32 +227,30 @@ void SourceCodePrint::write(const std::string& source_code) {
   }
 }
 
-void Compiler::Build(const Module& module,
-                     const std::string& code,
-                     const bool end) {
-  auto PatternMatch = adt::match{
+void Compiler::Build(const Module& module, const std::string& code) {
+  target_.arch.Match(
       [&](common::UnknownArch) { CINN_NOT_IMPLEMENTED; },
       [&](common::X86Arch) { CompileX86Module(module); },
       [&](common::ARMArch) { CINN_NOT_IMPLEMENTED; },
       [&](common::NVGPUArch) { CompileCudaModule(module, code); },
-      [&](common::HygonDCUArchHIP) { CompileHipModule(module, code); }};
-  std::visit(PatternMatch, target_.arch.variant());
-  if (end) {
-    RegisterDeviceModuleSymbol();
-    engine_->AddSelfModule();
-  }
+      [&](common::HygonDCUArchHIP) { CompileHipModule(module, code); });
 }
 
-void Compiler::AppendCX86(const Module& module, const bool end) {
+void Compiler::AppendCX86(const Module& module) {
   VLOG(3) << "Start Compiler::BuildCX86" << module;
   CompileX86Module(module);
-  if (end) {
-    RegisterDeviceModuleSymbol();
-    engine_->AddSelfModule();
-  }
   VLOG(3) << "Over Compiler::BuildCX86";
 }
 
+void Compiler::AppendBroadcastSwitchModule(const ir::Module& module) {
+  engine_->Link<CodeGenSwitchHost>(module);
+}
+
+void Compiler::EndCompile() {
+  RegisterDeviceModuleSymbol();
+  engine_->AddSelfModule();
+}
+
 std::string Compiler::GetSourceCode(const ir::Module& module) {
   return target_.arch.Match(
       [&](common::UnknownArch) -> std::string { CINN_NOT_IMPLEMENTED; },
@@ -305,13 +303,12 @@ std::string GetFileContent(const std::string& path) {
 }  // namespace
 
 void Compiler::RegisterDeviceModuleSymbol() {
-  auto PatternMatch =
-      adt::match{[&](common::UnknownArch) { CINN_NOT_IMPLEMENTED; },
-                 [&](common::X86Arch) { return; },
-                 [&](common::ARMArch) { return; },
-                 [&](common::NVGPUArch) { RegisterCudaModuleSymbol(); },
-                 [&](common::HygonDCUArchHIP) { CINN_NOT_IMPLEMENTED; }};
-  return std::visit(PatternMatch, target_.arch.variant());
+  return target_.arch.Match(
+      [&](common::UnknownArch) { CINN_NOT_IMPLEMENTED; },
+      [&](common::X86Arch) { return; },
+      [&](common::ARMArch) { return; },
+      [&](common::NVGPUArch) { RegisterCudaModuleSymbol(); },
+      [&](common::HygonDCUArchHIP) { CINN_NOT_IMPLEMENTED; });
 }
 
 void Compiler::RegisterCudaModuleSymbol() {
 
@@ -103,10 +103,13 @@ class Compiler final {
   /**
    * Compile and link to a CINN module.
    */
-  void Build(const ir::Module& module,
-             const std::string& code = "",
-             const bool end = true);
-  void AppendCX86(const ir::Module& module, const bool end = true);
+  void Build(const ir::Module& module, const std::string& code = "");
+
+  void AppendCX86(const ir::Module& module);
+
+  void AppendBroadcastSwitchModule(const ir::Module& module);
+
+  void EndCompile();
 
   void ExportObject(const std::string& path);
 
 
@@ -231,7 +231,7 @@ void ExecutionEngine::RegisterModuleRuntimeSymbols(
   module_symbols_ = std::forward<RuntimeSymbols>(module_symbols);
   auto *session = &jit_->getExecutionSession();
   for (const auto &sym : module_symbols_.All()) {
-    VLOG(0) << "Add symbol: {" << sym.first << ":" << sym.second << "}";
+    VLOG(3) << "Add symbol: {" << sym.first << ":" << sym.second << "}";
     llvm::cantFail(jit_->define(llvm::orc::absoluteSymbols(
         {{session->intern(sym.first),
           {llvm::pointerToJITTargetAddress(sym.second),
@@ -276,5 +276,7 @@ void ExecutionEngine::RegisterGlobalRuntimeSymbols() {
 template void ExecutionEngine::Link<CodeGenLLVM>(const ir::Module &module);
 template void ExecutionEngine::Link<CodeGenX86>(const ir::Module &module);
 template void ExecutionEngine::Link<CodeGenCUDA_Host>(const ir::Module &module);
+template void ExecutionEngine::Link<CodeGenSwitchHost>(
+    const ir::Module &module);
 
 }  // namespace cinn::backends
@@ -15,8 +15,13 @@
 #pragma once
 
 #include "paddle/cinn/hlir/framework/pir/compilation_task.h"
+
+#include "paddle/cinn/backends/codegen_device_util.h"
+#include "paddle/cinn/common/dim_expr_converter.h"
 #include "paddle/cinn/common/target.h"
 #include "paddle/cinn/hlir/framework/op_lowering.h"
+#include "paddle/cinn/hlir/framework/pir/op_lowering_group.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/common/enforce.h"
 namespace cinn {
 namespace hlir {
@@ -47,6 +52,99 @@ std::string GroupCompilationContext::PrintPredicate2Funcs() const {
   return ss.str();
 }
 
+void GroupCompilationContext::PrepareModuleBuilder() {
+  PADDLE_ENFORCE_EQ(predicates_.size(),
+                    lowered_funcs_.size(),
+                    phi::errors::InvalidArgument(
+                        "The size of predicates and lowered_funcs should be "
+                        "the same."));
+  PADDLE_ENFORCE_EQ(predicates_.size(),
+                    priorities_.size(),
+                    phi::errors::InvalidArgument(
+                        "The size of predicates and priorites should be "
+                        "the same."));
+  for (const ir::Expr& predicate : predicates_) {
+    module_builder_.AddPredicate(predicate);
+  }
+  for (const ir::LoweredFunc& func : lowered_funcs_) {
+    module_builder_.AddFunction(func);
+  }
+  for (int& priority : priorities_) {
+    module_builder_.AddPriority(priority);
+  }
+  module_builder_.SetInferShapeFunc(infer_shape_lowered_func_);
+
+  PADDLE_ENFORCE_EQ(CX86_predicates_.size(),
+                    CX86_lowered_funcs_.size(),
+                    phi::errors::InvalidArgument(
+                        "The size of predicates and lowered_funcs should be "
+                        "the same."));
+  for (const ir::Expr& predicate : CX86_predicates_) {
+    CX86_module_builder_.AddPredicate(predicate);
+  }
+  for (const ir::LoweredFunc& func : CX86_lowered_funcs_) {
+    CX86_module_builder_.AddFunction(func);
+  }
+}
+
+/**
+ * For functions belonging to different broadcast groups, int args and the name
+ * of the tensor args may be variate, but the number of the tensor args should
+ * be fixed. So we need to unify the tensor args and int args. For exmaple,
+ * func1(_var, _var_1, S4, S5); func2(_var, _var_2, S1) would be unified to
+ * func1(_var, _var_1, S4, S5, S1); func2(_var, _var_2, S4, S5, S1).
+ */
+void UnifyBroadcastGroupFuncArgs(
+    std::vector<GroupCompilationContext>* contexts,
+    pir::OpLoweringGroupPtr origin_group,
+    std::unordered_map<int, ir::Var>* symbolic_shape_var_index) {
+  std::unordered_map<ir::Var, pir::CINNKernelInfo::ArgDimIdx> new_args_map;
+  std::vector<ir::Argument> new_args_vec;
+  int total_args_num = 0;
+
+  const auto& AddTensorArgs = [&](GroupCompilationContext& context) {
+    const auto& func_args = context.lowered_funcs_[0]->args;
+    const auto& origin_int_args = context.group_->int_args_map();
+    for (size_t arg_idx = 0; arg_idx < func_args.size(); ++arg_idx) {
+      if (func_args[arg_idx].is_var()) {
+        new_args_map[func_args[arg_idx].var_arg()] =
+            origin_int_args.at(arg_idx);
+      } else {
+        new_args_vec.emplace_back(func_args[arg_idx]);
+      }
+    }
+    for (ir::LoweredFunc& func : context.lowered_funcs_) {
+      func->args = new_args_vec;
+    }
+  };
+  for (size_t i = 0; i < contexts->size(); ++i) {
+    AddTensorArgs((*contexts)[i]);
+    if (i == 0) total_args_num += new_args_vec.size();
+    new_args_vec.clear();
+  }
+
+  origin_group->mut_int_args_map().clear();
+  const auto& new_int_args_vec = [&]() -> std::vector<ir::Argument> {
+    std::vector<ir::Argument> res;
+    for (const auto& [arg, idx_info] : new_args_map) {
+      symbolic_shape_var_index->insert({total_args_num, arg});
+      origin_group->mut_int_args_map()[total_args_num++] = idx_info;
+      res.emplace_back(ir::Argument{arg});
+    }
+    return res;
+  }();
+
+  const auto& AddUnifiedIntArgs = [&](GroupCompilationContext& context) {
+    for (ir::LoweredFunc& func : context.lowered_funcs_) {
+      func->args.insert(
+          func->args.end(), new_int_args_vec.begin(), new_int_args_vec.end());
+    }
+  };
+  for (int i = 0; i < contexts->size(); ++i) {
+    AddUnifiedIntArgs((*contexts)[i]);
+  }
+}
+
 std::shared_ptr<pir::CompilationResult> CompilationTask::operator()() {
   Lowering();
   return CodegenAndJit();
@@ -60,49 +158,55 @@ void CompilationTask::Lowering() {
                              /* 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() {
-  ir::Module::Builder builder(cinn::common::UniqName("module"),
-                              context_->target_);
-  PADDLE_ENFORCE_EQ(context_->predicates_.size(),
-                    context_->lowered_funcs_.size(),
-                    phi::errors::InvalidArgument(
-                        "The size of predicates and lowered_funcs should be "
-                        "the same."));
-  PADDLE_ENFORCE_EQ(context_->predicates_.size(),
-                    context_->priorities_.size(),
-                    phi::errors::InvalidArgument(
-                        "The size of predicates and priorites should be "
-                        "the same."));
-  for (const ir::Expr& predicate : context_->predicates_) {
-    builder.AddPredicate(predicate);
-  }
-  for (const ir::LoweredFunc& func : context_->lowered_funcs_) {
-    builder.AddFunction(func);
-  }
-  for (int& priority : context_->priorities_) {
-    builder.AddPriority(priority);
-  }
-  builder.SetInferShapeFunc(context_->infer_shape_lowered_func_);
-  ir::Module ir_module = builder.Build();
+  if (context_->group_->IsBroadcastLeaf()) {
+    const auto& broadcast_condition_dimexprs =
+        context_->group_->GetBroadcastConditions();
 
-  ir::Module::Builder builder_CX86(cinn::common::UniqName("module"),
-                                   common::DefaultHostTarget());
-  PADDLE_ENFORCE_EQ(context_->CX86_predicates_.size(),
-                    context_->CX86_lowered_funcs_.size(),
-                    phi::errors::InvalidArgument(
-                        "The size of predicates and lowered_funcs should be "
-                        "the same."));
-  for (const ir::Expr& predicate : context_->CX86_predicates_) {
-    builder_CX86.AddPredicate(predicate);
-  }
-  for (const ir::LoweredFunc& func : context_->CX86_lowered_funcs_) {
-    builder_CX86.AddFunction(func);
+    using BranchType = pir::OpLoweringGroup::BranchType;
+    const auto& GetSingleBranchExprFromBroadcastCond =
+        [](const symbol::Broadcastable<symbol::DimExpr>& broadcast_expr,
+           const BranchType& condition) -> ir::Expr {
+      const auto& expr_converter = common::DimExprConverter();
+      ir::Expr lhs = expr_converter.ConvertToIrExpr(broadcast_expr->lhs);
+      ir::Expr rhs = expr_converter.ConvertToIrExpr(broadcast_expr->rhs);
+      ir::Expr one = ir::Expr(1);
+      ir::Expr condition_expr;
+      if (condition == BranchType::LHS_EQ_RHS) {
+        condition_expr = ir::EQ::Make(lhs, rhs);
+      } else {
+        condition_expr = ir::NE::Make(lhs, rhs);
+        ir::Expr eq_one_expr;
+        if (condition == BranchType::LHS_EQ_ONE) {
+          eq_one_expr = ir::EQ::Make(lhs, one);
+        } else {  // BranchType::RHS_EQ_ONE
+          eq_one_expr = ir::EQ::Make(rhs, one);
+        }
+        condition_expr = ir::And::Make(condition_expr, eq_one_expr);
+      }
+      return condition_expr;
+    };
+
+    const auto& ChangeBroadcastConditionToExpr = [&]() -> ir::Expr {
+      ir::Expr result = ir::Expr(true);
+      for (const auto& [broadcast_expr, condition] :
+           broadcast_condition_dimexprs) {
+        result = ir::And::Make(
+            result,
+            GetSingleBranchExprFromBroadcastCond(broadcast_expr, condition));
+      }
+      return result;
+    };
+
+    context_->broadcast_condition_ = ChangeBroadcastConditionToExpr();
   }
-  ir::Module ir_moduleCX86 = builder_CX86.Build();
+  VLOG(5) << "End to lowering: " << context_->PrintPredicate2Funcs();
+}
 
+std::shared_ptr<pir::CompilationResult> CompilationTask::CodegenAndJit() {
+  context_->PrepareModuleBuilder();
+  ir::Module ir_module = context_->module_builder_.Build();
+  ir::Module ir_moduleCX86 = context_->CX86_module_builder_.Build();
   return BuildPirCINNKernelInfo(ir_module, ir_moduleCX86);
 }
 
@@ -116,8 +220,47 @@ std::shared_ptr<pir::CompilationResult> CompilationTask::BuildPirCINNKernelInfo(
       context_->group_->FuncName() + "_infer_shape",
       context_->group_->int_args_map());
   VLOG(5) << "Start to compile module into cuda kernel...";
-  backend_resource->GetBackendCompiler()->Build(module, "", false);
+  backend_resource->GetBackendCompiler()->Build(module, "");
   backend_resource->GetBackendCompiler()->AppendCX86(CX86module);
+  backend_resource->GetBackendCompiler()->EndCompile();
+  compilation_result->SetBackendResource(backend_resource);
+  VLOG(5) << "End to compile module into cuda kernel.";
+  return compilation_result;
+}
+
+std::shared_ptr<pir::CompilationResult>
+CompilationTask::CompileBroadcastModules(
+    std::vector<GroupCompilationContext>* leaf_group_contexts,
+    const std::unordered_map<int, ir::Var>& symbolic_shape_var_index) {
+  auto compilation_result =
+      std::make_shared<pir::CompilationResult>(context_->target_);
+  auto backend_resource = std::make_shared<pir::BackendResource>(
+      context_->target_,
+      context_->group_->FuncName(),
+      context_->group_->FuncName() + "_infer_shape",
+      context_->group_->int_args_map());
+
+  std::vector<std::string> case_func_names;
+  std::vector<ir::Expr> broadcast_conditions;
+  for (auto& context : *leaf_group_contexts) {
+    context.PrepareModuleBuilder();
+    case_func_names.emplace_back(context.group_->FuncName());
+    broadcast_conditions.emplace_back(context.broadcast_condition_);
+    ir::Module ir_module = context.module_builder_.Build();
+    ir::Module ir_moduleCX86 = context.CX86_module_builder_.Build();
+    backend_resource->GetBackendCompiler()->Build(ir_module, "");
+    backend_resource->GetBackendCompiler()->AppendCX86(ir_moduleCX86);
+  }
+
+  ir::Module wrapper_module(
+      cinn::backends::CreateSwitchWithBroadcastConditionModule(
+          broadcast_conditions,
+          case_func_names,
+          context_->group_->FuncName(),
+          symbolic_shape_var_index));
+  backend_resource->GetBackendCompiler()->AppendBroadcastSwitchModule(
+      wrapper_module);
+  backend_resource->GetBackendCompiler()->EndCompile();
   compilation_result->SetBackendResource(backend_resource);
   VLOG(5) << "End to compile module into cuda kernel.";
   return compilation_result;