[ARITH] CanonicalSimplifier, better folding, eliminate store. (#3464)

src/arithmetic/canonical_simplify.cc

@@ -256,7 +256,12 @@ class SumExprNode : public CanonicalExprNode {
         SplitExpr& rhs = args[j];
         if (!lhs->IndexEqual(rhs)) break;
         if (lhs->upper_factor < rhs->lower_factor) break;
-        if (lhs->lower_factor == rhs->upper_factor &&
+        if (lhs->upper_factor == rhs->upper_factor &&
+            lhs->lower_factor == rhs->lower_factor) {
+          // folding same co-efficient.
+          rhs.CopyOnWrite()->scale += lhs->scale;
+          lhs.CopyOnWrite()->scale = 0;
+        } else if (lhs->lower_factor == rhs->upper_factor &&
                    lhs->scale % rhs->scale == 0 &&
                    lhs->lower_factor == (lhs->scale / rhs->scale) * rhs->lower_factor) {
           // Rules used in the proof:

src/arithmetic/stmt_simplify.cc

@@ -106,6 +106,19 @@ class StmtSimplifier : public IRMutator {
     }
   }
 
+  // eliminate useless stores
+  Stmt Mutate_(const Store* op, const Stmt& s) {
+    Stmt stmt = IRMutator::Mutate_(op, s);
+    op = stmt.as<Store>();
+    if (const Load* load = op->value.as<Load>()) {
+      if (load->buffer_var.same_as(op->buffer_var) &&
+          Equal(load->index, op->index)) {
+        return Evaluate::make(0);
+      }
+    }
+    return stmt;
+  }
+
  protected:
   Analyzer analyzer_;
   // variable domain

tests/python/unittest/test_arith_canonical_simplify.py

@@ -187,6 +187,18 @@ def test_simplify_if_then_else():
     ck.verify(res, 0)
 
 
+def test_complex_cases():
+    ck = CanonicalChecker()
+    x = tvm.var("x")
+    y = tvm.var("y")
+    res2 = (((((((((((x*128) + y) % 1296)/36)*2) + 1)/2)*36) +
+              ((((((x*128) + y) % 36)*2) + 1)/2))
+             - (((x*128) + y) % 1296)) + 1)
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 5))
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 127))
+    ck.verify(res2, 1)
+
+
 if __name__ == "__main__":
     test_simplify_if_then_else()
     test_div_simplify()
@@ -195,3 +207,4 @@ if __name__ == "__main__":
     test_mul_sum_simplify()
     test_split_index_simplify()
     test_canonical_mixed()
+    test_complex_cases()