[DLight][GPU] Skip fallback scheduling for zero extents

python/tvm/s_tir/dlight/gpu/fallback.py

@@ -53,6 +53,19 @@ def visit_for(node: tirx.For):
     return found
 
 
+def _has_zero_extent_loop(stmt: tirx.Stmt) -> bool:
+    """Check whether a statement contains a statically empty loop."""
+    found = False
+
+    def visit_for(node: tirx.For):
+        nonlocal found
+        if isinstance(node.extent, tirx.IntImm) and node.extent.value == 0:
+            found = True
+
+    tvm_ffi.structural_walk(stmt, (tirx.For, visit_for), order="post")
+    return found
+
+
 class Fallback(GPUScheduleRule):
     """
     A fallback schedule rule for all GPU operators. It will try to inline all the blocks first,
@@ -67,6 +80,8 @@ def apply(  # pylint: disable=too-many-locals,missing-docstring
     ) -> s_tir.Schedule:
         if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
             return None
+        if _has_zero_extent_loop(func.body):
+            return None
         max_threads_per_block = base.max_threads_per_block(target)
 
         sch = s_tir.Schedule(func)

tests/python/s_tir/dlight/test_gpu_fallback.py

@@ -69,6 +69,23 @@ def main(
     assert_structural_equal(mod, After)
 
 
+def test_fallback_skips_zero_extent_spatial():
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((4, 0), "float32"), B: T.Buffer((4, 0), "float32")):
+            for i, j in T.grid(4, 0):
+                with T.sblock("copy"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vi, vj]
+
+    with Target("nvidia/geforce-rtx-3090-ti"):
+        mod = dl.ApplyDefaultSchedule(  # pylint: disable=not-callable
+            dl.gpu.Fallback(),
+        )(Module)
+    assert_structural_equal(mod, Module)
+
+
 def test_fallback_reduction():
     @I.ir_module(s_tir=True)
     class Module: