[Relax] Fix divide-by-zero in reshape pattern detection

src/relax/analysis/tir_op_pattern_kind.cc

@@ -28,6 +28,8 @@
 #include <tvm/tirx/op.h>
 #include <tvm/tirx/stmt_functor.h>
 
+#include <algorithm>
+
 namespace tvm {
 namespace relax {
 
@@ -477,7 +479,12 @@ bool HasReshapePattern(const PrimFunc& func) {
         nontrivial_buffer = dst_buffer_;
       }
 
-      if (nontrivial_indices.defined()) {
+      // Skip check 1 on zero-extent iters: the inverse index map would divide by zero.
+      bool has_zero_extent = std::any_of(
+          block->iter_vars.begin(), block->iter_vars.end(),
+          [this](const IterVar& v) { return this->ana_->CanProveEqual(v->dom->extent, 0); });
+
+      if (nontrivial_indices.defined() && !has_zero_extent) {
         PrimType dtype =
             !block->iter_vars.empty() ? block->iter_vars[0]->var.ty() : PrimType::Int(64);
         tirx::Var fused_var("fused", dtype);

tests/python/relax/test_analysis.py

@@ -591,6 +591,22 @@ def reshape_scheduled(
     assert has_reshape_pattern(reshape_scheduled)
 
 
+def test_reshape_pattern_zero_extent():
+    @T.prim_func(s_tir=True)
+    def transpose_zero(
+        rxplaceholder: T.Buffer((3, 0, 4), "float32"),
+        T_transpose: T.Buffer((0, 3, 4), "float32"),
+    ):
+        for i0, i1, i2 in T.grid(0, 3, 4):
+            with T.sblock("T_transpose"):
+                ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+                T.reads(rxplaceholder[ax1, ax0, ax2])
+                T.writes(T_transpose[ax0, ax1, ax2])
+                T_transpose[ax0, ax1, ax2] = rxplaceholder[ax1, ax0, ax2]
+
+    assert not has_reshape_pattern(transpose_zero)
+
+
 def test_reshape_pattern_expand_dims():
     @T.prim_func(s_tir=True)
     def expand_dims(