[Fix][S-TIR][DLight] Fall back from non-affine reduction write-back

python/tvm/s_tir/dlight/base/transform.py

@@ -19,13 +19,17 @@
 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
 from tvm.target import Target
 
 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", "<unnamed>"),
+                exc_info=True,
+            )
+            continue
         if space is None:
             continue
         if isinstance(space, s_tir.Schedule):

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

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()

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