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 a540271e61 [DLight][CUDA] Fix undefined TX in GEMV broadcast epilogue
(#19970)
a540271e61 is described below
commit a540271e619bc6bf3f09079fd4634336fab8cec1
Author: Nanmur <[email protected]>
AuthorDate: Tue Jul 14 02:43:12 2026 +0800
[DLight][CUDA] Fix undefined TX in GEMV broadcast epilogue (#19970)
This PR fixes an undefined `TX` reference in the DLight GPU GEMV
inner-reduction schedule.
The broadcast epilogue path splits fused epilogue loops and binds the
inner loop to `threadIdx.x`, but it referenced `TX`, which is not
defined in the surrounding scope or passed into the helper. This PR uses
the existing `TR` tile factor, which is already used for the
`threadIdx.x` direction in this schedule.
A regression test is added to cover the GEMV broadcast epilogue path.
Fixes #19969
Tests:
- `python -m pytest tests/python/s_tir/dlight/test_gpu_gemv.py -q`
- `python -m pytest tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
-q`
- `python -m pytest tests/python/s_tir/dlight/test_gpu_reduction.py -q`
Note: The branch is based on `apache/tvm:main`. Local main-branch test
execution on this Windows machine could not be completed because the
available compiled TVM library is from a v0.25 build and does not match
the latest Python sources; the same tests passed in the matching local
v0.25 environment before rebasing the patch to main.
---
python/tvm/s_tir/dlight/gpu/gemv.py | 2 +-
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 2 +-
tests/python/s_tir/dlight/test_gpu_gemv.py | 43 +++++++++++++++++++
.../python/s_tir/dlight/test_gpu_low_batch_gemv.py | 49 ++++++++++++++++++++++
4 files changed, 94 insertions(+), 2 deletions(-)
diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py
b/python/tvm/s_tir/dlight/gpu/gemv.py
index 0400819974..c893c3449a 100644
--- a/python/tvm/s_tir/dlight/gpu/gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/gemv.py
@@ -288,7 +288,7 @@ class GEMV(GPUScheduleRule):
sch.reverse_compute_at(epilogue, bx)
sch.set_scope(block, 0, "shared")
_, _, *s = sch.get_loops(epilogue) # pylint:
disable=invalid-name
- _, tx = sch.split(sch.fuse(*s), factors=[None, TX])
+ _, tx = sch.split(sch.fuse(*s), factors=[None, TR])
sch.bind(tx, "threadIdx.x")
else:
sch.reverse_compute_at(epilogue, bx,
preserve_unit_loops=True)
diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
index 0e6de05325..e854ec1e26 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -492,7 +492,7 @@ class LowBatchGEMV(GPUScheduleRule):
sch.reverse_compute_at(epilogue, bx)
sch.set_scope(block, 0, "shared")
_, _, _, *s = sch.get_loops(epilogue) # pylint:
disable=invalid-name
- _, tx = sch.split(sch.fuse(*s), factors=[None, TX])
+ _, tx = sch.split(sch.fuse(*s), factors=[None, TS])
sch.bind(tx, TAG_S)
else:
sch.reverse_compute_at(epilogue, bx,
preserve_unit_loops=True)
diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py
b/tests/python/s_tir/dlight/test_gpu_gemv.py
index da62ffb1f4..14adab22ea 100644
--- a/tests/python/s_tir/dlight/test_gpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_gemv.py
@@ -1087,5 +1087,48 @@ def
test_gemv_cuda_target_without_max_shared_memory_per_block():
assert mod["main"].attrs["tirx.is_scheduled"] == 1
+def test_gemv_broadcast_epilogue():
+ # fmt: off
+ @T.prim_func(private=True, s_tir=True)
+ def before(
+ A: T.Buffer((1, 32, 1, 128), "float16"),
+ p_B: T.handle,
+ p_C: T.handle,
+ ):
+ T.func_attr({"tirx.noalias": True})
+ n = T.int32()
+ B = T.match_buffer(p_B, (1, 32, n, 128), "float16")
+ C = T.match_buffer(p_C, (1, 32, 2, 3, n), "float32")
+ C_temp = T.sblock_alloc_buffer((1, 32, 1, n), "float16")
+ for i0, i1, i2, i3, k in T.grid(1, 32, 1, n, 128):
+ with T.sblock("NT_matmul"):
+ v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1,
i2, i3, k])
+ T.reads(A[v_i0, v_i1, v_i2, v_k], B[v_i0, v_i1, v_i3, v_k])
+ T.writes(C_temp[v_i0, v_i1, v_i2, v_i3])
+ with T.init():
+ C_temp[v_i0, v_i1, v_i2, v_i3] = T.float16(0)
+ C_temp[v_i0, v_i1, v_i2, v_i3] = C_temp[
+ v_i0, v_i1, v_i2, v_i3
+ ] + A[v_i0, v_i1, v_i2, v_k] * B[v_i0, v_i1, v_i3, v_k]
+ for i0, i1, i2, i3, i4 in T.grid(1, 32, 2, 3, n):
+ with T.sblock("broadcast_epilogue"):
+ v_i0, v_i1, v_i2, v_i3, v_i4 = T.axis.remap(
+ "SSSSS", [i0, i1, i2, i3, i4]
+ )
+ T.reads(C_temp[v_i0, v_i1, 0, v_i4])
+ T.writes(C[v_i0, v_i1, v_i2, v_i3, v_i4])
+ C[v_i0, v_i1, v_i2, v_i3, v_i4] = T.Cast(
+ "float32", C_temp[v_i0, v_i1, 0, v_i4]
+ )
+
+ # fmt: on
+
+ mod = tvm.IRModule({"main": before})
+ with Target("nvidia/geforce-rtx-3090-ti"):
+ mod = dl.ApplyDefaultSchedule(dl.gpu.GEMV())(mod)
+
+ assert mod["main"].attrs["tirx.is_scheduled"] == 1
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
index 61f459c8d0..6f75191e50 100644
--- a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
@@ -17,9 +17,11 @@
# pylint: disable=missing-docstring
# ruff: noqa: E501
+from unittest import mock
import tvm.testing
from tvm.s_tir import dlight as dl
+from tvm.s_tir.dlight.gpu import low_batch_gemv
from tvm.script import tirx as T
from tvm.target import Target
@@ -557,5 +559,52 @@ def
test_low_batch_gemv_cuda_target_without_max_shared_memory_per_block():
assert mod["main"].attrs["tirx.is_scheduled"] == 1
+def test_low_batch_gemv_broadcast_epilogue():
+ # fmt: off
+ @T.prim_func(private=True, s_tir=True)
+ def before(
+ var_A: T.handle,
+ B: T.Buffer((T.int64(128), T.int64(128)), "float16"),
+ var_C: T.handle,
+ ):
+ T.func_attr({"tirx.noalias": True})
+ batch_size = T.int64()
+ A = T.match_buffer(var_A, (T.int64(1), batch_size, T.int64(1),
T.int64(128)), "float16")
+ C = T.match_buffer(var_C, (T.int64(1), batch_size, T.int64(2),
T.int64(3), T.int64(128)), "float32")
+ C_temp = T.sblock_alloc_buffer((T.int64(1), batch_size, T.int64(1),
T.int64(128)), "float16")
+ for i0, i1, i2, i3, k in T.grid(
+ T.int64(1), batch_size, T.int64(1), T.int64(128), T.int64(128)
+ ):
+ with T.sblock("NT_matmul"):
+ v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1,
i2, i3, k])
+ T.reads(A[v_i0, v_i1, v_i2, v_k], B[v_i3, v_k])
+ T.writes(C_temp[v_i0, v_i1, v_i2, v_i3])
+ with T.init():
+ C_temp[v_i0, v_i1, v_i2, v_i3] = T.float16(0)
+ C_temp[v_i0, v_i1, v_i2, v_i3] = (
+ C_temp[v_i0, v_i1, v_i2, v_i3]
+ + A[v_i0, v_i1, v_i2, v_k] * B[v_i3, v_k]
+ )
+ for i0, i1, i2, i3, i4 in T.grid(
+ T.int64(1), batch_size, T.int64(2), T.int64(3), T.int64(128)
+ ):
+ with T.sblock("broadcast_epilogue"):
+ v_i0, v_i1, v_i2, v_i3, v_i4 = T.axis.remap("SSSSS", [i0, i1,
i2, i3, i4])
+ T.reads(C_temp[v_i0, v_i1, T.int64(0), v_i4])
+ T.writes(C[v_i0, v_i1, v_i2, v_i3, v_i4])
+ C[v_i0, v_i1, v_i2, v_i3, v_i4] = T.Cast(
+ "float32", C_temp[v_i0, v_i1, T.int64(0), v_i4]
+ )
+ # fmt: on
+
+ mod = tvm.IRModule({"main": before})
+ with mock.patch.object(low_batch_gemv, "is_broadcast_epilogue",
return_value=True) as check:
+ with Target("nvidia/geforce-rtx-3090-ti"):
+ mod = dl.ApplyDefaultSchedule(dl.gpu.LowBatchGEMV(4))(mod)
+
+ check.assert_called_once()
+ assert mod["main"].attrs["tirx.is_scheduled"] == 1
+
+
if __name__ == "__main__":
tvm.testing.main()
