[Bugfix] Fix other div zero errors also in rewrite_simplify (#8983)

* fix div zero error in rewrite_simplify

* update the style to fix ci error

* remove useless code and comment

* fix div zero error of mod, floordiv, floormod in rewrite_simplify

* rewrite the test case of divison by zero to fix ci error

* remove useless tab

* retrigger ci

* remove useless blank to retrigger ci

Author: Sen Yang

src/arith/rewrite_simplify.cc

@@ -645,6 +645,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const ModNode* op) {
     if (truncmod(ramp(b1, c1, lanes), broadcast(c2, lanes)).Match(ret)) {
       int64_t c1val = c1.Eval()->value;
       int64_t c2val = c2.Eval()->value;
+      ICHECK(c2val != 0) << "division by zero";
       if (c1val % c2val == 0) {
         return broadcast(truncmod(b1, c2), lanes).Eval();
       }
@@ -724,6 +725,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorDivNode* op) {
     if (floordiv(ramp(b1, c1, lanes), broadcast(c2, lanes)).Match(ret)) {
       int64_t c1val = c1.Eval()->value;
       int64_t c2val = c2.Eval()->value;
+      ICHECK(c2val != 0) << "division by zero";
       if (c1val % c2val == 0) {
         return ramp(floordiv(b1, c2), floordiv(c1, c2), lanes).Eval();
       }
@@ -852,6 +854,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorModNode* op) {
     if (floormod(ramp(b1, c1, lanes), broadcast(c2, lanes)).Match(ret)) {
       int64_t c1val = c1.Eval()->value;
       int64_t c2val = c2.Eval()->value;
+      ICHECK(c2val != 0) << "division by zero";
       if (c1val % c2val == 0) {
         return broadcast(floormod(b1, c2), lanes).Eval();
       }

tests/python/unittest/test_arith_rewrite_simplify.py

@@ -934,9 +934,23 @@ def test_shift_left_simplify():
 
 def test_div_zero_simplify():
     ck = RewriteChecker()
+    ramp = tvm.tir.Ramp(1, 1, 2)
+    broadcast = tvm.tir.Broadcast(0, 2)
 
     with pytest.raises(tvm.error.TVMError) as cm:
-        ck.analyzer.rewrite_simplify(tvm.tir.Div(tvm.tir.Ramp(1, 1, 2), tvm.tir.Broadcast(0, 2)))
+        ck.analyzer.rewrite_simplify(tvm.tir.Div(ramp, broadcast))
+        assert "division by zero" in str(cm.execption)
+
+    with pytest.raises(tvm.error.TVMError) as cm:
+        ck.analyzer.rewrite_simplify(tvm.tir.Mod(ramp, broadcast))
+        assert "division by zero" in str(cm.execption)
+
+    with pytest.raises(tvm.error.TVMError) as cm:
+        ck.analyzer.rewrite_simplify(tvm.tir.FloorDiv(ramp, broadcast))
+        assert "division by zero" in str(cm.execption)
+
+    with pytest.raises(tvm.error.TVMError) as cm:
+        ck.analyzer.rewrite_simplify(tvm.tir.FloorMod(ramp, broadcast))
         assert "division by zero" in str(cm.execption)