From 88e918e277fd679343a801f9e468400174abf3b8 Mon Sep 17 00:00:00 2001 From: SamJSui Date: Tue, 18 Aug 2026 20:56:29 -0500 Subject: [PATCH] [Fix][DLight] Localize private scalar reduction buffers --- .../tvm/s_tir/dlight/gpu/general_reduction.py | 27 +++++++++- .../dlight/test_gpu_general_reduction.py | 53 +++++++++++++++++++ 2 files changed, 79 insertions(+), 1 deletion(-) diff --git a/python/tvm/s_tir/dlight/gpu/general_reduction.py b/python/tvm/s_tir/dlight/gpu/general_reduction.py index d83d8fe98b98..d3d758afc9cb 100644 --- a/python/tvm/s_tir/dlight/gpu/general_reduction.py +++ b/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 diff --git a/tests/python/s_tir/dlight/test_gpu_general_reduction.py b/tests/python/s_tir/dlight/test_gpu_general_reduction.py index 3ae666b5f3d3..383787031763 100644 --- a/tests/python/s_tir/dlight/test_gpu_general_reduction.py +++ b/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)