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 850807b904 [Fix][DLight] Reject GEMV accesses unsupported by
scheduling (#20122)
850807b904 is described below
commit 850807b9042563afa1101eb08af04362dfe5930e
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Mon Aug 17 13:35:39 2026 -0700
[Fix][DLight] Reject GEMV accesses unsupported by scheduling (#20122)
DLight's GEMV schedule rules can recognize reductions whose buffer
accesses do not satisfy the assumptions of subsequent scheduling steps.
For generic GEMV, normalization may produce a composite access source
that does not correspond to a block iterator. LowBatchGEMV can similarly
accept non-point access regions that are unsupported by
`Schedule.pad_einsum`. These cases currently raise an exception while
applying the schedule. This PR:
1. Validates that each normalized GEMV access source maps to a block
iterator.
2. Requires LowBatchGEMV read and write regions to be unit-extent point
accesses before applying `pad_einsum`.
3. Adds tests for composite transposed-convolution accesses and
non-einsum LowBatchGEMV accesses.
---
python/tvm/s_tir/dlight/analysis/gemv.py | 4 ++++
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 12 ++++++++++
tests/python/s_tir/dlight/test_gpu_gemv.py | 27 ++++++++++++++++++++++
.../python/s_tir/dlight/test_gpu_low_batch_gemv.py | 25 ++++++++++++++++++++
4 files changed, 68 insertions(+)
diff --git a/python/tvm/s_tir/dlight/analysis/gemv.py
b/python/tvm/s_tir/dlight/analysis/gemv.py
index f1134a0637..11e280e642 100644
--- a/python/tvm/s_tir/dlight/analysis/gemv.py
+++ b/python/tvm/s_tir/dlight/analysis/gemv.py
@@ -126,6 +126,10 @@ def normalize(
):
return None
iter_to_info = {i.var: i for i in block_info.iters}
+ if not access.args or any(
+ split_expr.source.source not in iter_to_info for split_expr in
access.args
+ ):
+ return None
batch_loops, s_loops, r_loops, c_loops = [], [], [], []
inner_axis = access.args[-1].source.source
is_inner_reduction = iter_to_info[inner_axis].kind == "R"
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 8dedbec4b6..b64235224d 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -53,6 +53,17 @@ def _get_reduction_expr(block: tirx.SBlock) -> tirx.Expr |
None:
return buffer_store.value.b
+def _has_pad_einsum_compatible_access(block: tirx.SBlock) -> bool:
+ """Check the point-access restriction required by
``Schedule.pad_einsum``."""
+ return all(
+ isinstance(dim.extent, tirx.IntImm)
+ and int(dim.extent) == 1
+ and isinstance(dim.min, tirx.IntImm | tirx.Var)
+ for region in [*block.reads, *block.writes]
+ for dim in region.region
+ )
+
+
def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) -> list[tirx.Buffer]
| None:
"""Check if the block is a low batch GEMM.
@@ -79,6 +90,7 @@ def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) ->
list[tirx.Buffer] |
conditions.append(len(block_stmt.reads) >= 2)
conditions.append(len(block_stmt.writes) == 1)
conditions.append(_get_reduction_expr(block_stmt) is not None)
+ conditions.append(_has_pad_einsum_compatible_access(block_stmt))
conditions.append(
len(collect_block_iter_vars_used_in_access_region(block_stmt,
block_stmt.writes[0].region))
> 0
diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py
b/tests/python/s_tir/dlight/test_gpu_gemv.py
index 14adab22ea..8cfcd1bb81 100644
--- a/tests/python/s_tir/dlight/test_gpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_gemv.py
@@ -24,6 +24,33 @@ from tvm.script import tirx as T
from tvm.target import Target
+def test_gemv_rejects_composite_normalized_axis():
+ @T.prim_func(private=True, s_tir=True)
+ def before(
+ p_data: T.handle,
+ weight: T.Buffer((64, 1, 512), "float32"),
+ p_output: T.handle,
+ n: T.int64,
+ ):
+ data = T.match_buffer(p_data, (1, 64, n), "float32")
+ output = T.match_buffer(p_output, (1, 1, n * 256), "float32")
+ for w, rc, rw in T.grid(n * 256, 64, 512):
+ with T.sblock("conv1d_transpose"):
+ vw, vrc, vrw = T.axis.remap("SRR", [w, rc, rw])
+ T.reads(data[0, vrc, (vw + vrw - 383) // 256], weight[vrc, 0,
511 - vrw])
+ T.writes(output[0, 0, vw])
+ with T.init():
+ output[0, 0, vw] = T.float32(0)
+ output[0, 0, vw] += (
+ data[0, vrc, (vw + vrw - 383) // 256] * weight[vrc, 0, 511
- vrw]
+ )
+
+ mod = tvm.IRModule({"main": before})
+ with Target("webgpu"):
+ scheduled = dl.ApplyDefaultSchedule(dl.gpu.GEMV())(mod)
+ tvm.ir.assert_structural_equal(scheduled["main"], before)
+
+
def test_gemv_basic():
# fmt: off
@T.prim_func(private=True, s_tir=True)
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 6f75191e50..528f45ca5e 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
@@ -559,6 +559,31 @@ 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_rejects_non_einsum_buffer_access():
+ @T.prim_func(private=True, s_tir=True)
+ def before(
+ var_A: T.handle,
+ var_B: T.handle,
+ var_C: T.handle,
+ ):
+ batch_size = T.int64()
+ A = T.match_buffer(var_A, (batch_size, 8), "float16")
+ B = T.match_buffer(var_B, (4, batch_size + 8), "float16")
+ C = T.match_buffer(var_C, (batch_size, 4), "float16")
+ for i, j, k in T.grid(batch_size, 4, 8):
+ with T.sblock("attention_score"):
+ vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+ T.reads(A[vi, vk], B[vj, T.max(vi + vk - 7, 0)])
+ T.writes(C[vi, vj])
+ with T.init():
+ C[vi, vj] = T.float16(0)
+ C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, T.max(vi + vk - 7,
0)]
+
+ with Target("webgpu") as target:
+ result = dl.gpu.LowBatchGEMV(4).apply(before, target, False)
+ assert result is None
+
+
def test_low_batch_gemv_broadcast_epilogue():
# fmt: off
@T.prim_func(private=True, s_tir=True)
