[CINN Reduce] Fix reduce small number (#63411)

jiahy0825 · web-flow · commit f33ce57d75cb · 2024-04-16T10:27:59.000+08:00
* [CINN Reduce] Fix reduce small number

* Fix SplitWarpNumber bug
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.cc b/paddle/cinn/ir/group_schedule/config/group_tile_config.cc
@@ -102,7 +102,8 @@ BuildPureStaticShapeConfig(
   } else if (base_info->reduce_numel <= 256) {
     // warp reduce
     int64_t reduce_block = Next2Power(base_info->reduce_numel);
-    int64_t spatial_inner_num = 256 / reduce_block;
+    int64_t spatial_inner_num =
+        std::min(256 / reduce_block, base_info->spatial_numel);
     int64_t tree_reduce_num = 32;
     int64_t warp_num = 8;
     BucketInfo bucket_info{/* sp_lower_bound = */ 1,
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
@@ -219,19 +219,21 @@ void TileFirstGeneralTactic::SplitWarpNumber(ir::IRSchedule* sch,
   };
   if (!IsWarpNumGT(1)) return;
 
-  const auto LimitWarpNum = [&](const ir::Expr& loop, ScheduleConfig* config) {
+  const auto GetMinimalWarpNum = [&](const ir::Expr& loop,
+                                     const ScheduleConfig& config) -> int {
     ir::Expr extent = loop.As<ir::For>()->extent;
     common::cas_intervals_t var_intervals =
         common::CollectVarIntervalsOfExprs({extent});
     common::SymbolicExprAnalyzer analyzer(var_intervals);
     const auto& proved_gt =
-        analyzer.ProveGT(ir::Expr(config->tile_config.warp_num), extent);
+        analyzer.ProveGT(ir::Expr(config.tile_config.warp_num * 32), extent);
     if (proved_gt.value_or(false)) {
       ir::Expr upper_bound = analyzer.UpperBound(extent);
       if (upper_bound.is_constant()) {
-        config->tile_config.warp_num = upper_bound.get_constant();
+        return (static_cast<int>(upper_bound.get_constant()) + 31) / 32;
       }
     }
+    return config.tile_config.warp_num;
   };
 
   auto loops = sch->GetLoops(block_id);
@@ -248,9 +250,9 @@ void TileFirstGeneralTactic::SplitWarpNumber(ir::IRSchedule* sch,
     }
   } else if (IsWarpReduce(context_->config)) {
     // get num warp from flatten num
-    LimitWarpNum(loops[0], &(context_->config));
-    int thread_y = context_->config.tile_config.warp_num * 32 /
-                   context_->config.tile_config.tree_reduce_num;
+    int minimal_warp_number = GetMinimalWarpNum(loops[0], context_->config);
+    int thread_y =
+        minimal_warp_number * 32 / context_->config.tile_config.tree_reduce_num;
     sch->Split(loops[0], std::vector<int>({-1, thread_y}));
 
     if (IsReduceBlock(context_->config, block_id) &&
diff --git a/test/ir/pir/cinn/test_cinn_reduce_sum.py b/test/ir/pir/cinn/test_cinn_reduce_sum.py
@@ -0,0 +1,57 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import sys
+import unittest
+from os.path import dirname
+
+import numpy as np
+
+import paddle
+
+sys.path.append(dirname(dirname(__file__)))
+import utils
+
+
+class ReduceSumSubGraph(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return paddle.max(x * 2, axis=[-1], keepdim=True)
+
+
+class TestReduceSumSubGraph(unittest.TestCase):
+    def setUp(self):
+        paddle.seed(2024)
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.x = paddle.randn([8, 4], dtype="float32")
+
+    def eval(self, use_cinn):
+        net = ReduceSumSubGraph()
+        net.eval()
+        net = utils.apply_to_static(net, use_cinn)
+        out = net(self.x)
+        return out
+
+    def test_eval(self):
+        cinn_out = self.eval(use_cinn=True)
+        dy_out = self.eval(use_cinn=False)
+
+        np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-6)
+
+
+if __name__ == '__main__':
+    unittest.main()
