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