[tvm] branch main updated: [BugFix][Arith] IterMapRewriter abort rewriting once failure (#15677)

[tqchen]

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

tqchen 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 1e6e2b35ae [BugFix][Arith] IterMapRewriter abort rewriting once 
failure (#15677)
1e6e2b35ae is described below

commit 1e6e2b35aec23146dc4232bb9887207aa70f52de
Author: Ruihang Lai <ruiha...@cs.cmu.edu>
AuthorDate: Wed Sep 6 07:33:18 2023 -0400

    [BugFix][Arith] IterMapRewriter abort rewriting once failure (#15677)
    
    This PR fixes an issue of the IterMapRewriter. Prior to this PR, the
    mutation function of the rewriter class returns the mutation results
    even when an invalid PrimExpr pattern was detected.
    
    Returning the mutation results when failure is not expected, and in
    such cases we should "abort" the mutation, and return the input
    PrimExpr which is not mutated, since insisting on returning the
    mutation results sometimes it incurs further error on other arith
    components like simplification.
    
    One unit test is added to ensure the rewriter behaves as expectation.
---
 src/arith/iter_affine_map.cc                        |  1 +
 tests/python/unittest/test_arith_iter_affine_map.py | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index 89a803d058..607be0a83d 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -322,6 +322,7 @@ class IterMapRewriter : public ExprMutator {
       ErrorLogger(this) << "IterMapExpr or subclasses should only result from 
calls in "
                         << "IterMapRewriter using DirectMutate.  "
                         << "Indirect return occurred in " << input_expr;
+      return input_expr;
     }
     return expr;
   }
diff --git a/tests/python/unittest/test_arith_iter_affine_map.py 
b/tests/python/unittest/test_arith_iter_affine_map.py
index 63bb79d2b2..1676855b31 100644
--- a/tests/python/unittest/test_arith_iter_affine_map.py
+++ b/tests/python/unittest/test_arith_iter_affine_map.py
@@ -1292,5 +1292,25 @@ def test_normalize_to_iter_sum():
     )
 
 
+def test_detect_iter_map_with_bufferload_recursion():
+    n = tvm.tir.Var("n", "int32")
+    m = tvm.tir.Var("m", "int32")
+    divisor = tvm.tir.Var("divisor", "int32")
+
+    i = tvm.tir.Var("i", "int32")
+    j = tvm.tir.Var("j", "int32")
+
+    buffer = tvm.tir.decl_buffer((n,), "int32", name="seqlen")
+
+    indices = [(buffer[i] + j) // divisor]
+    iter_vars = {
+        i: tvm.ir.Range(tvm.tir.const(0, "int32"), n),
+        j: tvm.ir.Range(tvm.tir.const(0, "int32"), m),
+    }
+
+    result = tvm.arith.detect_iter_map(indices, iter_vars)
+    assert len(result.indices) == 0
+
+
 if __name__ == "__main__":
     tvm.testing.main()