[CINN] Add the TileBroadcastTactic for NCHW broadcast (#70092)

Author: lshpku

paddle/cinn/ir/group_schedule/config/group_tile_config.cc

@@ -150,22 +150,17 @@ std::shared_ptr<ScheduleConfig::BaseInfo> InitBasicInfo(
     const std::shared_ptr<FusionGroupInfo>& group_info) {
   std::shared_ptr<ScheduleConfig::BaseInfo> base_info =
       std::make_shared<ScheduleConfig::BaseInfo>();
-  base_info->data_rank = group_info->loop_ranges.size();
+  base_info->reduce_axis = group_info->reduce_axis;
+  base_info->loop_ranges = group_info->loop_ranges;
   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) {
-    if (dim < 0) {
-      dim += base_info->data_rank;
-    }
-    base_info->reduce_axis.push_back(dim);
-    reduce_dim_loc.insert(dim);
-  }
+  std::set<int64_t> reduce_dim_loc(group_info->reduce_axis.begin(),
+                                   group_info->reduce_axis.end());
 
   base_info->spatial_numel = 1;
   base_info->reduce_numel = 1;
-  for (int64_t i = 0; i < base_info->data_rank; ++i) {
+  for (int64_t i = 0; i < base_info->loop_ranges.size(); ++i) {
     if (reduce_dim_loc.count(i)) {
       if (group_info->loop_ranges[i] == -1)
         base_info->has_dynamic_reduce = true;

paddle/cinn/ir/group_schedule/config/group_tile_config.h

@@ -33,8 +33,8 @@ using IterSpaceType = std::vector<std::pair<std::string, std::string>>;
 struct ScheduleConfig {
   struct BaseInfo {
     std::vector<int64_t> reduce_axis;
+    std::vector<int64_t> loop_ranges;
     std::vector<int64_t> loop_strides;
-    int64_t data_rank;
     int64_t reduce_numel;
     int64_t spatial_numel;
     bool has_dynamic_spatial{false};

paddle/cinn/ir/group_schedule/dy_shape_group_scheduler.cc

@@ -17,6 +17,7 @@
 #include "paddle/cinn/ir/group_schedule/config/schedule_config_manager.h"
 #include "paddle/cinn/ir/group_schedule/tactic/compute_at_reduction_tactic.h"
 #include "paddle/cinn/ir/group_schedule/tactic/compute_inline_tactic.h"
+#include "paddle/cinn/ir/group_schedule/tactic/tile_broadcast_tactic.h"
 #include "paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.h"
 #include "paddle/cinn/ir/ir_analyzer/ir_analyzer.h"
 #include "paddle/cinn/ir/op/ir_operators.h"
@@ -31,12 +32,10 @@ void DynamicShapeGroupScheduler::Init() {
   VLOG(4) << "original group func body: \n"
           << ir_sch_->GetModule().GetExprs()[0];
   InitBuckets();
+  tactics_.emplace_back(CreateTileBroadcastTactic());
   tactics_.emplace_back(CreateTileFirstGeneralTactic());
-  VLOG(4) << "CreateTileFirstGeneralTactic End";
   tactics_.emplace_back(CreateComputeInlineTactic());
-  VLOG(4) << "CreateTileCreateComputeInlineTactic End";
   tactics_.emplace_back(CreateComputeAtReductionTactic());
-  VLOG(4) << "CreateComputeAtReductionTactic End";
 }
 
 void DynamicShapeGroupScheduler::InitBuckets() {
@@ -156,7 +155,8 @@ void DynamicShapeGroupScheduler::ApplyTactics(BucketContext* bucket_context) {
               << "] on ScheduleBlockNode [" << node->id() << "] func body:\n"
               << bucket_context->ir_sch->GetModule().GetExprs().front();
     };
-    tactic->Init(&(bucket_context->schedule_context));
+    tactic->Init(&(bucket_context->schedule_context),
+                 bucket_context->ir_sch.get());
     bucket_context->schedule_block_graph->DFSTopoWalk(ApplyTacticFunc);
     bucket_context->schedule_block_graph->Update(*(bucket_context->ir_sch));
     VLOG(5) << "[End " << tactic->TacticName() << "] func body: "

paddle/cinn/ir/group_schedule/tactic/CMakeLists.txt

@@ -7,4 +7,5 @@ gather_srcs(cinnapi_src SRCS optimize_reduction_tactic.cc)
 gather_srcs(cinnapi_src SRCS compute_at_reduction_tactic.cc)
 gather_srcs(cinnapi_src SRCS bind_cuda_tactic.cc)
 gather_srcs(cinnapi_src SRCS arrange_storage_tactic.cc)
+gather_srcs(cinnapi_src SRCS tile_broadcast_tactic.cc)
 gather_srcs(cinnapi_src SRCS tile_first_general_tactic.cc)

paddle/cinn/ir/group_schedule/tactic/schedule_tactic.h

@@ -91,7 +91,17 @@ struct ScheduleContext {
 
 class ScheduleTactic {
  public:
-  virtual void Init(ScheduleContext* context) = 0;
+  // Attribute key to record which tile tactic has been applied on a graph.
+  // Exactly one tile tactic is applied on a graph during scheduling.
+  static constexpr char* kTileMethod = "tile_method";
+
+  virtual void Init(ScheduleContext* context) {
+    PADDLE_THROW(::common::errors::Unimplemented(
+        "ScheduleTactic subclass must implement one of the Init method."));
+  }
+  virtual void Init(ScheduleContext* context, ir::IRSchedule* sch) {
+    Init(context);
+  }
 
   virtual void Apply(ir::IRSchedule* sch, const std::string& block_id) = 0;