[TVM] Fix segfault for CanonicalSimplify(x % -1) (#2194)

src/arithmetic/canonical.cc

@@ -515,7 +515,15 @@ class Canonical::Internal : public IRMutator {
         n->elem.push_back(e);
       }
       Expr ret = Sum2Expr(ComExpr(n), v.type()) % v;
-      return Binary(ret.as<Mod>(), ret);
+      if (const Mod* mod = ret.as<Mod>()) {
+        return Binary(mod, ret);
+      } else {
+        // Sometimes the result is a constant, this may happen when value is -1
+        CHECK(is_const(ret)) << "CanonicalSimplify: "
+          << Sum2Expr(ComExpr(n), v.type()) << " % " << v << " is " << ret
+          << " which is neither Mod, nor a constant";
+        return ret;
+      }
     }
     ret_entry_.sum = pair[1];
     ret_entry_.max_level = stack_.back().max_level;

tests/python/unittest/test_arith_simplify.py

@@ -20,6 +20,16 @@ def test_simplify():
     zz = zz.a
     assert zz.a == x and zz.b.value == 4
 
+    n = tvm.var('n')
+    assert tvm.ir_pass.Equal(tvm.ir_pass.CanonicalSimplify(n % (-1)), tvm.const(0))
+    assert tvm.ir_pass.Equal(tvm.ir_pass.CanonicalSimplify(n % 1), tvm.const(0))
+    assert tvm.ir_pass.Equal(tvm.ir_pass.CanonicalSimplify(n / 1), n)
+    tvm.ir_pass.CanonicalSimplify(n / (-1))
+    # This is not true in the current implementation
+    #  assert tvm.ir_pass.Equal(tvm.ir_pass.CanonicalSimplify(n / (-1)),
+    #                           tvm.ir_pass.CanonicalSimplify(-n))
+
+
 def test_simplify_mod():
     """Not yet working, mock design"""
     ib = tvm.ir_builder.create()