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