[CINN] Fix bugs of calculation errors in grid reduce (#68937)

lshpku · web-flow · commit e1779ba939ef · 2024-10-25T14:36:13.000+08:00
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_util.cc b/paddle/cinn/hlir/framework/pir/trivial_op_util.cc
@@ -386,6 +386,10 @@ ExprSetFinder FilterLoadByTensor(const ir::Tensor& tensor) {
 ExprSetFinder ChildFors =
     Collector([](const ir::Expr* e) { return e->As<ir::For>(); }, "ChildFors");
 
+ExprSetFinder ChildIfThenElses =
+    Collector([](const ir::Expr* e) { return e->As<ir::IfThenElse>(); },
+              "ChildIfThenElses");
+
 ExprSetFinder FindFather(const ir::Expr& root) {
   const auto& f = [root](const auto& child) -> ExprSet {
     ExprSetFinder find_child =
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_util.h b/paddle/cinn/hlir/framework/pir/trivial_op_util.h
@@ -191,6 +191,8 @@ extern ExprSetFinder ChildTensorStores;
 
 extern ExprSetFinder ChildFors;
 
+extern ExprSetFinder ChildIfThenElses;
+
 ExprSetFinder IsForIterVar(const ir::Var& var);
 
 ExprSetFinder FilterLoadByTensor(const ir::Tensor& tensor);
diff --git a/paddle/cinn/optim/replace_cross_block_reduction.cc b/paddle/cinn/optim/replace_cross_block_reduction.cc
@@ -26,6 +26,8 @@
 
 namespace cinn {
 
+using hlir::framework::pir::trivial_fusion_detail::ExprSetFinderUtils::
+    ChildIfThenElses;
 using hlir::framework::pir::trivial_fusion_detail::ExprSetFinderUtils::
     ChildTensorLoads;
 
@@ -78,6 +80,14 @@ ir::Expr GetRightOperand(const ir::Expr& expr) {
       ::common::errors::InvalidArgument("Not a supported reduce op: %s", expr));
 }
 
+void RemoveGridReduceAxisFromIfCondition(ir::Expr* expr) {
+  std::vector<ir::Expr> if_exprs = ChildIfThenElses(*expr);
+  for (auto& if_expr : if_exprs) {
+    auto& condition = if_expr.As<ir::IfThenElse>()->condition;
+    optim::ReplaceVarWithExpr(&condition, ir::Var("blockIdx.y"), ir::Expr(0));
+  }
+}
+
 struct BaseMutator : public ir::IRMutator<> {
   void operator()(ir::Expr* expr) { Visit(expr); }
 
@@ -350,16 +360,22 @@ struct CrossBlockReductionReplacer : public BaseMutator {
     auto* load_node = rvalue.As<ir::Load>();
     ir::Tensor rf_tensor = load_node->tensor.as_tensor_ref();
 
-    // indices[0] of the load is the reduce axis of grid reduce. We set it to
-    // zero to compute the spatial part of load index.
-    // Note: it's hard to check whether indices[0] is the correct reduce axis
-    // in this pass, but it should be because this load is generated by doing
-    // rfactor with rf_axis=0.
-    load_node->indices[0] = ir::Expr(0);
-    ir::Expr spatial_index = load_node->index();
+    // The load's indices are like [ blockIdx.y, <spatial_index>... ].
+    // The loaded tensor's shape is like [ gridDim.y, <spatial_size>... ].
+    ir::Expr spatial_index = [&]() {
+      load_node->indices[0] = ir::Expr(0);
+      return load_node->index();
+    }();
+    ir::Expr spatial_size = [&]() {
+      load_node->indices[0] = ir::Expr(1);
+      for (int i = 1; i < load_node->indices.size(); i++) {
+        load_node->indices[i] = ir::Expr(0);
+      }
+      return load_node->index();
+    }();
 
-    store_node->value = lang::CallExtern(
-        func_name, {rf_tensor, num_spatial_threads, spatial_index});
+    store_node->value =
+        lang::CallExtern(func_name, {rf_tensor, spatial_size, spatial_index});
   }
 
   void Visit(const ir::_LoweredFunc_* expr, ir::Expr* op) override {
@@ -393,6 +409,10 @@ struct CrossBlockReductionReplacer : public BaseMutator {
 
     if (!IsGridReduce(expr)) {
       if (is_after_grid_reduce_) {
+        // For trivial ops that are after grid reduce, we need to remove the
+        // `blockIdx.y` axis from their if conditions, otherwise they may not
+        // be executed in the last done block.
+        RemoveGridReduceAxisFromIfCondition(op);
         *op = WrapInLastBlockDone(op);
       }
       return;
diff --git a/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh b/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh
@@ -720,21 +720,16 @@ EXPAND_REDUCE_FP16_MACRO(CINN_BLOCK_REDUCE_IMPL)
 
 #undef CINN_BLOCK_REDUCE_IMPL
 
-#define CINN_GRID_REDUCE_IMPL(REDUCE_TYPE, init_value, DTYPE)                    \
-  int tid = (threadIdx.z * blockDim.y + threadIdx.y) * blockDim.x + threadIdx.x; \
-  if (tid < spatial_threads) {                                                   \
-    DTYPE tmp_val = init_value;                                                  \
-    for (int y = 0; y < gridDim.y; y++) {                                        \
-      tmp_val = cinn_##REDUCE_TYPE(tmp_val, mem[(y * gridDim.x + blockIdx.x) * spatial_threads + tid]); \
-    }                                                                            \
-    mem[blockIdx.x * spatial_threads + tid] = tmp_val;                           \
-  }                                                                              \
-  __syncthreads();                                                               \
-  return mem[spatial_index];
-
-#define CINN_GRID_REDUCE_MACRO(REDUCE_TYPE, INITIAL_VALUE, DTYPE)                \
-  __device__ inline DTYPE cinn_grid_reduce_##REDUCE_TYPE(DTYPE* mem, int spatial_threads, int spatial_index) { \
-    CINN_GRID_REDUCE_IMPL(REDUCE_TYPE, (DTYPE)(INITIAL_VALUE), DTYPE);           \
+#define CINN_GRID_REDUCE_IMPL(REDUCE_TYPE, init_value, DTYPE)                       \
+  DTYPE tmp_val = init_value;                                                       \
+  for (int y = 0; y < gridDim.y; y++) {                                             \
+      tmp_val = cinn_##REDUCE_TYPE(tmp_val, mem[y * spatial_size + spatial_index]); \
+  }                                                                                 \
+  return tmp_val;
+
+#define CINN_GRID_REDUCE_MACRO(REDUCE_TYPE, INITIAL_VALUE, DTYPE)                   \
+  __device__ inline DTYPE cinn_grid_reduce_##REDUCE_TYPE(const DTYPE* mem, int spatial_size, int spatial_index) { \
+    CINN_GRID_REDUCE_IMPL(REDUCE_TYPE, (DTYPE)(INITIAL_VALUE), DTYPE);              \
   }
 
 EXPAND_REDUCE_INT32_MARCO(CINN_GRID_REDUCE_MACRO)
diff --git a/test/cpp/pir/cinn/replace_cross_block_reduction_test.cc b/test/cpp/pir/cinn/replace_cross_block_reduction_test.cc
@@ -78,7 +78,7 @@ TEST(CrossBlockReductionReplacer, SRLayout) {
             ScheduleBlock(B)
             {
               i0_0, i1 = axis.bind(i, reduce_k)
-              B[i0_0] = cinn_grid_reduce_sum_fp32(Tensor(A, [8,16]), 1, i0_0)
+              B[i0_0] = cinn_grid_reduce_sum_fp32(Tensor(A, [8,16]), 16, i0_0)
             }
           }
         }
@@ -166,7 +166,7 @@ TEST(CrossBlockReductionReplacer, RSLayout) {
               ScheduleBlock(B)
               {
                 i0_0, i1_0, i2 = axis.bind(i, j, reduce_k)
-                B[i0_0, i1_0] = cinn_grid_reduce_max_fp32(Tensor(A, [8,4,32]), 32, ((32 * i0_0) + i1_0))
+                B[i0_0, i1_0] = cinn_grid_reduce_max_fp32(Tensor(A, [8,4,32]), 128, ((32 * i0_0) + i1_0))
               }
             }
           }
diff --git a/test/ir/pir/cinn/test_grid_reduce.py b/test/ir/pir/cinn/test_grid_reduce.py
@@ -0,0 +1,97 @@
+# 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
+import utils
+
+import paddle
+
+
+class TestGridReduce(unittest.TestCase):
+    def eval(self, dy_compute, init_inputs):
+        paddle.seed(2024)
+        inputs = init_inputs()
+        dy_out = dy_compute(*inputs)
+
+        static_compute = utils.apply_to_static(dy_compute, use_cinn=True)
+        st_out = static_compute(*inputs)
+
+        for a, b in zip(
+            paddle.utils.flatten(dy_out), paddle.utils.flatten(st_out)
+        ):
+            numpy.testing.assert_allclose(a, b, atol=1e-3, rtol=1e-4)
+
+    def test_all_reduce(self):
+        def func(x):
+            return paddle.sum(x)
+
+        def init():
+            x = paddle.randn([32, 128, 256])
+            return (x,)
+
+        self.eval(func, init)
+
+    def test_continuous_reduce(self):
+        def func(x):
+            return paddle.sum(x, axis=(0, 2, 3))
+
+        def init():
+            x = paddle.randn([64, 128, 56, 56])
+            return (x,)
+
+        self.eval(func, init)
+
+    def test_discrete_reduce(self):
+        def func(x):
+            return paddle.sum(x, axis=(0, 1, 2))
+
+        def init():
+            x = paddle.randn([256, 28, 28, 80])
+            return (x,)
+
+        self.eval(func, init)
+
+    def test_multiple_reduce(self):
+        def func(x):
+            n = 512 * 14 * 14
+            sum_x = paddle.sum(x, axis=(0, 2, 3))
+            sum_x2 = paddle.sum(x * x, axis=(0, 2, 3))
+            mean_x = sum_x / n
+            mean_x2 = sum_x2 / n
+            mean_x_2 = mean_x * mean_x
+            return mean_x2 - mean_x_2
+
+        def init():
+            x = paddle.randn([512, 256, 14, 14])
+            return (x,)
+
+        self.eval(func, init)
+
+    def test_multiple_downstream(self):
+        def func(a, b):
+            a = paddle.sum(a, axis=(0, 1))
+            return a, a + b
+
+        def init():
+            a = paddle.randn([400, 300, 20, 10])
+            b = paddle.randn([20, 10])
+            return a, b
+
+        self.eval(func, init)
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -78,7 +78,7 @@ TEST(CrossBlockReductionReplacer, SRLayout) {`
`78`	`78`	`ScheduleBlock(B)`
`79`	`79`	`{`
`80`	`80`	`i0_0, i1 = axis.bind(i, reduce_k)`
`81`		`- B[i0_0] = cinn_grid_reduce_sum_fp32(Tensor(A, [8,16]), 1, i0_0)`
	`81`	`+ B[i0_0] = cinn_grid_reduce_sum_fp32(Tensor(A, [8,16]), 16, i0_0)`
`82`	`82`	`}`
`83`	`83`	`}`
`84`	`84`	`}`
`@@ -166,7 +166,7 @@ TEST(CrossBlockReductionReplacer, RSLayout) {`
`166`	`166`	`ScheduleBlock(B)`
`167`	`167`	`{`
`168`	`168`	`i0_0, i1_0, i2 = axis.bind(i, j, reduce_k)`
`169`		`- B[i0_0, i1_0] = cinn_grid_reduce_max_fp32(Tensor(A, [8,4,32]), 32, ((32 * i0_0) + i1_0))`
	`169`	`+ B[i0_0, i1_0] = cinn_grid_reduce_max_fp32(Tensor(A, [8,4,32]), 128, ((32 * i0_0) + i1_0))`
`170`	`170`	`}`
`171`	`171`	`}`
`172`	`172`	`}`