[Fix][DLight] Localize private scalar reduction buffers

python/tvm/s_tir/dlight/gpu/general_reduction.py

@@ -20,7 +20,7 @@
 from tvm import arith, ir, s_tir, tirx
 from tvm.target import Target
 
-from ..analysis import normalize_prim_func
+from ..analysis import get_root_block, normalize_prim_func
 from ..base import try_inline_contiguous_spatial
 from .base import GPUScheduleRule
 
@@ -65,6 +65,31 @@ def apply(  # pylint: disable=too-many-locals
             # Add a unit thread loop so the final write happens inside a valid
             # GPU thread environment.
             if num_last_block_iter == 0:
+                # Allocation planning can move a reduction buffer inside its
+                # producer kernel when all surrounding loops are trivial. Give
+                # buffers private to that kernel an explicit local scope, while
+                # preserving global scope for buffers accessed by another block.
+                blocks = [sch.get(info.block_rv) for info in block_infos]
+                alloc_buffers = list(sch.get(get_root_block(sch)).alloc_buffers)
+                analyzer = arith.Analyzer()
+                for block_index, (info, block) in enumerate(zip(block_infos[:-1], blocks[:-1])):
+                    loops = sch.get_loops(info.block_rv)
+                    if not all(analyzer.can_prove_equal(sch.get(loop).extent, 1) for loop in loops):
+                        continue
+
+                    other_block_buffers = [
+                        region.buffer
+                        for other_index, other_block in enumerate(blocks)
+                        if other_index != block_index
+                        for region in (*other_block.reads, *other_block.writes)
+                    ]
+                    for buffer_index, write in enumerate(block.writes):
+                        buffer = write.buffer
+                        is_allocated = any(buffer.same_as(other) for other in alloc_buffers)
+                        is_cross_block = any(buffer.same_as(other) for other in other_block_buffers)
+                        if buffer.scope() == "global" and is_allocated and not is_cross_block:
+                            sch.set_scope(block_infos[block_index].block_rv, buffer_index, "local")
+
                 # Put every block (both the running reductions and the final
                 # scalar write) inside a trivial GPU thread. The very first block
                 # gets a `blockIdx.x` wrapper so that kernels still have a unique

tests/python/s_tir/dlight/test_gpu_general_reduction.py

@@ -35,6 +35,59 @@ def _check(mod_before: IRModule, mod_after: IRModule):
     assert_structural_equal(mod, mod_after)
 
 
+def _make_scalar_argmin(length):
+    @I.ir_module(s_tir=True)
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main(x: T.Buffer((T.int64(length),), "float32"), x_red: T.Buffer((), "int64")):
+            T.func_attr({"tirx.noalias": True})
+            x_red_temp_v0 = T.sblock_alloc_buffer((), "int64")
+            x_red_temp_v1 = T.sblock_alloc_buffer(())
+            for k in range(T.int64(length)):
+                with T.sblock("x_red_temp"):
+                    v_k = T.axis.reduce(T.int64(length), k)
+                    T.reads(x[v_k])
+                    T.writes(x_red_temp_v0[()], x_red_temp_v1[()])
+                    with T.init():
+                        x_red_temp_v0[()] = T.int64(-1)
+                        x_red_temp_v1[()] = T.max_value("float32")
+                    value_is_smaller = x_red_temp_v1[()] < x[v_k]
+                    value_is_equal = x_red_temp_v1[()] == x[v_k]
+                    index_is_smaller = x_red_temp_v0[()] < v_k
+                    new_index: T.int64 = T.Select(
+                        value_is_smaller or (value_is_equal and index_is_smaller),
+                        x_red_temp_v0[()],
+                        v_k,
+                    )
+                    new_value: T.float32 = T.Select(
+                        value_is_smaller,
+                        x_red_temp_v1[()],
+                        x[v_k],
+                    )
+                    x_red_temp_v0[()] = new_index
+                    x_red_temp_v1[()] = new_value
+            with T.sblock("x_red"):
+                vi = T.axis.spatial(1, T.int64(0))
+                T.reads(x_red_temp_v0[()])
+                T.writes(x_red[()])
+                x_red[()] = x_red_temp_v0[()]
+
+    return Before
+
+
+def test_scalar_argmin_reduction_value_scope():
+    for length, expected_scope in ((1, "local"), (2, "global"), (3, "global")):
+        target = Target("nvidia/geforce-rtx-3090-ti")
+        with target:
+            mod = dl.ApplyDefaultSchedule(  # pylint: disable=not-callable
+                dl.gpu.GeneralReduction(),
+            )(_make_scalar_argmin(length))
+
+        index_temp, value_temp = mod["main"].body.block.alloc_buffers
+        assert index_temp.scope() == "global"
+        assert value_temp.scope() == expected_scope
+
+
 def test_softmax_1():
     # fmt: off
     @I.ir_module(s_tir=True)