From 688c22e0cf378ed5e79d4060572139aebbf5c062 Mon Sep 17 00:00:00 2001 From: Junius-Wynn Date: Mon, 27 Jul 2026 23:13:24 +0800 Subject: [PATCH] [Fix][S-TIR][DLight] Fall back from non-affine reduction write-back Issue #20048 reports that relax.build can abort with ScheduleError while compiling a valid CUDA conv2d. The GPU Reduction rule fuses all spatial loops, but its inner-spatial schedule derives the thread extent from the original innermost spatial extent. For some output shapes, rfactor and reverse_compute_at then produce non-quasi-affine write-back bindings, and bind aborts the default schedule chain. Guard the known unsafe write-back geometry and let GeneralReduction or Fallback handle those blocks, while preserving the dedicated Reduction schedule for affine cases. Also let default rule selection continue after ScheduleError so an applicability mistake in one rule does not become a hard compilation failure. Add regression coverage for the reported convolution-like access, a plain sum without mixed spatial/reduction indices, affine write-back shapes, and transform-level fallback. The tests directly verify that Reduction declines unsafe shapes and document the remaining dominant-read spatial-ordering edge case with a strict xfail. Fixes #20048 --- python/tvm/s_tir/dlight/base/transform.py | 17 +- python/tvm/s_tir/dlight/gpu/reduction.py | 52 +++++- .../python/s_tir/dlight/test_gpu_fallback.py | 33 ++++ .../python/s_tir/dlight/test_gpu_reduction.py | 166 ++++++++++++++++++ 4 files changed, 262 insertions(+), 6 deletions(-) diff --git a/python/tvm/s_tir/dlight/base/transform.py b/python/tvm/s_tir/dlight/base/transform.py index 18014e40c07d..d2034044947e 100644 --- a/python/tvm/s_tir/dlight/base/transform.py +++ b/python/tvm/s_tir/dlight/base/transform.py @@ -19,6 +19,8 @@ or a space for MetaSchedule tuning """ +import logging + from tvm import s_tir, tirx from tvm.ir import IRModule from tvm.ir.transform import PassContext, module_pass @@ -26,6 +28,8 @@ from .schedule_rule import ScheduleRule +logger = logging.getLogger(__name__) + def _is_scheduled(func: tirx.PrimFunc) -> bool: if not isinstance(func, tirx.PrimFunc): @@ -85,7 +89,18 @@ def _apply_rules( tunable: bool, ) -> list[s_tir.Schedule] | None: for rule in rules: - space = rule.apply(func, target, tunable) + try: + space = rule.apply(func, target, tunable) + except s_tir.schedule.ScheduleError: + # A rule that mis-judges its own applicability must not abort the whole + # pass: fall through so the remaining rules still get a chance. + logger.debug( + "Schedule rule %s failed on %s; trying the next rule.", + type(rule).__name__, + func.attrs.get("global_symbol", ""), + exc_info=True, + ) + continue if space is None: continue if isinstance(space, s_tir.Schedule): diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py b/python/tvm/s_tir/dlight/gpu/reduction.py index ced5c97531c6..e22fed23e60c 100644 --- a/python/tvm/s_tir/dlight/gpu/reduction.py +++ b/python/tvm/s_tir/dlight/gpu/reduction.py @@ -54,6 +54,44 @@ 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 thread extent up to 16 that divides the innermost spatial extent. + + A symbolic extent never compares equal to zero, so it falls through to 1, matching + the behaviour this loop had when it was inlined in `_sch_inner_spatial`. + """ + len_tx = 16 + while len_tx > 1 and s_factor % len_tx != 0: + len_tx -= 1 + return len_tx + + +def _inner_spatial_write_back_is_affine(block_info: SBlockInfo) -> bool: + """Whether `_sch_inner_spatial` can keep the write-back block bindings quasi-affine. + + `_sch_inner_spatial` fuses every spatial loop and then splits `len_tx` off the inner + end, where `len_tx` is derived from the innermost spatial extent alone. After + `rfactor` + `reverse_compute_at`, the write-back block has to recover each original + spatial index from the remaining fused outer loop. Each non-unit spatial extent left + outside the inner tile contributes one floordiv/floormod component to that recovery, + and once three or more components pile up the resulting bindings are no longer + recognized as quasi-affine. `bind` then fails the compact-dataflow precondition + (local reduction block condition #2) and aborts the whole default schedule chain. + """ + s_doms = [info.dom for info in block_info.iters if info.kind == "S"] + if not s_doms: + return False + # A symbolic extent may be 1 at runtime, but it has to be assumed non-unit here. + components = sum(1 for dom in s_doms[:-1] if not isinstance(dom, int) or dom > 1) + innermost = s_doms[-1] + if not isinstance(innermost, int) or _suggest_inner_spatial_tx(innermost) < innermost: + # The inner tile only covers part of the innermost extent, so its leftover + # outer part stays in the fused loop as one more component. A symbolic extent + # cannot be proven to be tiled exactly, so treat it the same way. + components += 1 + return components < 3 + + class Reduction(GPUScheduleRule): """A rule for Reduction.""" @@ -82,6 +120,9 @@ def apply( # pylint: disable=too-many-locals,too-many-branches,too-many-return- block_info = block_infos[0] block = block_info.block_rv block_stmt = sch.get(block) + # `_normalize` mutates `sch`, so decide up-front whether the inner-spatial + # schedule is even applicable to this block shape. + inner_spatial_ok = _inner_spatial_write_back_is_affine(block_info) # Step 1. Check reduction block if ( @@ -108,6 +149,10 @@ def apply( # pylint: disable=too-many-locals,too-many-branches,too-many-return- sch, target, block, c_factor, epilogue, loop_order, s_split_index ) else: + if not inner_spatial_ok: + # Let GeneralReduction / Fallback handle it instead of raising a + # ScheduleError, which would abort the entire schedule chain. + return None self._sch_inner_spatial( sch, target, block, block_info, c_factor, epilogue, loop_order, s_split_index ) @@ -247,14 +292,11 @@ def _sch_inner_spatial( ): # pylint: disable=invalid-name s, r, _ = sch.get_loops(block) - len_tx, len_ty = 16, 16 + len_ty = 16 s_factor = [i.dom for i in block_info.iters if i.kind == "S"][-1] # 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 eb94734596a7..fda8ee225cca 100644 --- a/tests/python/s_tir/dlight/test_gpu_fallback.py +++ b/tests/python/s_tir/dlight/test_gpu_fallback.py @@ -17,6 +17,7 @@ # pylint: disable=missing-docstring # ruff: noqa: E501, E741, F841 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 +259,37 @@ def cpu_func( assert_structural_equal(mod, After) +def test_schedule_error_falls_through_to_next_rule(): + # Keep this synthetic failure focused on the transform-layer contract. The + # reduction tests separately verify that Reduction declines the known bad shapes; + # this test verifies that an unexpected ScheduleError from any rule still lets the + # remaining rules schedule the original function. + @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) + # Looking up a block that does not exist raises a ScheduleError, standing in + # for a rule that misjudges the shape of the function it is given. + sch.get_sblock("no_such_block") + return sch + + with Target("nvidia/geforce-rtx-3090-ti"): + mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable + BrokenRule(), + dl.gpu.Fallback(), + )(Before) + assert mod["main"].attrs.get("tirx.is_scheduled") == 1 + + 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 ace05f93c387..e64d492f0ed3 100644 --- a/tests/python/s_tir/dlight/test_gpu_reduction.py +++ b/tests/python/s_tir/dlight/test_gpu_reduction.py @@ -17,6 +17,8 @@ # pylint: disable=missing-docstring,line-too-long,invalid-name,too-few-public-methods,too-many-locals # ruff: noqa: E501, F841 +import pytest + import tvm.testing from tvm.ir import assert_structural_equal from tvm.s_tir import dlight as dl @@ -926,6 +928,170 @@ def main(var_A: T.handle, var_B: T.handle, matmul: T.Buffer((T.int64(1), T.int64 assert_structural_equal(mod, Expected) +def test_reduction_inner_spatial_non_affine_write_back_falls_back(): + # `_sch_inner_spatial` derives its thread extent from the innermost spatial extent + # alone (20 -> len_tx=10) but fuses all three spatial loops. Recovering the original + # indices in the write-back block then needs three floordiv/floormod components, + # which is no longer quasi-affine, so `bind` would raise a ScheduleError. The rule + # must decline instead and let Fallback schedule the block. + @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"): + # Check the rule contract directly. The default-schedule pass also catches + # ScheduleError from a rule, so checking only the pass result would not prove + # that Reduction itself declined this shape before attempting `bind`. + 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(): + # Same failure without any spatial/reduction mixing in the read indices: the trigger + # is the spatial extents, not the access pattern. + @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"): + # This must be rejected by Reduction for the same geometric reason, even + # though none of the read indices mixes spatial and reduction variables. + assert dl.gpu.Reduction().apply(Before["main"], Target.current(), False) is None + # Must not raise, and must still end up scheduled by a later rule. + 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 + + +@pytest.mark.xfail( + strict=True, + reason="The affine guard uses block-iterator order instead of normalized access order", +) +def test_reduction_inner_spatial_reordered_access_declines(): + # The block iterators are ordered n, y, x, k, but the dominant read is ordered + # n, x, k, y. `_normalize` therefore fuses the spatial loops as n, x, y, while + # the current guard incorrectly treats x=16 as the innermost spatial extent. + # Reduction must decline this shape instead of reaching a non-affine `bind`. + @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(): + # len_tx == innermost spatial extent (16), so the write-back bindings stay affine and + # the dedicated Reduction rule must still claim the block. + @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_repeat_transpose_gemv(): # fmt: off