[CINN] Fix reduce large shape bug (#63689)

jiahy0825 · web-flow · commit df2d67bfaa90 · 2024-04-19T19:55:33.000+08:00
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.cc b/paddle/cinn/ir/group_schedule/config/group_tile_config.cc
@@ -117,7 +117,6 @@ BuildPureStaticShapeConfig(
         /* reduce_method = */ WarpReduceMethod()};
     return {{bucket_info, tile_config}};
   } else if (base_info->reduce_numel <= 2048) {
-    int64_t spatial_block = 1;
     int64_t reduce_block =
         int64_t(std::ceil(base_info->reduce_numel * 1.0 / 256.0)) * 256;
     int64_t warp_num = reduce_block / 256;
@@ -135,13 +134,9 @@ BuildPureStaticShapeConfig(
         /* reduce_method = */ BlockReduceMethod()};
     return {{bucket_info, tile_config}};
   } else {
-    int64_t spatial_block = 1;
-    int64_t reduce_block = 2048;
     int64_t warp_num = 8;
-    int64_t reduce_inner_num =
-        int64_t(std::ceil(base_info->reduce_numel * 1.0 / 256.0));
     int64_t spatial_inner_num = 1;
-    int64_t tree_reduce_num = reduce_block / reduce_inner_num;
+    int64_t tree_reduce_num = warp_num * 32;
     BucketInfo bucket_info{/* sp_lower_bound = */ 1,
                            /* sp_upper_bound = */ kMaxNumel,
                            /* rb_lower_bound = */ 2049,
@@ -150,7 +145,7 @@ BuildPureStaticShapeConfig(
         /* warp_num = */ warp_num,
         /* tree_reduce_num = */ tree_reduce_num,
         /* spatial_inner_num = */ spatial_inner_num,
-        /* reduce_method = */ NoneReduceMethod()};
+        /* reduce_method = */ BlockReduceMethod()};
     return {{bucket_info, tile_config}};
   }
 }
diff --git a/test/ir/pir/cinn/CMakeLists.txt b/test/ir/pir/cinn/CMakeLists.txt
@@ -2,6 +2,7 @@ add_subdirectory(adt)
 add_subdirectory(symbolic)
 add_subdirectory(inference)
 add_subdirectory(sub_graphs)
+add_subdirectory(performance)
 
 if(WITH_GPU)
   file(
diff --git a/test/ir/pir/cinn/performance/CMakeLists.txt b/test/ir/pir/cinn/performance/CMakeLists.txt
@@ -0,0 +1,40 @@
+if(WITH_GPU)
+  file(
+    GLOB CINN_PERFORMANCE_TEST
+    RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+    "test_*.py")
+
+  foreach(cinn_pir_test_name ${CINN_PIR_SYMBOLIC_TEST})
+    string(REGEX REPLACE ".py" "" cinn_pir_test_name ${cinn_pir_test_name})
+    add_test(
+      NAME ${cinn_pir_test_name}
+      COMMAND
+        ${CMAKE_COMMAND} -E env
+        PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
+        FLAGS_check_infer_symbolic=1 FLAGS_enable_pir_api=1
+        FLAGS_cinn_bucket_compile=True FLAGS_prim_enable_dynamic=true
+        FLAGS_pir_apply_shape_optimization_pass=1
+        FLAGS_group_schedule_tiling_first=1 FLAGS_cinn_new_group_scheduler=1
+        ${PYTHON_EXECUTABLE}
+        ${CMAKE_CURRENT_SOURCE_DIR}/${cinn_pir_test_name}.py
+      WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+    set_tests_properties(${cinn_pir_test_name} PROPERTIES LABELS
+                                                          "RUN_TYPE=CINN")
+
+    add_test(
+      NAME ${cinn_pir_test_name}_stride_read
+      COMMAND
+        ${CMAKE_COMMAND} -E env
+        PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
+        FLAGS_check_infer_symbolic=1 FLAGS_enable_pir_api=1
+        FLAGS_cinn_bucket_compile=True FLAGS_prim_enable_dynamic=true
+        FLAGS_pir_apply_shape_optimization_pass=1
+        FLAGS_group_schedule_tiling_first=1 FLAGS_cinn_new_group_scheduler=1
+        FLAGS_support_reduce_stride_read=1 ${PYTHON_EXECUTABLE}
+        ${CMAKE_CURRENT_SOURCE_DIR}/${cinn_pir_test_name}.py
+      WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+    set_tests_properties(${cinn_pir_test_name}_stride_read
+                         PROPERTIES LABELS "RUN_TYPE=CINN")
+  endforeach()
+
+endif()
diff --git a/test/ir/pir/cinn/performance/test_cinn_large_shape_reduce.py b/test/ir/pir/cinn/performance/test_cinn_large_shape_reduce.py
@@ -0,0 +1,69 @@
+# 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 unittest
+
+import numpy as np
+
+import paddle
+
+
+class LayerCase(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, var_0):
+        return var_0.sum([0, 2, 3])
+
+
+def create_paddle_inputs():
+    inputs = (paddle.rand(shape=[16, 128, 256, 256], dtype=paddle.float32),)
+    return inputs
+
+
+class TestLayer(unittest.TestCase):
+    def setUp(self):
+        self.inputs = create_paddle_inputs()
+        self.net = LayerCase()
+
+    def train(self, net, to_static, with_prim=False, with_cinn=False):
+        if to_static:
+            paddle.set_flags({'FLAGS_prim_all': with_prim})
+            if with_cinn:
+                build_strategy = paddle.static.BuildStrategy()
+                build_strategy.build_cinn_pass = True
+                net = paddle.jit.to_static(
+                    net, build_strategy=build_strategy, full_graph=True
+                )
+            else:
+                net = paddle.jit.to_static(net, full_graph=True)
+        paddle.seed(123)
+        outs = net(*self.inputs)
+        return outs
+
+    def test_ast_prim_cinn(self):
+        st_out = self.train(self.net, to_static=True)
+        cinn_out = self.train(
+            self.net, to_static=True, with_prim=True, with_cinn=True
+        )
+        for st, cinn in zip(
+            paddle.utils.flatten(st_out), paddle.utils.flatten(cinn_out)
+        ):
+            np.testing.assert_allclose(
+                st.numpy(), cinn.numpy(), atol=1e-6, rtol=1e-5
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()
