This is an automated email from the ASF dual-hosted git repository.
junrushao 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 1534983 [ARITH] fix zero iter bug in arith (#8494)
1534983 is described below
commit 1534983d39a05225a8721cccbab7de966473e7ff
Author: Bohan Hou <[email protected]>
AuthorDate: Mon Jul 19 05:38:36 2021 +0800
[ARITH] fix zero iter bug in arith (#8494)
* fix
* black
---
src/arith/iter_affine_map.cc | 1 -
tests/python/unittest/test_arith_iter_affine_map.py | 9 +++++++++
2 files changed, 9 insertions(+), 1 deletion(-)
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index 96cb928..cd48227 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -515,7 +515,6 @@ class IterMapRewriter : public ExprMutator {
*/
Optional<IterSplitExpr> TryFuseIters(IterSumExpr expr) {
if (!is_zero(expr->base)) return NullOpt;
- if (expr->args.size() == 1) return expr->args[0];
// select the iterators in order
std::vector<bool> visited(expr->args.size(), false);
std::vector<IterSplitExpr> flattened_iters, grouped_iters;
diff --git a/tests/python/unittest/test_arith_iter_affine_map.py
b/tests/python/unittest/test_arith_iter_affine_map.py
index b34acb9..c307034 100644
--- a/tests/python/unittest/test_arith_iter_affine_map.py
+++ b/tests/python/unittest/test_arith_iter_affine_map.py
@@ -285,6 +285,15 @@ def test_predicate():
)
assert len(res) == 0
+ # zero iter
+ xo = tvm.tir.Var("xo", "int32"), 1
+ xi = tvm.tir.Var("xi", "int32"), 129
+ y = tvm.tir.Var("y", "int32"), 128
+
+ res = tvm.arith.detect_iter_map(
+ [xo[0] * 129 + xi[0], y[0]], var_dom([xo, xi, y]), xo[0] * 129 + xi[0]
< 128
+ )
+
def convert_division(divisions):
if divisions is None or len(divisions) == 0:
