PaddlePaddle
diff --git a/‎paddle/cinn/hlir/pe/reduction.cc
Copy file name to clipboard
+34Lines changed: 34 additions & 0 deletions b/‎paddle/cinn/hlir/pe/reduction.cc
Copy file name to clipboard
+34Lines changed: 34 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/pe/reduction.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/cinn/hlir/pe/reduction.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/cinn/ir/group_schedule/config/group_tile_config.cc
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions b/‎paddle/cinn/ir/group_schedule/config/group_tile_config.cc
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions
diff --git a/‎paddle/cinn/ir/group_schedule/config/group_tile_config.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/ir/group_schedule/config/group_tile_config.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
@@ -1205,6 +1205,40 @@ std::string DiscreteReduceExternalFuncName(const ir::Expr& op,
   return "";
 }
 
+std::string IntervalReduceExternalFuncName(const ir::Expr& op,
+                                           const ir::Expr& tensor) {
+  CHECK_NOTNULL(tensor.as_tensor());
+  if (op.As<ir::Add>()) {
+    if (tensor.as_tensor()->type().is_bool()) {
+      return "cinn_interval_reduce_any_internal_shm";
+    }
+    return "cinn_interval_reduce_sum" +
+           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
+  } else if (op.As<ir::Mul>()) {
+    if (tensor.as_tensor()->type().is_bool()) {
+      return "cinn_interval_reduce_all_internal_shm";
+    }
+    return "cinn_interval_reduce_prod" +
+           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
+  } else if (op.As<ir::Max>()) {
+    return "cinn_interval_reduce_max" +
+           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
+  } else if (op.As<ir::Min>()) {
+    return "cinn_interval_reduce_min" +
+           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
+  } else if (op.As<ir::And>()) {
+    return "cinn_interval_reduce_all_internal_shm";
+  } else if (op.As<ir::Or>()) {
+    return "cinn_interval_reduce_any_internal_shm";
+  } else {
+    std::stringstream ss;
+    ss << op;
+    PADDLE_THROW(::common::errors::InvalidArgument(
+        "Reduce type %s is not supported yet!", ss.str()));
+  }
+  return "";
+}
+
 }  // namespace pe
 }  // namespace hlir
 }  // namespace cinn
@@ -474,6 +474,9 @@ std::string CrossThreadReduceExternalFuncName(const ir::Expr& op,
 std::string DiscreteReduceExternalFuncName(const ir::Expr& op,
                                            const ir::Expr& tensor);
 
+std::string IntervalReduceExternalFuncName(const ir::Expr& op,
+                                           const ir::Expr& tensor);
+
 std::string Type2StrForReduce(cinn::common::Type type);
 }  // namespace pe
 }  // namespace hlir
 
@@ -96,21 +96,21 @@ std::shared_ptr<ScheduleConfig::BaseInfo> InitBasicInfo(
   base_info->reduce_tensor_names = group_info->reduce_var_names;
   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->data_space = group_info->data_space;
   base_info->loop_strides = group_info->loop_strides;
 
   std::set<int64_t> reduce_dim_loc;
   for (int64_t dim : group_info->reduce_axis) {
     if (dim < 0) {
-      dim += base_info->data_rank;
+      dim += base_info->data_space.size();
     }
     base_info->reduce_axis.push_back(dim);
     reduce_dim_loc.insert(dim);
   }
 
   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->data_space.size(); ++i) {
     if (reduce_dim_loc.count(i)) {
       if (group_info->data_space[i] == -1) base_info->has_dynamic_reduce = true;
       base_info->reduce_numel *= group_info->data_space[i];
@@ -121,7 +121,7 @@ std::shared_ptr<ScheduleConfig::BaseInfo> InitBasicInfo(
     }
   }
   base_info->is_reduce_all =
-      (base_info->reduce_axis.size() == base_info->data_rank);
+      (base_info->reduce_axis.size() == base_info->data_space.size());
 
   for (int64_t i = 0; i < group_info->data_space.size(); ++i) {
     if (group_info->data_space[i] == 1) continue;
 
@@ -32,7 +32,7 @@ struct ScheduleConfig {
   struct BaseInfo {
     std::vector<int64_t> reduce_axis;
     std::vector<int64_t> loop_strides;
-    int64_t data_rank;
+    std::vector<int64_t> data_space;
     int64_t reduce_numel;
     int64_t spatial_numel;
     bool has_dynamic_spatial{false};