(tvm) branch main updated: [Relax] Fix divide-by-zero in reshape pattern detection (#19958)

[tlopex]

This is an automated email from the ASF dual-hosted git repository.

tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 82fa8bf8f6 [Relax] Fix divide-by-zero in reshape pattern detection 
(#19958)
82fa8bf8f6 is described below

commit 82fa8bf8f688e27f60d7d52c0d4e0da3c9ebab9b
Author: Guan-Ming Chiu <105915352+guan404m...@users.noreply.github.com>
AuthorDate: Fri Jul 10 02:39:55 2026 +0800

    [Relax] Fix divide-by-zero in reshape pattern detection (#19958)
    
    ## Why
    
    Fixes #17745. `has_reshape_pattern` builds an inverse index map that
    divides by each iter extent, so a zero-extent iter crashed with
    divide-by-zero.
    
    ## How
    
    - Skip the fused-var check when any block iter has zero extent; such
    blocks touch no elements, so they are not reshape patterns.
    - Added `test_reshape_pattern_zero_extent` in
    `tests/python/relax/test_analysis.py`.
---
 src/relax/analysis/tir_op_pattern_kind.cc |  9 ++++++++-
 tests/python/relax/test_analysis.py       | 16 ++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/relax/analysis/tir_op_pattern_kind.cc 
b/src/relax/analysis/tir_op_pattern_kind.cc
index b43969e866..89c40c1e4e 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/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);
diff --git a/tests/python/relax/test_analysis.py 
b/tests/python/relax/test_analysis.py
index 7be99442ac..52f60ac08e 100644
--- a/tests/python/relax/test_analysis.py
+++ b/tests/python/relax/test_analysis.py
@@ -591,6 +591,22 @@ def test_reshape_pattern_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(