From 156d82c93142a871f23958544091e09400d95392 Mon Sep 17 00:00:00 2001 From: SamJSui Date: Tue, 18 Aug 2026 08:36:52 -0500 Subject: [PATCH] [Fix][DLight] Handle rank-one GEMV cache loads The DLight GPU GEMV rule assumed that the final two loops around a local cache read were both owned by the cache stage. Rank-one vector inputs create only one cache loop, so the rule included the shared placement loop and attempted to fuse an imperfect nest. Record the cache loop count before compute_at and recover that innermost suffix after placement. Use the sole cache loop directly for rank-one inputs while preserving the existing final-two-loop fusion for higher-rank inputs. Add a focused regression that verifies rank-one GEMV scheduling and vectorized local loading. --- python/tvm/s_tir/dlight/gpu/gemv.py | 11 +++++---- tests/python/s_tir/dlight/test_gpu_gemv.py | 28 ++++++++++++++++++++++ 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py b/python/tvm/s_tir/dlight/gpu/gemv.py index 1c451b964faa..1a71fa697e52 100644 --- a/python/tvm/s_tir/dlight/gpu/gemv.py +++ b/python/tvm/s_tir/dlight/gpu/gemv.py @@ -164,13 +164,14 @@ def apply( and shared_mem_usage.value <= max_smem ) - # vectorize load A - # (TODO) this is now actually problematic since the number of loops is dependent on the - # number of dimensions of A_q Aq_local = sch.cache_read(rf, read_buffer_index=1, storage_scope="local") + num_cache_loops = len(sch.get_loops(block=Aq_local)) sch.compute_at(Aq_local, r, preserve_unit_loops=True) - s_local, r_local = sch.get_loops(block=Aq_local)[-2:] - fused_load = sch.fuse(s_local, r_local) + cache_loops = sch.get_loops(block=Aq_local)[-num_cache_loops:] + if len(cache_loops) == 1: + fused_load = cache_loops[0] + else: + fused_load = sch.fuse(*cache_loops[-2:]) aq_vec_len = max(1, VEC_LOAD // get_bytes(sch.get(Aq_local).reads[0].buffer.dtype)) fused_load, vec_load = sch.split( fused_load, factors=[None, aq_vec_len], preserve_unit_iters=True diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py b/tests/python/s_tir/dlight/test_gpu_gemv.py index 8cfcd1bb8169..e7454f5505dd 100644 --- a/tests/python/s_tir/dlight/test_gpu_gemv.py +++ b/tests/python/s_tir/dlight/test_gpu_gemv.py @@ -1114,6 +1114,34 @@ def before( assert mod["main"].attrs["tirx.is_scheduled"] == 1 +def test_gemv_rank_one_vector_input(): + @T.prim_func(private=True, s_tir=True) + def before( + matrix: T.Buffer((2, 2), "float32"), + vector: T.Buffer((2,), "float32"), + output: T.Buffer((2,), "float32"), + ): + T.func_attr({"tirx.noalias": True}) + for i, k in T.grid(2, 2): + with T.sblock("gemv"): + vi, vk = T.axis.remap("SR", [i, k]) + T.reads(matrix[vi, vk], vector[vk]) + T.writes(output[vi]) + with T.init(): + output[vi] = T.float32(0) + output[vi] += matrix[vi, vk] * vector[vk] + + 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 + sch = tvm.s_tir.Schedule(mod) + vector_local = sch.get_sblock("vector_local") + vector_load_loop = sch.get(sch.get_loops(vector_local)[-1]) + assert vector_load_loop.kind == tvm.tirx.ForKind.VECTORIZED + + def test_gemv_broadcast_epilogue(): # fmt: off @T.prim_func(private=True, s_tir=True)