diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py b/python/tvm/s_tir/dlight/gpu/gemv.py index 1a71fa697e52..e54bd9d144f8 100644 --- a/python/tvm/s_tir/dlight/gpu/gemv.py +++ b/python/tvm/s_tir/dlight/gpu/gemv.py @@ -19,7 +19,7 @@ from functools import reduce -from tvm import s_tir, tirx +from tvm import arith, s_tir, tirx from tvm.target import Target from ..analysis import ( @@ -42,7 +42,7 @@ def apply( # pylint: disable=too-many-locals,too-many-branches,too-many-return- func: tirx.PrimFunc, target: Target, _: bool, - ) -> None | s_tir.Schedule | list[s_tir.Schedule]: + ) -> s_tir.Schedule | list[s_tir.Schedule] | None: if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target): return None sch = s_tir.Schedule(func) @@ -118,6 +118,8 @@ def apply( UNROLL, SUPPORT_WARP_SHUFFLE, ): + analyzer = arith.Analyzer() + # rfactor: reduce to tx * vec_c _, s, r, c = sch.get_loops(block=gemv) s = sch.fuse(_, s) @@ -226,7 +228,8 @@ def apply( factors=[None, get_max_factor(TILE_S, [1, 2, 4, 8])], preserve_unit_iters=True, ) - assert sch.get(ts_o).extent.value == 1 + if not analyzer.can_prove_equal(sch.get(ts_o).extent, 1): + return None ts = sch.fuse(ts_o, ts_i) sch.reorder(ts, tr, tile_s, vec_s, vec_c) sch.bind(ts, TAG_S) @@ -240,7 +243,8 @@ def apply( ts_o, ts_i, tile_s = sch.split( ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True ) - assert sch.get(ts_o).extent.value == 1 + if not analyzer.can_prove_equal(sch.get(ts_o).extent, 1): + return None ts = sch.fuse(ts_o, ts_i) sch.reorder(tile_s, ts, tr) sch.bind(ts, TAG_S) @@ -298,7 +302,8 @@ def apply( ts_o, ts_i, tile_s = sch.split( ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True ) - assert sch.get(ts_o).extent.value == 1 + if not analyzer.can_prove_equal(sch.get(ts_o).extent, 1): + return None ts = sch.fuse(ts_o, ts_i) sch.bind(ts, TAG_S) sch.set_scope(block, 0, "local") diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py b/tests/python/s_tir/dlight/test_gpu_gemv.py index e7454f5505dd..c7f3c112608d 100644 --- a/tests/python/s_tir/dlight/test_gpu_gemv.py +++ b/tests/python/s_tir/dlight/test_gpu_gemv.py @@ -1185,5 +1185,41 @@ def before( assert mod["main"].attrs["tirx.is_scheduled"] == 1 +def test_gemv_falls_back_for_non_unit_outer_spatial_tile(): + @T.prim_func(private=True, s_tir=True) + def before( + data: T.Buffer((1, 1, 3, 10), "float32"), + weight: T.Buffer((1, 1, 1, 2), "float32"), + output: T.Buffer((1, 1, 3, 9), "float32"), + ): + T.func_attr({"tirx.noalias": True}) + padded = T.sblock_alloc_buffer((1, 1, 3, 10), "float32") + for i0, i1, i2, i3 in T.grid(1, 1, 3, 10): + with T.sblock("pad"): + v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) + T.reads(data[v0, v1, v2, v3]) + T.writes(padded[v0, v1, v2, v3]) + padded[v0, v1, v2, v3] = data[v0, v1, v2, v3] + for nn, ff, yy, xx, rc, ry, rx in T.grid(1, 1, 3, 9, 1, 1, 2): + with T.sblock("conv2d"): + vnn, vff, vyy, vxx, vrc, vry, vrx = T.axis.remap( + "SSSSRRR", [nn, ff, yy, xx, rc, ry, rx] + ) + T.reads(padded[vnn, vrc, vyy + vry, vxx + vrx], weight[vff, vrc, vry, vrx]) + T.writes(output[vnn, vff, vyy, vxx]) + with T.init(): + output[vnn, vff, vyy, vxx] = T.float32(0) + output[vnn, vff, vyy, vxx] += ( + padded[vnn, vrc, vyy + vry, vxx + vrx] * weight[vff, vrc, vry, vrx] + ) + + mod = tvm.IRModule({"main": before}) + target = Target({"kind": "cuda", "max_num_threads": 1024, "max_shared_memory_per_block": 49152}) + with target: + mod = dl.ApplyDefaultSchedule(dl.gpu.GEMV(), dl.gpu.Fallback())(mod) + + assert mod["main"].attrs["tirx.is_scheduled"] == 1 + + if __name__ == "__main__": tvm.testing.main()