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-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new d7347e6cc [FIX] Drop cancelled iterator-map splits before fusion (#191)
d7347e6cc is described below
commit d7347e6cc39453e2c4a115068400abed9e1f8101
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Nov 4 15:54:28 2026 -0500
[FIX] Drop cancelled iterator-map splits before fusion (#191)
Drop iterator splits whose combined scale is zero before fusion,
avoiding a divide-by-zero for predicates such as `i - i >= 0`.
Constraints with no remaining iterator terms keep the existing
normalization failure: `detect_iter_map` returns an unsuccessful result,
while `iter_map_simplify` retains its safe fallback. This change does
not add proof rules for empty constraints.
Addresses apache/tvm#20843. Related: apache/tvm-tirx#127. The zero-scale
removal follows ShawnNotFound's change in that PR.
---------
Co-authored-by: ShawnNotFound <[email protected]>
---
src/sym/iter_affine_map.cc | 7 ++++++-
tests/python/sym/test_iter_affine_map.py | 11 +++++++++++
2 files changed, 17 insertions(+), 1 deletion(-)
diff --git a/src/sym/iter_affine_map.cc b/src/sym/iter_affine_map.cc
index 61b5a21dd..8068407f5 100644
--- a/src/sym/iter_affine_map.cc
+++ b/src/sym/iter_affine_map.cc
@@ -1246,7 +1246,12 @@ class IterMapRewriter : public tvm::ExprMutator {
} else {
rhs.CopyOnWrite()->scale = lvalue->scale - rhs->scale;
}
- lhs->args.Set(i, rhs);
+ // Iterator fusion divides by the scale, so discard cancelled terms.
+ if (IsZero(rhs->scale)) {
+ lhs->args.erase(lhs->args.begin() + i);
+ } else {
+ lhs->args.Set(i, rhs);
+ }
return;
}
}
diff --git a/tests/python/sym/test_iter_affine_map.py
b/tests/python/sym/test_iter_affine_map.py
index ca0764b80..ebbf241c6 100644
--- a/tests/python/sym/test_iter_affine_map.py
+++ b/tests/python/sym/test_iter_affine_map.py
@@ -1353,6 +1353,17 @@ def
test_iter_map_simplify_predicate_fallback_requires_no_padding():
tvm_tirx.ir.assert_structural_equal(simplified, [index])
+def test_cancelled_iter_map_terms():
+ i = tvm_tirx.tirx.Var("i", "int32")
+ dom_map = var_dom([(i, 4)])
+ predicate = i - i >= 0
+ result = tvm_tirx.sym.detect_iter_map([i], dom_map, predicate=predicate)
+ assert len(result.indices) == 0
+ assert len(result.errors) != 0
+ simplified = tvm_tirx.sym.iter_map_simplify([i], dom_map,
predicate=predicate)
+ tvm_tirx.ir.assert_structural_equal(simplified, [i])
+
+
def test_iter_map_simplify_symbolic_reshape():
n = tvm_tirx.tirx.Var("n", "int64")
fused = tvm_tirx.tirx.Var("fused", "int64")
