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 46c37fc523 [Fix][S-TIR][DLight] Guard non-affine reduction write-back
(#20057)
46c37fc523 is described below
commit 46c37fc523a2ae751272cc131376633cb502908b
Author: J.W (Jun) <[email protected]>
AuthorDate: Fri Aug 28 09:57:55 2026 +0400
[Fix][S-TIR][DLight] Guard non-affine reduction write-back (#20057)
Fixes #20048.
## Motivation
A valid Relax `conv2d` can fail during `relax.build(..., target="cuda")`
when the DLight GPU Reduction rule selects its inner-spatial schedule.
The rule fuses all spatial loops, but `_sch_inner_spatial` chooses its
`threadIdx.x` tile from only the innermost spatial extent. For the
reported shape, the non-unit spatial extents are `(2, 2, 20)` and the
selected tile extent is 10. After `rfactor` and `reverse_compute_at`,
the write-back block has to recover the original axes from the remaining
fused loop. The resulting bindings are not quasi-affine, so `bind`
rejects the block's compact-dataflow precondition.
The same geometry can occur without mixed spatial/reduction indexing,
for example in a plain reduction generated by `relax.op.sum`. Therefore
the fix is based on the fused spatial geometry, rather than on a
particular access pattern.
## Changes
This PR adds an applicability check inside `Reduction`. It counts the
non-unit spatial components that the write-back block must recover after
fusion and rejects the inner-spatial schedule when that recovery would
not remain quasi-affine. Unsupported reductions return `None`, allowing
`GeneralReduction` or `Fallback` to handle them; affine cases continue
to use the dedicated Reduction schedule.
The check derives spatial domains from the normalized dominant-read
access, which is the same order consumed by `_normalize`. Reordered or
transposed accesses therefore use the same loop order for both
applicability and scheduling. Thread-extent selection is shared with the
scheduling path, and symbolic extents are handled conservatively.
Unexpected `ScheduleError` exceptions are intentionally not caught by
`_apply_rules`. `None` remains the contract for a non-applicable rule,
while an exception from a rule remains visible as a scheduling bug
instead of being silently converted into a fallback.
## Testing
The changes were tested in a Python 3.12.13 environment with CUDA 13.0:
```bash
python -m pytest \
tests/python/s_tir/dlight/test_gpu_reduction.py \
tests/python/s_tir/dlight/test_gpu_fallback.py -q
```
Result: `25 passed`.
The regression tests cover the reported convolution-like access, a plain
reduction without mixed spatial/reduction indices, reordered
dominant-read access, an affine write-back shape, and propagation of an
unexpected `ScheduleError` from a broken rule.
The Relax reproduction from #20048 now builds successfully with
`target="cuda"`, and the previously failing plain `relax.op.sum` shape
is handled by a later schedule rule.
---
python/tvm/s_tir/dlight/gpu/reduction.py | 94 +++++++++++--
tests/python/s_tir/dlight/test_gpu_fallback.py | 29 ++++
tests/python/s_tir/dlight/test_gpu_reduction.py | 170 ++++++++++++++++++++++++
3 files changed, 279 insertions(+), 14 deletions(-)
diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py
b/python/tvm/s_tir/dlight/gpu/reduction.py
index ced5c97531..ff22cdfe98 100644
--- a/python/tvm/s_tir/dlight/gpu/reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/reduction.py
@@ -54,6 +54,65 @@ def _has_reduction_loop(block_info):
return any([info.kind == "R" for info in block_info.iters])
+def _suggest_inner_spatial_tx(s_factor: int | tirx.Expr) -> int:
+ """Pick the largest divisor no greater than 16."""
+ if not isinstance(s_factor, int):
+ return 1
+ len_tx = 16
+ while len_tx > 1 and s_factor % len_tx != 0:
+ len_tx -= 1
+ return len_tx
+
+
+def _get_spatial_domains_in_access_order(
+ block_info: SBlockInfo, access: arith.IterSumExpr
+) -> list[int | tirx.Expr] | None:
+ """Return normalized spatial extents in access order."""
+ iter_to_info = {info.var: info for info in block_info.iters}
+ spatial_domains = []
+ seen = set()
+ for split_expr in access.args:
+ var = split_expr.source.source
+ info = iter_to_info.get(var)
+ if info is None:
+ return None
+ if info.kind == "S":
+ if var in seen:
+ return None
+ seen.add(var)
+ # `_normalize` fuses the outer loop produced by this split.
+ if split_expr.lower_factor > 1:
+ extent = split_expr.extent
+ extent = int(extent) if isinstance(extent, tirx.IntImm) else
extent
+ spatial_domains.append(extent)
+ else:
+ spatial_domains.append(info.dom)
+
+ # `_normalize` appends omitted unit spatial loops.
+ for info in block_info.iters:
+ if info.kind == "S" and info.var not in seen:
+ if not isinstance(info.dom, int) or info.dom != 1:
+ return None
+ spatial_domains.append(info.dom)
+
+ if len(spatial_domains) != sum(info.kind == "S" for info in
block_info.iters):
+ return None
+ return spatial_domains
+
+
+def _inner_spatial_write_back_is_affine(spatial_domains: list[int |
tirx.Expr]) -> bool:
+ """Check whether the write-back bindings remain quasi-affine."""
+ if not spatial_domains:
+ return False
+ # Each non-unit outer dimension adds a div/mod component to the binding.
+ components = sum(1 for dom in spatial_domains[:-1] if not isinstance(dom,
int) or dom > 1)
+ innermost = spatial_domains[-1]
+ if not isinstance(innermost, int) or _suggest_inner_spatial_tx(innermost)
< innermost:
+ # A partial innermost tile adds one more component.
+ components += 1
+ return components < 3
+
+
class Reduction(GPUScheduleRule):
"""A rule for Reduction."""
@@ -92,13 +151,15 @@ class Reduction(GPUScheduleRule):
):
return None
# Step 2. Normalize the block, merge spatial and reduction iters
+ access = arith.normalize_to_iter_sum(
+ detect_dominant_read(block_stmt),
+ input_iters={i.var: i.dom for i in block_stmt.iter_vars},
+ )
+ spatial_domains = _get_spatial_domains_in_access_order(block_info,
access)
+ if spatial_domains is None:
+ return None
is_inner_reduction, c_factor, loop_order, s_split_index =
self._normalize(
- sch,
- block_info,
- arith.normalize_to_iter_sum(
- detect_dominant_read(block_stmt),
- input_iters={i.var: i.dom for i in block_stmt.iter_vars},
- ),
+ sch, block_info, access
)
if is_inner_reduction is None and c_factor is None:
return None
@@ -108,8 +169,17 @@ class Reduction(GPUScheduleRule):
sch, target, block, c_factor, epilogue, loop_order,
s_split_index
)
else:
+ if not _inner_spatial_write_back_is_affine(spatial_domains):
+ return None
self._sch_inner_spatial(
- sch, target, block, block_info, c_factor, epilogue,
loop_order, s_split_index
+ sch,
+ target,
+ block,
+ spatial_domains[-1],
+ c_factor,
+ epilogue,
+ loop_order,
+ s_split_index,
)
return sch
@@ -239,7 +309,7 @@ class Reduction(GPUScheduleRule):
sch: s_tir.Schedule,
_: Target,
block: s_tir.schedule.SBlockRV,
- block_info: SBlockInfo,
+ s_factor: int | tirx.Expr,
unroll_spatial_factor: int | None,
epilogue_info: SBlockInfo | None,
loop_order,
@@ -247,14 +317,10 @@ class Reduction(GPUScheduleRule):
):
# pylint: disable=invalid-name
s, r, _ = sch.get_loops(block)
- len_tx, len_ty = 16, 16
- s_factor = [i.dom for i in block_info.iters if i.kind == "S"][-1]
+ len_ty = 16
# get perfect spatial factor, spatial factor should be divide the
innermost spatial loop so
# that the block after r_factor and be reversed compute at the
original scope
- while len_tx > 1:
- if s_factor % len_tx == 0:
- break
- len_tx -= 1
+ len_tx = _suggest_inner_spatial_tx(s_factor)
_, _ = sch.split(s, factors=[None, len_tx])
_, ty = sch.split(r, factors=[None, len_ty])
# Schedule the RF block
diff --git a/tests/python/s_tir/dlight/test_gpu_fallback.py
b/tests/python/s_tir/dlight/test_gpu_fallback.py
index eb94734596..76bb99c8a0 100644
--- a/tests/python/s_tir/dlight/test_gpu_fallback.py
+++ b/tests/python/s_tir/dlight/test_gpu_fallback.py
@@ -16,7 +16,10 @@
# under the License.
# pylint: disable=missing-docstring
# ruff: noqa: E501, E741, F841
+import pytest
+
import tvm.testing
+from tvm import s_tir
from tvm.ir import assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm.script import ir as I
@@ -258,5 +261,31 @@ def test_gpu_fallback_ignores_non_gpu_functions():
assert_structural_equal(mod, After)
+def test_schedule_error_propagates_from_rule():
+ # ScheduleError indicates a broken rule and must propagate.
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((128,), "float32"), C: T.Buffer((128,),
"float32")):
+ for i in range(128):
+ with T.sblock("copy"):
+ vi = T.axis.remap("S", [i])
+ T.reads(A[vi])
+ T.writes(C[vi])
+ C[vi] = A[vi] * T.float32(2)
+
+ class BrokenRule(dl.base.ScheduleRule):
+ def apply(self, func, target, tunable):
+ sch = s_tir.Schedule(func)
+ sch.get_sblock("no_such_block")
+ return sch
+
+ with Target("nvidia/geforce-rtx-3090-ti"),
pytest.raises(s_tir.ScheduleError):
+ dl.ApplyDefaultSchedule( # pylint: disable=not-callable
+ BrokenRule(),
+ dl.gpu.Fallback(),
+ )(Before)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/s_tir/dlight/test_gpu_reduction.py
b/tests/python/s_tir/dlight/test_gpu_reduction.py
index ace05f93c3..c34dd6822a 100644
--- a/tests/python/s_tir/dlight/test_gpu_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_reduction.py
@@ -926,6 +926,176 @@ def test_reduction_inner_spatial_choose_perfect_factor():
assert_structural_equal(mod, Expected)
+def test_reduction_inner_spatial_non_affine_write_back_falls_back():
+ # The partial innermost tile leaves three div/mod components in the
write-back binding.
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((2, 4, 20), "float32"),
+ W: T.Buffer((3,), "float32"),
+ C: T.Buffer((2, 2, 20), "float32"),
+ ):
+ for n, y, x, k in T.grid(2, 2, 20, 3):
+ with T.sblock("conv"):
+ vn, vy, vx, vk = T.axis.remap("SSSR", [n, y, x, k])
+ T.reads(A[vn, vy + vk, vx], W[vk])
+ T.writes(C[vn, vy, vx])
+ with T.init():
+ C[vn, vy, vx] = T.float32(0)
+ C[vn, vy, vx] += A[vn, vy + vk, vx] * W[vk]
+
+ @I.ir_module(s_tir=True)
+ class Expected:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((2, 4, 20), "float32"),
+ W: T.Buffer((3,), "float32"),
+ C: T.Buffer((2, 2, 20), "float32"),
+ ):
+ T.func_attr({"tirx.is_scheduled": True})
+ for ax0_ax1_ax2_fused_0 in T.thread_binding(1,
thread="blockIdx.x"):
+ for ax0_ax1_ax2_fused_1 in T.thread_binding(1024,
thread="threadIdx.x"):
+ with T.sblock("conv_init"):
+ v0 = T.axis.spatial(
+ 2, (ax0_ax1_ax2_fused_0 * 1024 +
ax0_ax1_ax2_fused_1) // 40
+ )
+ v1 = T.axis.spatial(
+ 2, (ax0_ax1_ax2_fused_0 * 1024 +
ax0_ax1_ax2_fused_1) % 40 // 20
+ )
+ v2 = T.axis.spatial(
+ 20, (ax0_ax1_ax2_fused_0 * 1024 +
ax0_ax1_ax2_fused_1) % 20
+ )
+ T.where(ax0_ax1_ax2_fused_0 * 1024 +
ax0_ax1_ax2_fused_1 < 80)
+ T.reads()
+ T.writes(C[v0, v1, v2])
+ C[v0, v1, v2] = T.float32(0)
+ for ax3 in range(3):
+ with T.sblock("conv_update"):
+ v0 = T.axis.spatial(
+ 2, (ax0_ax1_ax2_fused_0 * 1024 +
ax0_ax1_ax2_fused_1) // 40
+ )
+ v1 = T.axis.spatial(
+ 2,
+ (ax0_ax1_ax2_fused_0 * 1024 +
ax0_ax1_ax2_fused_1) % 40 // 20,
+ )
+ v2 = T.axis.spatial(
+ 20, (ax0_ax1_ax2_fused_0 * 1024 +
ax0_ax1_ax2_fused_1) % 20
+ )
+ v3 = T.axis.reduce(3, ax3)
+ T.where(ax0_ax1_ax2_fused_0 * 1024 +
ax0_ax1_ax2_fused_1 < 80)
+ T.reads(C[v0, v1, v2], A[v0, v1 + v3, v2], W[v3])
+ T.writes(C[v0, v1, v2])
+ C[v0, v1, v2] += A[v0, v1 + v3, v2] * W[v3]
+
+ with Target("nvidia/geforce-rtx-3090-ti"):
+ # Reduction must decline instead of raising ScheduleError.
+ assert dl.gpu.Reduction().apply(Before["main"], Target.current(),
False) is None
+ mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
+ dl.gpu.Reduction(),
+ dl.gpu.Fallback(),
+ )(Before)
+ assert_structural_equal(mod, Expected)
+
+
+def test_reduction_inner_spatial_non_affine_without_mixed_access():
+ # The failure depends on the spatial extents, not mixed spatial/reduction
indexing.
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((2, 2, 3, 20), "float32"),
+ C: T.Buffer((2, 2, 20), "float32"),
+ ):
+ for n, y, x, k in T.grid(2, 2, 20, 3):
+ with T.sblock("sum"):
+ vn, vy, vx, vk = T.axis.remap("SSSR", [n, y, x, k])
+ T.reads(A[vn, vy, vk, vx])
+ T.writes(C[vn, vy, vx])
+ with T.init():
+ C[vn, vy, vx] = T.float32(0)
+ C[vn, vy, vx] += A[vn, vy, vk, vx]
+
+ with Target("nvidia/geforce-rtx-3090-ti"):
+ assert dl.gpu.Reduction().apply(Before["main"], Target.current(),
False) is None
+ mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
+ dl.gpu.Reduction(),
+ dl.gpu.GeneralReduction(),
+ dl.gpu.Fallback(),
+ )(Before)
+ assert "tirx.is_scheduled" in mod["main"].attrs
+
+
+def test_reduction_inner_spatial_reordered_access_declines():
+ # `_normalize` fuses the spatial loops in dominant-read order: n, x, y.
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((2, 16, 3, 20), "float32"),
+ C: T.Buffer((2, 20, 16), "float32"),
+ ):
+ for n, y, x, k in T.grid(2, 20, 16, 3):
+ with T.sblock("sum"):
+ vn, vy, vx, vk = T.axis.remap("SSSR", [n, y, x, k])
+ T.reads(A[vn, vx, vk, vy])
+ T.writes(C[vn, vy, vx])
+ with T.init():
+ C[vn, vy, vx] = T.float32(0)
+ C[vn, vy, vx] += A[vn, vx, vk, vy]
+
+ with Target("nvidia/geforce-rtx-3090-ti"):
+ assert dl.gpu.Reduction().apply(Before["main"], Target.current(),
False) is None
+
+
+def test_reduction_inner_spatial_affine_write_back_still_applies():
+ # An exact innermost tile leaves only two div/mod components.
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((2, 4, 16), "float32"),
+ W: T.Buffer((3,), "float32"),
+ C: T.Buffer((2, 2, 16), "float32"),
+ ):
+ for n, y, x, k in T.grid(2, 2, 16, 3):
+ with T.sblock("conv"):
+ vn, vy, vx, vk = T.axis.remap("SSSR", [n, y, x, k])
+ T.reads(A[vn, vy + vk, vx], W[vk])
+ T.writes(C[vn, vy, vx])
+ with T.init():
+ C[vn, vy, vx] = T.float32(0)
+ C[vn, vy, vx] += A[vn, vy + vk, vx] * W[vk]
+
+ with Target("nvidia/geforce-rtx-3090-ti"):
+ sch = dl.gpu.Reduction().apply(Before["main"], Target.current(), False)
+ assert sch is not None, "Reduction rule should still handle affine
write-back blocks"
+
+
+def test_reduction_inner_spatial_uses_normalized_split_extent():
+ # x // 2 normalizes the x extent from 8 to 4, which must also be used for
tx.
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((2, 3, 3, 4), "float32"),
+ C: T.Buffer((2, 3, 8), "float32"),
+ ):
+ for n, y, x, k in T.grid(2, 3, 8, 3):
+ with T.sblock("sum"):
+ vn, vy, vx, vk = T.axis.remap("SSSR", [n, y, x, k])
+ T.reads(A[vn, vy, vk, vx // 2])
+ T.writes(C[vn, vy, vx])
+ with T.init():
+ C[vn, vy, vx] = T.float32(0)
+ C[vn, vy, vx] += A[vn, vy, vk, vx // 2]
+
+ with Target("nvidia/geforce-rtx-3090-ti"):
+ sch = dl.gpu.Reduction().apply(Before["main"], Target.current(), False)
+ assert sch is not None, "Reduction should use the normalized split extent
for tx"
+ assert 'thread_binding(4, thread="threadIdx.x")' in sch.mod.script()
+
+
def test_repeat_transpose_gemv():
# fmt: off
