From e6a192f772313c7b76d6d11d0d2bc8998e8cb474 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] Guard 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. Detect the unsafe geometry before applying the inner-spatial schedule and return None so the next applicable schedule rule can handle the block. Derive the spatial order from the normalized dominant read so reordered accesses use the same order as Reduction itself. Keep unexpected ScheduleError failures visible at the rule-selection layer instead of catching them globally. Add regression coverage for the reported convolution-like access, a plain reduction without mixed spatial/reduction indices, reordered access, affine write-back, and propagation of unexpected ScheduleError. Fixes #20048 --- python/tvm/s_tir/dlight/gpu/reduction.py | 113 +++++++++++-- .../python/s_tir/dlight/test_gpu_fallback.py | 33 ++++ .../python/s_tir/dlight/test_gpu_reduction.py | 159 ++++++++++++++++++ 3 files changed, 291 insertions(+), 14 deletions(-) diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py b/python/tvm/s_tir/dlight/gpu/reduction.py index ced5c97531c6..a755a5f7e7fd 100644 --- a/python/tvm/s_tir/dlight/gpu/reduction.py +++ b/python/tvm/s_tir/dlight/gpu/reduction.py @@ -54,6 +54,82 @@ 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. + + Symbolic extents are handled conservatively because divisibility cannot be proven at + schedule time. + """ + 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 spatial iterator domains in the order used by ``_normalize``. + + The write-back loop is reconstructed from the dominant read's normalized access, rather + than from the block iterator declaration order. These orders differ for transposed or + otherwise reordered accesses, and using the declaration order can make the applicability + check disagree with the schedule that follows it. + """ + 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) + spatial_domains.append(info.dom) + + # ``_normalize`` appends spatial unit loops that do not occur in the dominant read. + # Any omitted non-unit loop would make the schedule inapplicable for a different reason. + 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: + """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. + """ + if not spatial_domains: + 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 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: + # 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.""" @@ -92,13 +168,15 @@ def apply( # pylint: disable=too-many-locals,too-many-branches,too-many-return- ): 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 +186,19 @@ 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_write_back_is_affine(spatial_domains): + # The write-back bindings would not be quasi-affine for this shape. Let + # GeneralReduction / Fallback handle it instead. + 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 +328,7 @@ def _sch_inner_spatial( 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 +336,10 @@ def _sch_inner_spatial( ): # 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 eb94734596a7..6dbf18a80313 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,35 @@ def cpu_func( assert_structural_equal(mod, After) +def test_schedule_error_propagates_from_rule(): + # An unexpected ScheduleError is a rule bug, rather than an applicability result. + # The reduction tests separately verify that Reduction declines known bad shapes by + # returning None before it schedules the block. + @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"), 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 ace05f93c387..7708bcedaf24 100644 --- a/tests/python/s_tir/dlight/test_gpu_reduction.py +++ b/tests/python/s_tir/dlight/test_gpu_reduction.py @@ -926,6 +926,165 @@ 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 + + +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. The + # applicability check must use this same order and decline the non-affine write-back. + @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