diff --git a/python/tvm/s_tir/dlight/gpu/fallback.py b/python/tvm/s_tir/dlight/gpu/fallback.py index e64f8544ecf1..ebb6998e1874 100644 --- a/python/tvm/s_tir/dlight/gpu/fallback.py +++ b/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) diff --git a/tests/python/s_tir/dlight/test_gpu_fallback.py b/tests/python/s_tir/dlight/test_gpu_fallback.py index 76bb99c8a016..74a6de92e2d3 100644 --- a/tests/python/s_tir/dlight/test_gpu_fallback.py +++ b/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: