PaddlePaddle
diff --git a/‎paddle/cinn/backends/codegen_cuda_host.cc
Copy file name to clipboard
+154Lines changed: 154 additions & 0 deletions b/‎paddle/cinn/backends/codegen_cuda_host.cc
Copy file name to clipboard
+154Lines changed: 154 additions & 0 deletions
diff --git a/‎paddle/cinn/backends/codegen_cuda_host.h
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions b/‎paddle/cinn/backends/codegen_cuda_host.h
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions
diff --git a/‎paddle/cinn/backends/codegen_cuda_util.cc
Copy file name to clipboard
+90Lines changed: 90 additions & 0 deletions b/‎paddle/cinn/backends/codegen_cuda_util.cc
Copy file name to clipboard
+90Lines changed: 90 additions & 0 deletions
diff --git a/‎paddle/cinn/backends/codegen_cuda_util.h
Copy file name to clipboard
+52-2Lines changed: 52 additions & 2 deletions b/‎paddle/cinn/backends/codegen_cuda_util.h
Copy file name to clipboard
+52-2Lines changed: 52 additions & 2 deletions
diff --git a/‎paddle/cinn/hlir/framework/op_lowering.h
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions b/‎paddle/cinn/hlir/framework/op_lowering.h
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions
@@ -182,5 +182,159 @@ llvm::Value* CodeGenCUDA_Host::LowerGPUKernelLauncher(
   return function;
 }
 
+llvm::Value* CodeGenCUDA_Host::LowerHostFunc(const ir::_LoweredFunc_* func) {
+  // Create the function
+  // @{
+  auto* function_type = GenFunctionTypeFromCinnFunction(func, true);
+  f_ = llvm::Function::Create(
+      function_type, llvm::Function::ExternalLinkage, func->name, m_);
+  f_->setCallingConv(llvm::CallingConv::C);
+  f_->setHasUWTable();
+
+  std::vector<llvm::Value*> ll_function_args;
+  std::transform(f_->arg_begin(),
+                 f_->arg_end(),
+                 std::back_inserter(ll_function_args),
+                 [](auto& arg) { return std::addressof(arg); });
+  // @}
+
+  llvm::BasicBlock* entry = llvm::BasicBlock::Create(
+      /*Context=*/b_->getContext(),
+      /*Name=*/"entry",
+      /*Parent=*/f_,
+      /*InsertBefore=*/nullptr);
+  b_->SetInsertPoint(entry);
+  CodeGenLLVM::Visit(&func->body);
+  RetVoid();
+
+  return f_;
+}
+
+llvm::Value* CodeGenCUDA_Host::LowerCUDAKernelCall(const ir::Call* call_ir) {
+  std::vector<llvm::Value*> ll_function_args;
+  std::transform(f_->arg_begin(),
+                 f_->arg_end(),
+                 std::back_inserter(ll_function_args),
+                 [](auto& arg) { return std::addressof(arg); });
+  auto* kernel_args = ll_function_args[0];
+  auto* kernel_args_count = ll_function_args[1];
+  llvm::Value* kernel_stream = nullptr;
+  if (ll_function_args.size() == 3) {
+    kernel_stream = ll_function_args[2];
+    CHECK_EQ(kernel_stream->getType(), ll_void_p_ty());  // void* stream
+  }
+  CHECK_EQ(kernel_args->getType(), ll_void_p_ty());       // void* args
+  CHECK_EQ(kernel_args_count->getType(), ll_int32_ty());  // int32
+
+  std::unordered_map<std::string, llvm::Value*> global_args = {
+      {KERNEL_ARGS, kernel_args},
+      {KERNEL_ARGS_NUM, kernel_args_count},
+      {KERNEL_STREAM, kernel_stream}};
+
+  auto ret_type = CinnTypeToLLVMType(Void(), m_);
+  std::vector<llvm::Type*> args_type;
+  for (auto r_arg : call_ir->read_args) {
+    if (r_arg.is_var()) {
+      if (r_arg.as_var()->type().is_cpp_handle() ||
+          r_arg.as_var()->type().is_string()) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<void*>(), m_));
+      } else if (r_arg.as_var()->type().is_int(32)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<int32_t>(), m_));
+      } else {
+        CINN_NOT_IMPLEMENTED;
+      }
+    } else {
+      if (r_arg.type().is_bool()) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<bool>(), m_));
+      } else if (r_arg.type().is_uint(8)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<uint8_t>(), m_));
+      } else if (r_arg.type().is_uint(16)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<uint16_t>(), m_));
+      } else if (r_arg.type().is_uint(32)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<uint32_t>(), m_));
+      } else if (r_arg.type().is_uint(64)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<uint64_t>(), m_));
+      } else if (r_arg.type().is_int(8)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<int8_t>(), m_));
+      } else if (r_arg.type().is_int(16)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<int16_t>(), m_));
+      } else if (r_arg.type().is_int(32)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<int32_t>(), m_));
+      } else if (r_arg.type().is_int(64)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<int64_t>(), m_));
+      } else if (r_arg.type().is_float(32)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<float>(), m_));
+      } else if (r_arg.type().is_float(64)) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<double>(), m_));
+      } else if (r_arg.type().is_bfloat16()) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<bfloat16>(), m_));
+      } else if (r_arg.type().is_float16()) {
+        args_type.push_back(CinnTypeToLLVMType(type_of<float16>(), m_));
+      } else {
+        CINN_NOT_IMPLEMENTED;
+      }
+    }
+  }
+  auto func_type = llvm::FunctionType::get(ret_type, args_type, false);
+  auto call_func = m_->getOrInsertFunction(call_ir->name, func_type);
+
+  std::vector<llvm::Value*> call_args;
+  for (auto& r_arg : call_ir->read_args) {
+    if (r_arg.is_var()) {
+      if (r_arg.as_var()->type().is_string()) {
+        auto kvalue = m_->getOrInsertGlobal(r_arg.as_var()->name + "_ptr_",
+                                            b_->getInt8PtrTy());
+        call_args.push_back(b_->CreateLoad(
+            b_->getInt8PtrTy(), kvalue, r_arg.as_var()->name + "_ptr_load"));
+      } else if (r_arg.as_var()->type().is_cpp_handle() ||
+                 r_arg.as_var()->type().is_int(32)) {
+        CHECK(global_args.count(r_arg.as_var()->name));
+        call_args.push_back(global_args[r_arg.as_var()->name]);
+      } else {
+        CINN_NOT_IMPLEMENTED;
+      }
+    } else {
+      if (r_arg.type().is_bool()) {
+        call_args.push_back(b_->getInt1(r_arg.as_bool()));
+      } else if (r_arg.type().is_int(8)) {
+        call_args.push_back(b_->getInt8(r_arg.as_int8()));
+      } else if (r_arg.type().is_int(16)) {
+        call_args.push_back(b_->getInt16(r_arg.as_int16()));
+      } else if (r_arg.type().is_int(32)) {
+        call_args.push_back(b_->getInt32(r_arg.as_int32()));
+      } else if (r_arg.type().is_int(64)) {
+        call_args.push_back(b_->getInt64(r_arg.as_int64()));
+      } else if (r_arg.type().is_uint(8)) {
+        call_args.push_back(b_->getInt8(r_arg.as_uint8()));
+      } else if (r_arg.type().is_uint(16)) {
+        call_args.push_back(b_->getInt16(r_arg.as_uint16()));
+      } else if (r_arg.type().is_uint(32)) {
+        call_args.push_back(b_->getInt32(r_arg.as_uint32()));
+      } else if (r_arg.type().is_uint(64)) {
+        call_args.push_back(b_->getInt64(r_arg.as_uint64()));
+      } else if (r_arg.type().is_float(32)) {
+        call_args.push_back(llvm::ConstantFP::get(
+            b_->getFloatTy(), llvm::APFloat(r_arg.as_float())));
+      } else if (r_arg.type().is_float(64)) {
+        call_args.push_back(llvm::ConstantFP::get(
+            b_->getDoubleTy(), llvm::APFloat(r_arg.as_double())));
+      } else if (r_arg.type().is_bfloat16()) {
+        call_args.push_back(llvm::ConstantFP::get(
+            b_->getBFloatTy(),
+            llvm::APFloat(static_cast<float>(r_arg.as_bfloat16()))));
+      } else if (r_arg.type().is_float16()) {
+        call_args.push_back(llvm::ConstantFP::get(
+            b_->getHalfTy(),
+            llvm::APFloat(static_cast<float>(r_arg.as_float16()))));
+      } else {
+        CINN_NOT_IMPLEMENTED;
+      }
+    }
+  }
+  b_->CreateCall(call_func, call_args);
+
+  return nullptr;
+}
+
 }  // namespace backends
 }  // namespace cinn
@@ -23,6 +23,8 @@
 
 #include "paddle/cinn/backends/llvm/codegen_llvm.h"
 
+PD_DECLARE_bool(cinn_bucket_compile);
+
 namespace cinn {
 namespace backends {
 
@@ -38,9 +40,16 @@ class CodeGenCUDA_Host : public CodeGenLLVM {
 
   using CodeGenLLVM::Visit;
   llvm::Value *Visit(const ir::_LoweredFunc_ *func) override {
+    if (FLAGS_cinn_bucket_compile) {
+      return LowerHostFunc(func);
+    }
     return LowerGPUKernelLauncher(func);
   }
 
+  llvm::Value *Visit(const ir::Call *op) override {
+    return LowerCUDAKernelCall(op);
+  }
+
  private:
   /**
    * Lower a CUDA kernel launcher.
@@ -56,6 +65,10 @@ class CodeGenCUDA_Host : public CodeGenLLVM {
    *
    */
   llvm::Value *LowerGPUKernelLauncher(const ir::_LoweredFunc_ *func);
+
+  llvm::Value *LowerHostFunc(const ir::_LoweredFunc_ *func);
+
+  llvm::Value *LowerCUDAKernelCall(const ir::Call *op);
 };
 
 }  // namespace backends
 
@@ -17,14 +17,104 @@
 #include "paddle/cinn/backends/cuda_util.h"
 #include "paddle/cinn/ir/ir_mutator.h"
 
+PD_DECLARE_bool(cinn_bucket_compile);
 namespace cinn {
 namespace backends {
 
 std::tuple<ir::Module, ir::Module> SplitCudaAndHostModule(ir::Module module) {
+  if (FLAGS_cinn_bucket_compile) {
+    detail::CollectBucketStrategyHostFunctionVisitor visitor(module->name);
+    Expr expr(module);
+    return visitor(&expr);
+  }
   detail::CollectHostFunctionVisitor visitor(module->name);
   Expr expr(module);
   return visitor(&expr);
 }
 
+struct PredicatePrinter : public ir::IrPrinter {
+  explicit PredicatePrinter(std::ostream &os) : ir::IrPrinter(os) {}
+
+ private:
+  void Visit(const ir::Add *x) { PrintBinaryOp("ADD", x); }
+  void Visit(const ir::Sub *x) { PrintBinaryOp("SUB", x); }
+  void Visit(const ir::Mul *x) { PrintBinaryOp("MUL", x); }
+  void Visit(const ir::Div *x) { PrintBinaryOp("DIV", x); }
+  void Visit(const ir::Mod *x) { PrintBinaryOp("MOD", x); }
+  void Visit(const ir::EQ *x) { PrintBinaryOp("EQ", x); }
+  void Visit(const ir::NE *x) { PrintBinaryOp("NE", x); }
+  void Visit(const ir::LT *x) { PrintBinaryOp("LT", x); }
+  void Visit(const ir::LE *x) { PrintBinaryOp("LE", x); }
+  void Visit(const ir::GT *x) { PrintBinaryOp("GT", x); }
+  void Visit(const ir::GE *x) { PrintBinaryOp("GE", x); }
+  void Visit(const ir::And *x) { PrintBinaryOp("AND", x); }
+  void Visit(const ir::Or *x) { PrintBinaryOp("OR", x); }
+
+  template <typename IRN>
+  void PrintBinaryOp(const std::string &op, const ir::BinaryOpNode<IRN> *x) {
+    str_ += "_FPA_";
+    ir::IrPrinter::Visit(x->a());
+    str_ += op;
+    ir::IrPrinter::Visit(x->b());
+    str_ += "_BPA_";
+  }
+};
+
+std::string Predicate2String(ir::Expr predicate) {
+  std::stringstream ss;
+  PredicatePrinter cond_printer(ss);
+  cond_printer.Print(predicate);
+  return ss.str();
+}
+
+std::string
+detail::CollectBucketStrategyHostFunctionVisitor::GenDeviceKernelName(
+    const std::string &fn_name, ir::Expr predicate) {
+  std::string cond_str = Predicate2String(predicate);
+  VLOG(3) << "predicate string: " << cond_str;
+  return fn_name + "__COND_" + cond_str + "__kernel";
+}
+
+void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
+    ir::Expr func, ir::Expr predicate) {
+  ir::_LoweredFunc_ *func_node = func.as_lowered_func();
+  CHECK(func_node);
+  if (!func_node->cuda_axis_info.valid()) {
+    func_node->cuda_axis_info.set_valid(true);
+  }
+  // process device func
+  device_module_builder.AddFunctionWithoutOptim(
+      CreateDeviceFunction(func, predicate).as_lowered_func_ref());
+  // process host func
+  ir::Var kernel_ptr(GenDeviceKernelName(func_node->name, predicate),
+                     type_of<std::string>());
+  ir::Expr call_extern_api =
+      ir::Call::Make(Void(),
+                     runtime::intrinsic::call_cuda_kernel,
+                     {kernel_ptr,
+                      kernel_args_,
+                      kernel_args_num_,
+                      Expr(func_node->cuda_axis_info.grid_dim(0)),   // grid_x
+                      Expr(func_node->cuda_axis_info.grid_dim(1)),   // grid_y
+                      Expr(func_node->cuda_axis_info.grid_dim(2)),   // grid_z
+                      Expr(func_node->cuda_axis_info.block_dim(0)),  // block_x
+                      Expr(func_node->cuda_axis_info.block_dim(1)),  // block_y
+                      Expr(func_node->cuda_axis_info.block_dim(2)),  // block_z
+                      kernel_stream_},
+                     {},
+                     ir::CallType::Extern,
+                     ir::FunctionRef(),
+                     0);
+  buckets_.emplace_back(ir::IfThenElse::Make(predicate, call_extern_api));
+}
+
+Expr detail::CollectBucketStrategyHostFunctionVisitor::CreateDeviceFunction(
+    ir::Expr expr, ir::Expr predicate) {
+  auto copied = ir::ir_utils::IRCopy(expr);
+  auto *lowered_func = copied.as_lowered_func();
+  lowered_func->name = GenDeviceKernelName(lowered_func->name, predicate);
+  return copied;
+}
+
 }  // namespace backends
 }  // namespace cinn
@@ -57,7 +57,7 @@ struct CollectHostFunctionVisitor : public ir::IRMutator<> {
                            device_module_builder.Build());
   }
 
- private:
+ protected:
   void Visit(const ir::_LoweredFunc_* op, Expr* expr) override {
     if (op->body.As<ir::Call>()) {
       host_module_builder.AddFunctionWithoutOptim(expr->as_lowered_func_ref());
@@ -137,11 +137,61 @@ struct CollectHostFunctionVisitor : public ir::IRMutator<> {
     return fn + "_kernel";
   }
 
- private:
+ protected:
   ir::Module::Builder host_module_builder;
   ir::Module::Builder device_module_builder;
 };
 
+struct CollectBucketStrategyHostFunctionVisitor
+    : public CollectHostFunctionVisitor {
+  explicit CollectBucketStrategyHostFunctionVisitor(
+      const std::string& module_name)
+      : CollectHostFunctionVisitor(module_name),
+        kernel_args_(KERNEL_ARGS, type_of<void*>()),
+        kernel_args_num_(KERNEL_ARGS_NUM, type_of<int>()),
+        kernel_stream_(KERNEL_STREAM, type_of<void*>()) {}
+
+  std::tuple<ir::Module, ir::Module> operator()(Expr* expr) {
+    ir::IRMutator<>::Visit(expr, expr);
+    return std::make_tuple(host_module_builder.Build(),
+                           device_module_builder.Build());
+  }
+
+ private:
+  void Visit(const ir::_Module_* op, Expr* expr) {
+    CHECK_EQ(op->functions.size(), op->predicates.size());
+    for (int i = 0; i < op->functions.size(); ++i) {
+      ProcessLoweredFunc(op->functions[i], op->predicates[i]);
+    }
+
+    std::vector<ir::Argument> arguments = {
+        ir::Argument(kernel_args_, ir::Argument::IO::kOutput),
+        ir::Argument(kernel_args_num_, ir::Argument::IO::kInput),
+        ir::Argument(kernel_stream_, ir::Argument::IO::kOutput)};
+    ir::Expr host_func =
+        ir::_LoweredFunc_::Make(op->functions[0].as_lowered_func()->name,
+                                arguments,
+                                ir::Block::Make(buckets_),
+                                {});
+    host_module_builder.AddFunctionWithoutOptim(
+        host_func.as_lowered_func_ref());
+  }
+
+  void ProcessLoweredFunc(ir::Expr func, ir::Expr predicate);
+
+  Expr CreateDeviceFunction(ir::Expr expr, ir::Expr predicate);
+
+  inline std::string GenDeviceKernelName(const std::string& fn_name,
+                                         ir::Expr predicate);
+
+ private:
+  std::vector<ir::Expr> buckets_;
+
+  ir::Var kernel_args_;
+  ir::Var kernel_args_num_;
+  ir::Var kernel_stream_;
+};
+
 }  // namespace detail
 
 }  // namespace backends
 
@@ -20,6 +20,7 @@
 #include "paddle/cinn/hlir/framework/graph.h"
 #include "paddle/cinn/hlir/framework/op_lowering_impl.h"
 #include "paddle/cinn/hlir/framework/op_lowering_impl_base.h"
+#include "paddle/cinn/ir/group_schedule/base_group_scheduler.h"
 #include "paddle/cinn/lang/packed_func.h"
 #ifndef CINN_WITH_ONLY
 #include "paddle/cinn/hlir/framework/pir/op_lowering_impl.h"
@@ -46,6 +47,15 @@ class OpLowerer {
         group, apply_op_schedule, apply_group_schedule, apply_pass);
   }
 
+  std::vector<std::pair<ir::SymbolicPredicate, ir::LoweredFunc>> BucketLower(
+      const T& group,
+      bool apply_op_schedule = false,
+      bool apply_group_schedule = true,
+      bool apply_pass = true) {
+    return impl_->BucketLower(
+        group, apply_op_schedule, apply_group_schedule, apply_pass);
+  }
+
  private:
   std::shared_ptr<OpLowererImplBase<T>> impl_;
 };