[DLight][CUDA] Fall back from unsupported GEMV spatial tiles

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

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