This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 4f849cfa88 [Fix][DLight] Localize private scalar reduction buffers
(#20160)
4f849cfa88 is described below
commit 4f849cfa88b17039fa009e76ada18eb11f756d33
Author: Sam Sui <[email protected]>
AuthorDate: Fri Aug 28 02:55:25 2026 -0500
[Fix][DLight] Localize private scalar reduction buffers (#20160)
DLight's scalar-output path returns before the normal reduction-buffer
scope assignment. For length-one `argmin`, the paired index and value
temporaries therefore remain global. The index is read by the final
result block and must remain global, while the value temporary ends in
the producer kernel and becomes an illegal internal global allocation
during CUDA code generation.
This change assigns local scope only to root-allocated global buffers
that are produced under unit-extent loops and are not accessed by
another block. The index temporary therefore retains global scope.
A focused regression covers the length-one case and verifies that
non-unit reductions retain global scope.
Fixes #20060.
Testing:
- `python -m pytest
tests/python/s_tir/dlight/test_gpu_general_reduction.py -q` — 8 passed
- `python -m pytest tests/python/s_tir/dlight -q` — 101 passed, 1
skipped, 1 xfailed
- Adjacent Relax tests — 52 passed
- Exact issue reproducer — passed
- Changed-file pre-commit hooks — passed
---
python/tvm/s_tir/dlight/gpu/general_reduction.py | 27 ++++++++++-
.../s_tir/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 d83d8fe98b..d3d758afc9 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 @@ class GeneralReduction(GPUScheduleRule):
# 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 3ae666b5f3..3837870317 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)
