[PASS] More improvement of canonical (#314)

apache · Aug 12, 2017 · afa2086 · afa2086
1 parent 3c2569a
commit afa2086
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/src/arithmetic/canonical.cc b/src/arithmetic/canonical.cc
@@ -191,6 +191,9 @@ class Canonical::Internal : public IRMutator {
     ret_entry_.max_level = stack_.back().max_level;
     stack_.pop_back();
     CHECK(expr.defined());
+    if (const IntImm* op = expr.as<IntImm>()) {
+      return Mutate_(op, expr);
+    }
     return expr;
   }
   // call produce to get a cache entry.

diff --git a/tests/python/unittest/test_pass_simplify.py b/tests/python/unittest/test_pass_simplify.py
@@ -27,6 +27,16 @@ def test_basic():
     assert str(ret.value) == "(m - 1)"
 
 
+def test_canonical():
+    x = tvm.var("x")
+    z = tvm.const(3)
+    ret = tvm.ir_pass.CanonicalSimplify(x / (z*z) - x / (z*z))
+    assert(tvm.ir_pass.Equal(ret, 0))
+
+    ret = tvm.ir_pass.CanonicalSimplify(x / (z+z) - x / (z+z))
+    assert(tvm.ir_pass.Equal(ret, 0))
+
 if __name__ == "__main__":
     test_basic()
     test_simplify()
+    test_canonical()