align reduce loops to input (#66419)

lshpku · web-flow · commit 9f7638bce358 · 2024-07-26T14:44:09.000+08:00
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -94,6 +94,7 @@ std::shared_ptr<GroupInfo> OpLowererImpl::GetGroupInfo(
     const std::unordered_map<::pir::Value, ir::Tensor>& tensor_map) {
   std::shared_ptr<GroupInfo> group_info = std::make_shared<GroupInfo>();
   group_info->data_space = fusion_group_info.loop_ranges;
+  group_info->loop_transform_map = fusion_group_info.loop_transform_map;
   group_info->reduce_axis = fusion_group_info.reduce_axis;
   group_info->reduce_var_names =
       std::set<std::string>(fusion_group_info.reduce_var_name.begin(),
@@ -102,14 +103,6 @@ std::shared_ptr<GroupInfo> OpLowererImpl::GetGroupInfo(
   for (auto& val : group->output_values()) {
     group_info->direct_output_var_names.insert(ValueName(val));
   }
-
-  group->WalkOps([&group_info](::pir::Operation* op) {
-    if (CompatibleInfo::OpKind(*op) == OpPatternKind::kReduction) {
-      group_info->raw_reduce_axis = cinn::fusion::GetReduceAxisIdx(op);
-      group_info->raw_data_rank =
-          cinn::fusion::GetCompitableRank(op->operand_source(0));
-    }
-  });
   return group_info;
 }
 
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.h b/paddle/cinn/hlir/framework/pir/op_lowering_impl.h
@@ -50,8 +50,7 @@ typedef bool (OpLowererImpl::*ScheduleDetermineFunction)(::pir::Operation*);
 struct GroupInfo {
   std::vector<int64_t> data_space;
   std::vector<int64_t> reduce_axis;
-  int64_t raw_data_rank;
-  std::vector<int64_t> raw_reduce_axis;
+  std::vector<int64_t> loop_transform_map;
   std::set<std::string> reduce_var_names;
   std::set<std::string> shared_var_names;
   std::set<std::string> direct_output_var_names;
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc b/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
@@ -618,6 +618,55 @@ std::vector<ir::Expr> OperationFusion(
   return output;
 }
 
+void InitLoopTransformMap(FusionGroupInfo* group_info, const ir::Expr& body) {
+  using trivial_fusion_detail::ExprSetFinderUtils::ChildScheduleBlockRealizes;
+  using trivial_fusion_detail::ExprSetFinderUtils::ChildTensorLoads;
+  using trivial_fusion_detail::ExprSetFinderUtils::
+      ScheduleBlockRealizeIsSplitTransform;
+
+  std::vector<ir::Expr> split_transform_block =
+      (ChildScheduleBlockRealizes * ScheduleBlockRealizeIsSplitTransform)(body);
+  if (split_transform_block.empty()) {
+    return;
+  }
+  CHECK_EQ(split_transform_block.size(), 1);
+  VLOG(4) << "found split_transform block: " << body;
+
+  auto* block = split_transform_block[0].As<ir::ScheduleBlockRealize>();
+  auto& iter_values = block->iter_values;
+  auto& iter_vars = block->schedule_block.As<ir::ScheduleBlock>()->iter_vars;
+  CHECK_EQ(iter_values.size(), iter_vars.size());
+
+  // Map iter_vars (e.g. inner_block_0) to loops.
+  std::unordered_map<ir::Var, size_t> iter_var_map;
+  std::vector<ir::Var> for_iters = trivial_fusion_detail::GetAllForIters(body);
+  for (size_t i = 0; i < iter_vars.size(); i++) {
+    auto iter = std::find(
+        for_iters.begin(), for_iters.end(), iter_values[i].as_var_ref());
+    CHECK(iter != for_iters.end());
+    iter_var_map[iter_vars[i]] = std::distance(for_iters.begin(), iter);
+  }
+
+  // Map loops to the load indices of the first input.
+  // TODO(liangshuhao): Choose the most appropriate input when there are
+  // multiple inputs.
+  const auto& all_loads = ChildTensorLoads(split_transform_block[0]);
+  group_info->loop_transform_map.assign(for_iters.size(), -1);
+  for (const auto& load : all_loads) {
+    const auto& indices = load.As<ir::Load>()->indices;
+    VLOG(4) << "load indices: " << utils::Join(indices, ", ");
+    for (size_t i = 0; i < indices.size(); i++) {
+      if (indices[i].as_var()) {
+        size_t loop_idx = iter_var_map[indices[i].as_var_ref()];
+        group_info->loop_transform_map[loop_idx] = i;
+      }
+    }
+    VLOG(4) << "loop_transform_map: "
+            << utils::Join(group_info->loop_transform_map, ", ");
+    break;
+  }
+}
+
 FusionGroupInfo GetFusionGroupInfo(
     const std::vector<ir::Expr>& op_compute_bodies) {
   using trivial_fusion_detail::AppendBound;
@@ -635,6 +684,8 @@ FusionGroupInfo GetFusionGroupInfo(
   };
 
   for (const auto& body : op_compute_bodies) {
+    InitLoopTransformMap(&group_info, body);
+
     if (IsReduceBody(body)) {
       ReduceOp op = ReduceOp(body);
       if (group_info.reduce_var_name.empty()) {
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_impl.h b/paddle/cinn/hlir/framework/pir/trivial_op_impl.h
@@ -165,6 +165,7 @@ std::vector<ir::Var> GetAllForIters(const ir::Expr& expr);
 
 struct FusionGroupInfo {
   std::vector<int64_t> loop_ranges;
+  std::vector<int64_t> loop_transform_map;
   std::vector<int64_t> reduce_axis;
   std::vector<std::string> reduce_var_name;
 
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_util.cc b/paddle/cinn/hlir/framework/pir/trivial_op_util.cc
@@ -282,6 +282,15 @@ ExprSetFinder ScheduleBlockRealizeIsInit = FilterMaker(
     },
     "ScheduleBlockRealizeIsInit");
 
+ExprSetFinder ScheduleBlockRealizeIsSplitTransform = FilterMaker(
+    [](const ir::Expr& e) -> bool {
+      return (e.As<ir::ScheduleBlockRealize>() &&
+              e.As<ir::ScheduleBlockRealize>()
+                      ->schedule_block.As<ir::ScheduleBlock>()
+                      ->name.find("_split_transform") != std::string::npos);
+    },
+    "ScheduleBlockRealizeIsSplitTransform");
+
 ExprSetFinder IsFor = FilterMaker(
     [](const ir::Expr& e) -> bool { return e.As<ir::For>(); }, "IsFor");
 
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_util.h b/paddle/cinn/hlir/framework/pir/trivial_op_util.h
@@ -163,6 +163,8 @@ extern ExprSetFinder ScheduleBlockRealizeIsNotInit;
 
 extern ExprSetFinder ScheduleBlockRealizeIsInit;
 
+extern ExprSetFinder ScheduleBlockRealizeIsSplitTransform;
+
 extern ExprSetFinder IsFor;
 
 extern ExprSetFinder ChildScheduleBlocks;
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.cc b/paddle/cinn/ir/group_schedule/config/group_tile_config.cc
@@ -97,7 +97,7 @@ std::shared_ptr<ScheduleConfig::BaseInfo> InitBasicInfo(
   base_info->shared_var_names = group_info->shared_var_names;
   base_info->direct_output_var_names = group_info->direct_output_var_names;
   base_info->data_rank = group_info->data_space.size();
-  base_info->raw_data_rank = group_info->raw_data_rank;
+  base_info->loop_transform_map = group_info->loop_transform_map;
 
   std::set<int64_t> reduce_dim_loc;
   for (int64_t dim : group_info->reduce_axis) {
@@ -108,13 +108,6 @@ std::shared_ptr<ScheduleConfig::BaseInfo> InitBasicInfo(
     reduce_dim_loc.insert(dim);
   }
 
-  for (int64_t dim : group_info->raw_reduce_axis) {
-    if (dim < 0) {
-      dim += base_info->data_rank;
-    }
-    base_info->raw_reduce_axis.push_back(dim);
-  }
-
   base_info->spatial_numel = 1;
   base_info->reduce_numel = 1;
   for (int64_t i = 0; i < base_info->data_rank; ++i) {
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.h b/paddle/cinn/ir/group_schedule/config/group_tile_config.h
@@ -31,9 +31,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> raw_reduce_axis;
+    std::vector<int64_t> loop_transform_map;
     int64_t data_rank;
-    int64_t raw_data_rank;
     int64_t reduce_numel;
     int64_t spatial_numel;
     bool has_dynamic_spatial{false};
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
@@ -46,25 +46,16 @@ bool IsWarpReduce(const ScheduleConfig& config) {
 }
 
 bool UseContinuousDataTile(const ScheduleConfig& config) {
-  const auto& ReduceAxisContinuous = [&]() {
-    const auto& raw_reduce_axis = config.base_info->raw_reduce_axis;
-    const auto raw_data_rank = config.base_info->raw_data_rank;
-    if (raw_reduce_axis.empty()) {
-      return true;
-    }
-    for (size_t i = 1; i < raw_reduce_axis.size(); i++) {
-      if (raw_reduce_axis[i] != raw_reduce_axis[i - 1] + 1) {
-        return false;
-      }
-    }
-    return raw_reduce_axis.back() + 1 == raw_data_rank;
-  };
-  const auto& IsLastAxisReduce = [&]() {
-    const auto& raw_reduce_axis = config.base_info->raw_reduce_axis;
-    const auto raw_data_rank = config.base_info->raw_data_rank;
-    return raw_reduce_axis.back() + 1 == raw_data_rank;
-  };
-  return ReduceAxisContinuous() || IsLastAxisReduce();
+  int64_t last_axis = 0;
+  int64_t last_reduce_axis = 0;
+  for (auto axis : config.base_info->loop_transform_map) {
+    last_axis = std::max(axis, last_axis);
+  }
+  for (auto axis : config.base_info->reduce_axis) {
+    last_reduce_axis =
+        std::max(config.base_info->loop_transform_map[axis], last_reduce_axis);
+  }
+  return last_axis == last_reduce_axis;
 }
 
 class TileFirstGeneralTactic final : public ScheduleTactic {
@@ -78,6 +69,7 @@ class TileFirstGeneralTactic final : public ScheduleTactic {
   std::string TacticName() const override { return "TileFirstGeneralTactic"; }
 
  private:
+  void AlignToReduceInput(ir::IRSchedule* sch, const std::string& block_id);
   void MergeFlattenAxis(ir::IRSchedule* sch, const std::string& block_id);
   void MergeDiscreteFlattenAxis(ir::IRSchedule* sch,
                                 const std::string& block_id);
@@ -119,9 +111,16 @@ void TileFirstGeneralTactic::Init(ScheduleContext* context) {
   }
   vec_spatial_axis_first_.clear();
   vec_spatial_axis_last_.clear();
-  if (!context_->config.base_info->raw_reduce_axis.empty()) {
+
+  if (!context_->config.base_info->reduce_axis.empty()) {
+    int64_t first_reduce_axis = context_->config.base_info->data_rank - 1;
+    for (auto axis : context_->config.base_info->reduce_axis) {
+      first_reduce_axis =
+          std::min(context_->config.base_info->loop_transform_map[axis],
+                   first_reduce_axis);
+    }
     for (int32_t i = 0; i < reduce_start_idx; ++i) {
-      if (i < context_->config.base_info->raw_reduce_axis.front()) {
+      if (i < first_reduce_axis) {
         vec_spatial_axis_first_.push_back(i);
       } else {
         vec_spatial_axis_last_.push_back(i);
@@ -132,12 +131,19 @@ void TileFirstGeneralTactic::Init(ScheduleContext* context) {
 
 void TileFirstGeneralTactic::Apply(ir::IRSchedule* sch,
                                    const std::string& block_id) {
+  if (ir::IsReduceInitTensorName(block_id)) return;
+
+  AlignToReduceInput(sch, block_id);
+  VLOG(6) << "After AlignToReduceInput on block: [" << block_id
+          << "], loop nest:\n"
+          << sch->GetLoops(block_id)[0];
+
   if (UseContinuousDataTile(context_->config)) {
     VLOG(4) << "Using ApplyContinuousDataTile";
     ApplyContinuousDataTile(sch, block_id);
     return;
   }
-  if (ir::IsReduceInitTensorName(block_id)) return;
+
   MergeReduceAxis(sch, block_id);
   VLOG(6) << "After MergeReduceAxis on block: [" << block_id
           << "], loop nest:\n"
@@ -166,8 +172,6 @@ void TileFirstGeneralTactic::Apply(ir::IRSchedule* sch,
 
 void TileFirstGeneralTactic::ApplyContinuousDataTile(
     ir::IRSchedule* sch, const std::string& block_id) {
-  if (ir::IsReduceInitTensorName(block_id)) return;
-
   const auto sp_thread = context_->config.tile_config.warp_num * 32 /
                          context_->config.tile_config.tree_reduce_num;
   const auto sp_loop = context_->config.tile_config.spatial_inner_num;
@@ -275,6 +279,44 @@ void TileFirstGeneralTactic::ApplyContinuousDataTile(
   SetReduceType(sch, block_id);
 }
 
+void TileFirstGeneralTactic::AlignToReduceInput(ir::IRSchedule* sch,
+                                                const std::string& block_id) {
+  auto& loop_transform_map = context_->config.base_info->loop_transform_map;
+  if (loop_transform_map.empty()) {
+    return;
+  }
+
+  std::vector<ir::Expr> loops = sch->GetLoops(block_id);
+  std::vector<int64_t> loop_perm(loops.size());
+  std::iota(loop_perm.begin(), loop_perm.end(), 0);
+
+  const auto IsReduce = [&](int64_t axis) {
+    auto& reduce_axis = context_->config.base_info->reduce_axis;
+    return std::find(reduce_axis.begin(), reduce_axis.end(), axis) !=
+           reduce_axis.end();
+  };
+
+  std::sort(loop_perm.begin(), loop_perm.end(), [&](int64_t a, int64_t b) {
+    if (IsReduce(a) == IsReduce(b)) {
+      return loop_transform_map[a] < loop_transform_map[b];
+    }
+    return IsReduce(b);
+  });
+  VLOG(4) << "loop_perm: " << utils::Join(loop_perm, ", ");
+
+  // Reorder S/R loops seperately, otherwise reduce_init will be de-inlined.
+  std::vector<Expr> sp_loops, rd_loops;
+  for (auto i : loop_perm) {
+    if (IsReduce(i)) {
+      rd_loops.push_back(loops[i]);
+    } else if (loop_transform_map[i] != -1) {
+      sp_loops.push_back(loops[i]);
+    }
+  }
+  sch->Reorder(sp_loops);
+  sch->Reorder(rd_loops);
+}
+
 void TileFirstGeneralTactic::MergeFlattenAxis(ir::IRSchedule* sch,
                                               const std::string& block_id) {
   if (vec_flatten_axis_.size() >= 2) {
