[CINN] Adjust apply continuous reduce principle (#65336)

jiahy0825 · web-flow · commit 61575200b933 · 2024-06-21T14:22:15.000+08:00
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,17 +46,25 @@ bool IsWarpReduce(const ScheduleConfig& config) {
 }
 
 bool UseContinuousDataTile(const ScheduleConfig& config) {
-  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;
+  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;
     }
-  }
-  return raw_reduce_axis.back() + 1 == raw_data_rank;
+    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();
 }
 
 class TileFirstGeneralTactic final : public ScheduleTactic {
