[CINN] Support grid reduce for single SR reduce (#68220)

lshpku · web-flow · commit 4c1ba9a2e500 · 2024-10-15T11:23:39.000+08:00
diff --git a/paddle/cinn/backends/codegen_device_util.cc b/paddle/cinn/backends/codegen_device_util.cc
@@ -262,6 +262,24 @@ void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
                      ir::CallType::Extern,
                      ir::FunctionRef(),
                      0);
+
+  // create memset calls for temp_spaces if needed
+  std::vector<ir::Expr> call_kernel_stmts;
+  for (auto &temp_space : func_node->temp_spaces) {
+    if (temp_space.need_zero_init()) {
+      ir::Expr size = common::cast(temp_space.size(), common::UInt(64));
+      ir::Expr call_get_arg =
+          lang::CallExtern(runtime::intrinsic::get_item_in_cuda_kernel_args,
+                           {kernel_args_, ir::Expr(temp_space.arg_idx())});
+      ir::Expr call_memset = lang::CallExtern(
+          runtime::intrinsic::call_cuda_memset,
+          {call_get_arg, ir::Expr(1), ir::Expr(0), size, kernel_stream_});
+      call_kernel_stmts.push_back(call_memset);
+    }
+  }
+  call_kernel_stmts.push_back(call_extern_api);
+  call_extern_api = ir::Block::Make(call_kernel_stmts);
+
   if (buckets_.empty()) {
     buckets_.emplace_back(ir::IfThenElse::Make(predicate, call_extern_api));
   } else {
@@ -270,6 +288,26 @@ void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
     buckets_.emplace_back(
         ir::IfThenElse::Make(predicate, call_extern_api, false_expr));
   }
+
+  // create infer shape calls for temp_spaces
+  std::vector<ir::Expr> temp_space_infer_shape_stmts;
+  for (int i = 0; i < func_node->temp_spaces.size(); ++i) {
+    ir::Var tensor_shape_args(TENSOR_SHAPE_ARGS, type_of<int64_t **>());
+    ir::Expr size =
+        common::cast(func_node->temp_spaces[i].size(), common::Int(64));
+    ir::Expr call_set_value =
+        lang::CallExtern(runtime::intrinsic::infer_shape_set_value,
+                         {ir::Expr(func_node->num_output_tensors + i),
+                          ir::Expr(0),
+                          size,
+                          tensor_shape_args});
+    temp_space_infer_shape_stmts.push_back(call_set_value);
+  }
+  if (!temp_space_infer_shape_stmts.empty()) {
+    ir::Expr if_body = ir::Block::Make(temp_space_infer_shape_stmts);
+    temp_space_infer_shape_body_ =
+        ir::IfThenElse::Make(predicate, if_body, temp_space_infer_shape_body_);
+  }
 }
 
 void detail::CollectBucketStrategyHostFunctionVisitor::ProcessArgs(
diff --git a/paddle/cinn/backends/codegen_device_util.h b/paddle/cinn/backends/codegen_device_util.h
@@ -280,6 +280,9 @@ struct CollectBucketStrategyHostFunctionVisitor
     infer_shape_func_body_stmts.insert(
         infer_shape_func_body_stmts.end(),
         op->infer_shape_func.as_lowered_func()->body);
+    if (temp_space_infer_shape_body_.defined()) {
+      infer_shape_func_body_stmts.push_back(temp_space_infer_shape_body_);
+    }
 
     std::vector<ir::Argument> infer_shape_arguments = {
         ir::Argument(kernel_args_, ir::Argument::IO::kOutput),
@@ -307,6 +310,7 @@ struct CollectBucketStrategyHostFunctionVisitor
  private:
   std::vector<ir::Expr> buckets_;
   std::vector<ir::Expr> arg_defs_;
+  ir::Expr temp_space_infer_shape_body_;
 
   ir::Var kernel_args_;
   ir::Var kernel_args_num_;
diff --git a/paddle/cinn/hlir/framework/pir/compilation_cache.cc b/paddle/cinn/hlir/framework/pir/compilation_cache.cc
@@ -54,6 +54,7 @@ pir::CINNKernelInfo BackendResource::GenerateKernelInfo() const {
   kernel_info.infer_shape_fn_ptr = GetInferFuncPtr();
   kernel_info.CX86_fn_ptr = GetCX86HostFuncPtr();
   kernel_info.symbol_args_map = GetSymbolArgsMap();
+  kernel_info.temp_space_sizes = GetTempSpaceSizes();
   return kernel_info;
 }
 }  // namespace pir
diff --git a/paddle/cinn/hlir/framework/pir/compilation_cache.h b/paddle/cinn/hlir/framework/pir/compilation_cache.h
@@ -33,10 +33,12 @@ class BackendResource final {
       const Target& target,
       const std::string& host_fn_name,
       const std::string& infer_fn_name,
-      const std::map<int, CINNKernelInfo::SymbolArgBindInfo>& symbol_args_map)
+      const std::map<int, CINNKernelInfo::SymbolArgBindInfo>& symbol_args_map,
+      const std::vector<int64_t>& temp_space_sizes)
       : host_fn_name_(host_fn_name),
         infer_fn_name_(infer_fn_name),
-        symbol_args_map_(symbol_args_map) {
+        symbol_args_map_(symbol_args_map),
+        temp_space_sizes_(temp_space_sizes) {
     backend_compiler_ = backends::Compiler::Create(target);
   }
 
@@ -47,6 +49,9 @@ class BackendResource final {
       const {
     return symbol_args_map_;
   }
+  const std::vector<int64_t>& GetTempSpaceSizes() const {
+    return temp_space_sizes_;
+  }
   const std::shared_ptr<backends::Compiler>& GetBackendCompiler() const {
     return backend_compiler_;
   }
@@ -57,6 +62,7 @@ class BackendResource final {
   std::string host_fn_name_;
   std::string infer_fn_name_;
   std::map<int, CINNKernelInfo::SymbolArgBindInfo> symbol_args_map_;
+  std::vector<int64_t> temp_space_sizes_;
 
   std::shared_ptr<backends::Compiler> backend_compiler_{nullptr};
 };
diff --git a/paddle/cinn/hlir/framework/pir/compilation_task.cc b/paddle/cinn/hlir/framework/pir/compilation_task.cc
@@ -216,7 +216,8 @@ std::shared_ptr<pir::CompilationResult> CompilationTask::BuildPirCINNKernelInfo(
       context_->target_,
       context_->group_->FuncName(),
       context_->group_->FuncName() + "_infer_shape",
-      context_->group_->symbol_args_map());
+      context_->group_->symbol_args_map(),
+      context_->group_->temp_space_sizes());
   VLOG(5) << "Start to compile module into cuda kernel...";
   backend_resource->GetBackendCompiler()->Build(module, "");
   backend_resource->GetBackendCompiler()->AppendCX86(CX86module);
@@ -236,7 +237,8 @@ CompilationTask::CompileBroadcastModules(
       context_->target_,
       context_->group_->FuncName(),
       context_->group_->FuncName() + "_infer_shape",
-      context_->group_->symbol_args_map());
+      context_->group_->symbol_args_map(),
+      context_->group_->temp_space_sizes());
 
   std::vector<std::string> case_func_names;
   std::vector<ir::Expr> broadcast_conditions;
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_group.h b/paddle/cinn/hlir/framework/pir/op_lowering_group.h
@@ -185,6 +185,14 @@ class OpLoweringGroup {
     return this->symbol_args_map_;
   }
 
+  const std::vector<int64_t>& temp_space_sizes() const {
+    return this->temp_space_sizes_;
+  }
+
+  std::vector<int64_t>& mut_temp_space_sizes() {
+    return this->temp_space_sizes_;
+  }
+
  private:
   using alignment_schedule_info_t = std::unordered_map<
       ::pir::Operation*,
@@ -231,6 +239,7 @@ class OpLoweringGroup {
   std::vector<std::string> output_names_;
   std::vector<::pir::Value> output_values_;
   std::map<int, CINNKernelInfo::SymbolArgBindInfo> symbol_args_map_;
+  std::vector<int64_t> temp_space_sizes_;
 
   alignment_schedule_info_t alignment_schedule_info_;
   std::vector<int64_t> reduce_axis_;
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -117,6 +117,12 @@ BucketLoweredFuncsWrapper OpLowererImpl::BucketLower(
   func_bodies = OperationFusion(ops, func_bodies, group->fusion_tracker_ptr);
   std::shared_ptr<FusionGroupInfo> fusion_group_info =
       GetFusionGroupInfo(func_bodies);
+  // TODO(liangshuhao): grid reduce is disabled for broadcast-leaf group now,
+  // because grid reduce introduces extra func args that currently cannot be
+  // unified with other broadcast-leaf groups.
+  if (group->IsBroadcastLeaf()) {
+    fusion_group_info->can_apply_grid_reduce = false;
+  }
 
   if (FLAGS_cinn_check_tensor_buffer_map) {
     optim::CheckTensorBufferMap(func_bodies, "BucketLower OpFusion");
@@ -386,9 +392,21 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
       func = optim::Optimize(Expr(func), common::DefaultHostTarget(), false)
                  .as_lowered_func_ref();
     }
+    func->num_output_tensors = infer_shape_arg_tensor->size();
     lowered_funcs.push_back(std::move(func));
   }
 
+  // collect temp space sizes
+  if (lowered_funcs.size() > 1) {
+    for (auto& temp_space : lowered_funcs[0]->temp_spaces) {
+      int64_t size = -1;
+      if (temp_space.size().is_constant()) {
+        size = temp_space.size().as_int64();
+      }
+      group->mut_temp_space_sizes().push_back(size);
+    }
+  }
+
   return lowered_funcs;
 }
 
diff --git a/paddle/cinn/hlir/framework/pir/utils.h b/paddle/cinn/hlir/framework/pir/utils.h
@@ -66,6 +66,11 @@ struct CINNKernelInfo {
   //     3: ArgValueIdx{1, 6}
   //   }
   std::map<int, SymbolArgBindInfo> symbol_args_map;
+
+  // Sizes in bytes of the temporary global spaces needed by the kernel.
+  // These spaces are allocated before the kernel is launched, appended to the
+  // kernel's argument list, and released when the kernel completes.
+  std::vector<int64_t> temp_space_sizes;
 };
 
 struct CompatibleInfo {
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.cc b/paddle/cinn/ir/group_schedule/config/group_tile_config.cc
@@ -95,6 +95,7 @@ std::shared_ptr<ScheduleConfig::BaseInfo> InitBasicInfo(
       std::make_shared<ScheduleConfig::BaseInfo>();
   base_info->data_rank = group_info->loop_ranges.size();
   base_info->loop_strides = group_info->loop_strides;
+  base_info->can_apply_grid_reduce = group_info->can_apply_grid_reduce;
 
   std::set<int64_t> reduce_dim_loc;
   for (int64_t dim : group_info->reduce_axis) {
@@ -179,6 +180,19 @@ BuildPureStaticShapeConfig(
           /* spatial_inner_num = */ 1,
           /* reduce_method = */ BlockReduceMethod()};
       return {{bucket_info, tile_config}};
+    } else if (base_info->can_apply_grid_reduce &&
+               base_info->reduce_numel > 65536) {
+      BucketInfo bucket_info{/* sp_lower_bound = */ 1,
+                             /* sp_upper_bound = */ 1,
+                             /* rb_lower_bound = */ 2049,
+                             /* rb_upper_bound = */ kMaxNumel};
+      ScheduleConfig::TileConfig tile_config{
+          /* warp_num = */ 32,
+          /* tree_reduce_num = */ 1024,
+          /* spatial_inner_num = */ 1,
+          /* reduce_method = */ BlockReduceMethod(),
+          /* grid_reduce_num = */ 8};
+      return {{bucket_info, tile_config}};
     } else {
       BucketInfo bucket_info{/* sp_lower_bound = */ 1,
                              /* sp_upper_bound = */ 1,
@@ -244,6 +258,20 @@ BuildPureStaticShapeConfig(
         /* spatial_inner_num = */ spatial_inner_num,
         /* reduce_method = */ BlockReduceMethod()};
     return {{bucket_info, tile_config}};
+  } else if (base_info->can_apply_grid_reduce &&
+             base_info->spatial_numel <= 256 &&
+             base_info->reduce_numel > 65536) {
+    BucketInfo bucket_info{/* sp_lower_bound = */ 1,
+                           /* sp_upper_bound = */ 256,
+                           /* rb_lower_bound = */ 65537,
+                           /* rb_upper_bound = */ kMaxNumel};
+    ScheduleConfig::TileConfig tile_config{
+        /* warp_num = */ 32,
+        /* tree_reduce_num = */ 1024,
+        /* spatial_inner_num = */ 1,
+        /* reduce_method = */ BlockReduceMethod(),
+        /* grid_reduce_num = */ 8};
+    return {{bucket_info, tile_config}};
   } else {
     int64_t warp_num = 32;
     int64_t spatial_inner_num = 1;
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.h b/paddle/cinn/ir/group_schedule/config/group_tile_config.h
@@ -40,6 +40,7 @@ struct ScheduleConfig {
     bool has_dynamic_spatial{false};
     bool has_dynamic_reduce{false};
     bool is_reduce_all{false};
+    bool can_apply_grid_reduce{false};
     IterSpaceType iter_space_type;
   };
 
@@ -48,6 +49,7 @@ struct ScheduleConfig {
     int64_t tree_reduce_num{1};
     int64_t spatial_inner_num{1};
     ReduceMethod reduce_method{NoneReduceMethod()};
+    int64_t grid_reduce_num{1};
   };
 
   std::shared_ptr<BaseInfo> base_info;
diff --git a/paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc b/paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc
@@ -164,9 +164,11 @@ void TileFirstGeneralTactic::ApplyContinuousDataTile(
                          context_->config.tile_config.tree_reduce_num;
   const auto sp_loop = context_->config.tile_config.spatial_inner_num;
   const auto rd_thread = context_->config.tile_config.tree_reduce_num;
+  const auto rd_block = context_->config.tile_config.grid_reduce_num;
   VLOG(4) << "ApplyContinuousDataTile sp_thread=" << sp_thread;
   VLOG(4) << "ApplyContinuousDataTile sp_loop=" << sp_loop;
   VLOG(4) << "ApplyContinuousDataTile rd_thread=" << rd_thread;
+  VLOG(4) << "ApplyContinuousDataTile rd_block=" << rd_block;
   VLOG(4) << "ApplyContinuousDataTile vec_flatten_axis: "
           << utils::Join(vec_flatten_axis_, ", ");
   VLOG(4) << "ApplyContinuousDataTile vec_reduce_axis: "
@@ -204,21 +206,14 @@ void TileFirstGeneralTactic::ApplyContinuousDataTile(
   VLOG(4) << "After SplitSptial on block: [" << block_id << "], loop nest:\n"
           << sch->GetModule().GetExprs().front();
 
-  // Split reduce axes -> [rd_loop, rd_thread]
+  // Split reduce axes -> [rd_loop, rd_block, rd_thread]
+  std::string global_rf_block;
   if (vec_reduce_axis_.size() > 0) {
     auto loops = sch->GetLoops(block_id);
-    // [S..S, R] => [S..S, R(-1), R(thread)]
-    sch->Split(loops[current_reduce_axis], {-1, rd_thread});
-    VLOG(4) << "Before ReorderReduction on block: [" << block_id
-            << "], loop nest:\n"
-            << sch->GetModule().GetExprs().front();
+    sch->Split(loops[current_reduce_axis], {-1, rd_block * rd_thread});
 
     loops = sch->GetLoops(block_id);
-    // [S..S, R(-1), R(thread)] => [S..S, R(thread), R(-1)]
     sch->Reorder({loops[current_reduce_axis + 1], loops[current_reduce_axis]});
-    VLOG(4) << "Before FactorizeReduction on block: [" << block_id
-            << "], loop nest:\n"
-            << sch->GetModule().GetExprs().front();
 
     if (IsReductionSBlock(sch->GetBlock(block_id))) {
       loops = sch->GetLoops(block_id);
@@ -228,6 +223,24 @@ void TileFirstGeneralTactic::ApplyContinuousDataTile(
                                   /* with_write_back_block_init = */ false);
       map_rf_block_[block_id] = rf_tensor.as_tensor_ref()->name;
     }
+
+    if (rd_block > 1) {
+      loops = sch->GetLoops(block_id);
+      sch->Split(loops[current_reduce_axis], {rd_block, rd_thread});
+
+      if (IsReductionSBlock(sch->GetBlock(block_id))) {
+        loops = sch->GetLoops(map_rf_block_[block_id]);
+        sch->Split(loops[current_reduce_axis], {rd_block, rd_thread});
+
+        loops = sch->GetLoops(block_id);
+        ir::Expr rf_tensor =
+            sch->FactorizeReduction(loops[current_reduce_axis],
+                                    /* rf_axis = */ 0,
+                                    /* with_write_back_block_init = */ false);
+        global_rf_block = rf_tensor.as_tensor_ref()->name;
+        rf_tensor.as_tensor_ref()->WithBuffer("global", "_" + global_rf_block);
+      }
+    }
   }
   VLOG(4) << "After SplitReduce on block: [" << block_id << "], loop nest:\n"
           << sch->GetModule().GetExprs().front();
@@ -248,16 +261,23 @@ void TileFirstGeneralTactic::ApplyContinuousDataTile(
       }
     }
     if (!vec_reduce_axis_.empty() && current_reduce_axis > 0) {
-      // [S(blockIdx.x), optional(inner_loop), S(threadIdx.y), R..R] =>
-      // [S(blockIdx.x), optional(inner_loop), S(threadIdx.y), R(threadIdx.x),
-      // R(inner_loop)]
-      sch->Bind(loops[current_reduce_axis], rd_axis_type);
+      if (rd_block > 1) {
+        sch->Bind(loops[current_reduce_axis], "blockIdx.y");
+        if (loops.size() > current_reduce_axis + 1) {
+          sch->Bind(loops[current_reduce_axis + 1], rd_axis_type);
+        }
+      } else {
+        sch->Bind(loops[current_reduce_axis], rd_axis_type);
+      }
     }
   };
   DoBind(sch->GetLoops(block_id));
   if (map_rf_block_.count(block_id) > 0) {
     DoBind(sch->GetLoops(map_rf_block_[block_id]));
   }
+  if (!global_rf_block.empty()) {
+    DoBind(sch->GetLoops(global_rf_block));
+  }
   VLOG(4) << "After BindCudaInfo on block: [" << block_id << "], loop nest:\n"
           << sch->GetModule().GetExprs().front();
 
diff --git a/paddle/cinn/ir/lowered_func.h b/paddle/cinn/ir/lowered_func.h
@@ -157,6 +157,10 @@ struct _LoweredFunc_ : ExprNode<_LoweredFunc_> {
   //! argument list.
   std::vector<TempSpaceInfo> temp_spaces;
 
+  //! Number of output tensors that appear in the function's argument list.
+  //! This number doesn't include temp_spaces.
+  int num_output_tensors;
+
   //! Body of this function.
   Expr body;
 
diff --git a/paddle/cinn/ir/utils/ir_copy.cc b/paddle/cinn/ir/utils/ir_copy.cc
@@ -297,6 +297,8 @@ struct IRCopyVisitor : public ir::IRVisitorRequireReImpl<Expr> {
     func->args = op->args;
     func->body = Visit(&op->body);
     func->temp_bufs = op->temp_bufs;
+    func->temp_spaces = op->temp_spaces;
+    func->num_output_tensors = op->num_output_tensors;
 
     func->device_api = op->device_api;
 
diff --git a/paddle/cinn/optim/replace_cross_block_reduction.cc b/paddle/cinn/optim/replace_cross_block_reduction.cc
diff --git a/paddle/fluid/framework/new_executor/instruction/cinn_jit_instruction.cc b/paddle/fluid/framework/new_executor/instruction/cinn_jit_instruction.cc
diff --git a/paddle/fluid/framework/new_executor/instruction/cinn_jit_instruction.h b/paddle/fluid/framework/new_executor/instruction/cinn_jit_instruction.h

Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,7 @@ pir::CINNKernelInfo BackendResource::GenerateKernelInfo() const {`
`54`	`54`	`kernel_info.infer_shape_fn_ptr = GetInferFuncPtr();`
`55`	`55`	`kernel_info.CX86_fn_ptr = GetCX86HostFuncPtr();`
`56`	`56`	`kernel_info.symbol_args_map = GetSymbolArgsMap();`
	`57`	`+ kernel_info.temp_space_sizes = GetTempSpaceSizes();`
`57`	`58`	`return kernel_info;`
`58`	`59`	`}`
`59`	`60`	`} // namespace pir`