[PASS] More improvement of canonical (#314)

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.

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()