From e3738a7e0e756efef44a162a72af96980ad3c799 Mon Sep 17 00:00:00 2001 From: Nanmur <1067841003@qq.com> Date: Thu, 9 Jul 2026 12:14:17 +0800 Subject: [PATCH] [DLight][CUDA] Fix undefined TX in GEMV schedules The GEMV and LowBatchGEMV inner-reduction schedules handle broadcast epilogues by splitting fused epilogue loops and binding the inner loop to a CUDA thread axis. These paths referenced TX, but TX is not defined in either schedule scope or passed into the inner scheduling helpers. Use the existing tile factors for the corresponding bound axes: TR for GEMV's threadIdx.x path and TS for LowBatchGEMV's TAG_S path. Add regression tests covering both broadcast epilogue branches. --- 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 ++++++++++++++++ .../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 0400819974cb..c893c3449a95 100644 --- a/python/tvm/s_tir/dlight/gpu/gemv.py +++ b/python/tvm/s_tir/dlight/gpu/gemv.py @@ -288,7 +288,7 @@ def apply( 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 0e6de0532525..e854ec1e2639 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 @@ def apply( 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 da62ffb1f4ee..14adab22eae3 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 before( 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 61f459c8d07c..6f75191e504b 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 before(var_A: T.handle, B: T.Buffer((T.int64(128), T.int64(128)), "float16") 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()