[arith][BugFix] Fix simplify input PrimExpr of DetectClipBound (#12150)

apache · Jul 22, 2022 · 8c42a83 · 8c42a83
1 parent 4fc79b5
commit 8c42a83
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/arith/detect_linear_equation.cc b/src/arith/detect_linear_equation.cc
@@ -245,15 +245,15 @@ void SplitCommExpr(const PrimExpr& e, std::vector<PrimExpr>* ret) {
 // e must be connected by and.
 Array<PrimExpr> DetectClipBound(const PrimExpr& e, const Array<Var>& vars) {
   std::vector<PrimExpr> splits;
-  SplitCommExpr<tir::AndNode>(e, &splits);
+  Analyzer analyzer;
+  SplitCommExpr<tir::AndNode>(analyzer.Simplify(e), &splits);
   std::unordered_map<const VarNode*, IntervalEntry> rmap;
   for (Var v : vars) {
     rmap[v.get()] = IntervalEntry();
   }
   for (PrimExpr cond : splits) {
     if (!DetectClipBound(cond, &rmap)) return Array<PrimExpr>();
   }
-  Analyzer analyzer;
   Array<PrimExpr> ret;
   for (Var v : vars) {
     IntervalEntry e = rmap[v.get()];

diff --git a/tests/python/unittest/test_arith_detect_clip_bound.py b/tests/python/unittest/test_arith_detect_clip_bound.py
@@ -31,6 +31,12 @@ def test_basic():
     m = tvm.arith.detect_clip_bound(tvm.tir.all(a + 10 * c <= 20, b - 1 > 0), [a, b])
     tvm.testing.assert_prim_expr_equal(m[1], 20 - 10 * c)
     tvm.testing.assert_prim_expr_equal(m[2], 2)
+    m = tvm.arith.detect_clip_bound(tvm.tir.all(tvm.tir.Not(a * 1 > b * 6), a - 1 > 0), [a])
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6)
+    m = tvm.arith.detect_clip_bound(tvm.tir.all(tvm.tir.Min(a, b) > 3, a - 10 < 0), [a, b])
+    tvm.testing.assert_prim_expr_equal(m[0], 4)
+    tvm.testing.assert_prim_expr_equal(m[1], 9)
+    tvm.testing.assert_prim_expr_equal(m[2], 4)
 
 
 if __name__ == "__main__":