diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py b/python/tvm/s_tir/dlight/gpu/reduction.py index ced5c97531c6..afbedcb7993c 100644 --- a/python/tvm/s_tir/dlight/gpu/reduction.py +++ b/python/tvm/s_tir/dlight/gpu/reduction.py @@ -107,10 +107,16 @@ def apply( # pylint: disable=too-many-locals,too-many-branches,too-many-return- self._sch_inner_reduction( sch, target, block, c_factor, epilogue, loop_order, s_split_index ) - else: + elif ( self._sch_inner_spatial( sch, target, block, block_info, c_factor, epilogue, loop_order, s_split_index ) + is None + ): + # `_sch_inner_spatial` bails (returns None) on nest shapes it does not model; + # `sch` may already carry partial, unbound edits at that point, so it must not + # be handed back as if scheduling succeeded. Let `Fallback` take over instead. + return None return sch def _normalize( # pylint: disable=too-many-branches @@ -249,6 +255,7 @@ def _sch_inner_spatial( 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] + # 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: @@ -268,10 +275,18 @@ def _sch_inner_spatial( sch.decompose_reduction(rf, r) # Schedule the write back block sch.reverse_compute_at(block, bx, preserve_unit_loops=True) - _, r, *s = sch.get_loops(block) + _, r, *s_loops = sch.get_loops(block) + + if len([i for i in block_info.iters if i.kind == "S"]) > 2: + return None + if unroll_spatial_factor: - assert len(s) == len(loop_order) - new_order_s = [s[loop_order[i]] for i in range(len(s))] + if len(s_loops) != len(loop_order): + # `loop_order` is indexed by the original block's spatial loops; when + # `reverse_compute_at` regenerates a different number of them the + # premutation is meaningless. Bail rather than mis-permute. + return None + new_order_s = [s_loops[loop_order[i]] for i in range(len(s_loops))] sch.reorder(*new_order_s) new_order_s[s_split_index], c = sch.split( new_order_s[s_split_index], factors=[None, unroll_spatial_factor] @@ -280,7 +295,7 @@ def _sch_inner_spatial( s = sch.fuse(*new_order_s) sch.reorder(s, c, r) else: - s = sch.fuse(*s) + s = sch.fuse(*s_loops) sch.reorder(s, r) sch.bind(s, "threadIdx.x") sch.bind(r, "threadIdx.y") @@ -303,3 +318,4 @@ def _sch_inner_spatial( tx, _ = sch.split(sch.fuse(*s), factors=[len_tx, None]) sch.bind(tx, "threadIdx.x") # pylint: enable=invalid-name + return sch diff --git a/tests/python/s_tir/dlight/test_gpu_reduction.py b/tests/python/s_tir/dlight/test_gpu_reduction.py index ace05f93c387..7566f6c02ba3 100644 --- a/tests/python/s_tir/dlight/test_gpu_reduction.py +++ b/tests/python/s_tir/dlight/test_gpu_reduction.py @@ -1179,5 +1179,172 @@ def matmul(lv43: T.Buffer((T.int64(1), T.int64(32), T.int64(1)), "float16"), lv4 assert_structural_equal(mod, Before) +def test_reduction_4d_non_contiguous_reduce_axis0(): + # Test reduction with kind-based loop classification in _sch_inner_spatial. + # With 3+ spatial dims, `blockIdx.x` alone fully absorbs 2+ of them, so the + # write-back block ends up with 2+ degenerate (extent-1) spatial loops; fusing + # and binding those would require a non-canonical floordiv/subtraction binding. + # `_sch_inner_spatial` detects this and bails, so `Fallback` schedules it instead. + # fmt: off + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((4, 8, 16, 32), "float32"), A_red: T.Buffer((8, 16, 32), "float32")): + T.func_attr({"tirx.noalias": True}) + for ax0, ax1, ax2, k in T.grid(8, 16, 32, 4): + with T.sblock("A_red"): + v_ax0, v_ax1, v_ax2, v_k = T.axis.remap("SSSR", [ax0, ax1, ax2, k]) + T.reads(A[v_k, v_ax0, v_ax1, v_ax2]) + T.writes(A_red[v_ax0, v_ax1, v_ax2]) + with T.init(): + A_red[v_ax0, v_ax1, v_ax2] = T.float32(0) + A_red[v_ax0, v_ax1, v_ax2] = A_red[v_ax0, v_ax1, v_ax2] + A[v_k, v_ax0, v_ax1, v_ax2] + # fmt: on + + target = Target("nvidia/geforce-rtx-3090-ti") + with target: + mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable + dl.gpu.Reduction(), dl.gpu.Fallback() + )(Before) + assert mod["main"].attrs["tirx.is_scheduled"] == 1 + + +def test_reduction_4d_non_contiguous_reduce_axis1(): + # Test reduction with kind-based loop classification on 4D tensor, reduce axis 1. + # Same 3+-spatial-dims shape as axis0 above: `_sch_inner_spatial` bails on the + # 2+ degenerate write-back loops and `Fallback` schedules it instead. + # fmt: off + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((4, 8, 16, 32), "float32"), A_red: T.Buffer((4, 16, 32), "float32")): + T.func_attr({"tirx.noalias": True}) + for ax0, ax1, ax2, k in T.grid(4, 16, 32, 8): + with T.sblock("A_red"): + v_ax0, v_ax1, v_ax2, v_k = T.axis.remap("SSSR", [ax0, ax1, ax2, k]) + T.reads(A[v_ax0, v_k, v_ax1, v_ax2]) + T.writes(A_red[v_ax0, v_ax1, v_ax2]) + with T.init(): + A_red[v_ax0, v_ax1, v_ax2] = T.float32(0) + A_red[v_ax0, v_ax1, v_ax2] = A_red[v_ax0, v_ax1, v_ax2] + A[v_ax0, v_k, v_ax1, v_ax2] + # fmt: on + + target = Target("nvidia/geforce-rtx-3090-ti") + with target: + mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable + dl.gpu.Reduction(), dl.gpu.Fallback() + )(Before) + assert mod["main"].attrs["tirx.is_scheduled"] == 1 + + +def test_reduction_4d_non_contiguous_reduce_axis1_keepdims(): + # Test reduction with keepdims=True on non-contiguous axis. + # Unit dimension inserted at reduce position is collapsed by NormalizePrimFunc. + # leaving the same 3+-spatial-dims shape as axis0/axis1 above: `_sch_inner_spatial` + # bails on the 2+ degenerate write-back loops and `Fallback` schedules it instead. + # fmt: off + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((4, 8, 16, 32), "float32"), A_red: T.Buffer((4, 1, 16, 32), "float32")): + T.func_attr({"tirx.noalias": True}) + for ax0, ax1u, ax1, ax2, k in T.grid(4, 1, 16, 32, 8): + with T.sblock("A_red"): + v_ax0, v_ax1u, v_ax1, v_ax2, v_k = T.axis.remap("SSSSR", [ax0, ax1u, ax1, ax2, k]) + T.reads(A[v_ax0, v_k, v_ax1, v_ax2]) + T.writes(A_red[v_ax0, v_ax1u, v_ax1, v_ax2]) + with T.init(): + A_red[v_ax0, v_ax1u, v_ax1, v_ax2] = T.float32(0) + A_red[v_ax0, v_ax1u, v_ax1, v_ax2] = A_red[v_ax0, v_ax1u, v_ax1, v_ax2] + A[v_ax0, v_k, v_ax1, v_ax2] + # fmt: on + + target = Target("nvidia/geforce-rtx-3090-ti") + with target: + mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable + dl.gpu.Reduction(), dl.gpu.Fallback() + )(Before) + assert mod["main"].attrs["tirx.is_scheduled"] == 1 + + +def test_reduction_4d_non_contiguous_reduce_axis2(): + # Test reduction with kind-based loop classification on 4D tensor, reduce axis 2. + # Same 3+-spatial-dims shape as the other 4D tests above: `_sch_inner_spatial` bails + # on the 2+ degenerate write-back loops and `Fallback` schedules it instead. + # fmt: off + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((4, 8, 16, 32), "float32"), A_red: T.Buffer((4, 8, 32), "float32")): + T.func_attr({"tirx.noalias": True}) + for ax0, ax1, ax2, k in T.grid(4, 8, 32, 16): + with T.sblock("A_red"): + v_ax0, v_ax1, v_ax2, v_k = T.axis.remap("SSSR", [ax0, ax1, ax2, k]) + T.reads(A[v_ax0, v_ax1, v_k, v_ax2]) + T.writes(A_red[v_ax0, v_ax1, v_ax2]) + with T.init(): + A_red[v_ax0, v_ax1, v_ax2] = T.float32(0) + A_red[v_ax0, v_ax1, v_ax2] = A_red[v_ax0, v_ax1, v_ax2] + A[v_ax0, v_ax1, v_k, v_ax2] + # fmt: on + + target = Target("nvidia/geforce-rtx-3090-ti") + with target: + mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable + dl.gpu.Reduction(), dl.gpu.Fallback() + )(Before) + assert mod["main"].attrs["tirx.is_scheduled"] == 1 + + +def test_reduction_4d_contiguous_reduce_axis3_regression(): + # Regression guard: contiguous (last-axis) reduction on 4D tensor. + # Takes _sch_inner_reduction path (not affected by _sch_inner_spatial changes); + # must keep passing to ensure no unintended side effects. + # fmt: off + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((4, 8, 16, 32), "float32"), A_red: T.Buffer((4, 8, 16), "float32")): + T.func_attr({"tirx.noalias": True}) + for ax0, ax1, ax2, k in T.grid(4, 8, 16, 32): + with T.sblock("A_red"): + v_ax0, v_ax1, v_ax2, v_k = T.axis.remap("SSSR", [ax0, ax1, ax2, k]) + T.reads(A[v_ax0, v_ax1, v_ax2, v_k]) + T.writes(A_red[v_ax0, v_ax1, v_ax2]) + with T.init(): + A_red[v_ax0, v_ax1, v_ax2] = T.float32(0) + A_red[v_ax0, v_ax1, v_ax2] = A_red[v_ax0, v_ax1, v_ax2] + A[v_ax0, v_ax1, v_ax2, v_k] + # fmt: on + + target = Target("nvidia/geforce-rtx-3090-ti") + with target: + mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable + assert mod["main"].attrs["tirx.is_scheduled"] == 1 + + +def test_reduction_3d_non_contiguous_reduce_axis0_regression(): + # Regression guard: non-contiguous reduction on 3D tensor with 2 spatial dimensions. + # Takes _sch_inner_spatial path as 4D tests but with fewer spatial dimensions. + # ensures existing working cases remain unaffected by kind-based classification. + # fmt: off + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((4, 8, 32), "float32"), A_red: T.Buffer((8, 32), "float32")): + T.func_attr({"tirx.noalias": True}) + for ax0, ax1, k in T.grid(8, 32, 4): + with T.sblock("A_red"): + v_ax0, v_ax1, v_k = T.axis.remap("SSR", [ax0, ax1, k]) + T.reads(A[v_k, v_ax0, v_ax1]) + T.writes(A_red[v_ax0, v_ax1]) + with T.init(): + A_red[v_ax0, v_ax1] = T.float32(0) + A_red[v_ax0, v_ax1] = A_red[v_ax0, v_ax1] + A[v_k, v_ax0, v_ax1] + # fmt: on + + target = Target("nvidia/geforce-rtx-3090-ti") + with target: + mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable + assert mod["main"].attrs["tirx.is_scheduled"] == 1 + + if __name__ == "__main__": tvm.testing.main()