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 d536d413e6 [Fix][DLight] Handle rank-one GEMV cache loads (#20158)
d536d413e6 is described below
commit d536d413e66cfd018d45f5400301b0d747e87b81
Author: Sam Sui <[email protected]>
AuthorDate: Tue Aug 18 17:55:39 2026 -0500
[Fix][DLight] Handle rank-one GEMV cache loads (#20158)
DLight's GPU GEMV rule assumes that the final two loops around a local
cache
read are both owned by the cache stage. Rank-one vector inputs create
only one
cache loop, causing the rule to include the shared placement loop and
attempt
to fuse an imperfect loop nest.
This change records the cache loop count before `compute_at` and
recovers that
innermost suffix afterward. A single cache loop is used directly, while
the
existing final-two-loop fusion is preserved for higher-rank inputs.
A focused regression verifies that rank-one GEMV scheduling completes
and that
the local vector load remains vectorized.
Fixes #20118
Testing:
- `python -m pytest tests/python/s_tir/dlight/test_gpu_gemv.py -q` — 14
passed
- `python -m pytest tests/python/relax/test_pipeline.py -q` — 24 passed
- Exact issue reproducer — passed
- RTX 4070 Ti SUPER (SM89): FP32/FP16 widths 2, 3, 32, 33, and 128 —
NumPy parity
- `pre-commit run --all-files` — passed
---
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 1c451b964f..1a71fa697e 100644
--- a/python/tvm/s_tir/dlight/gpu/gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/gemv.py
@@ -164,13 +164,14 @@ class GEMV(GPUScheduleRule):
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 8cfcd1bb81..e7454f5505 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
test_gemv_cuda_target_without_max_shared_memory_per_block():
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)
