[SimplifyExpr] Simplify consecutive adds with constants (#9671)

Author: Ivy Zhang

src/relay/transforms/simplify_expr.cc

@@ -585,6 +585,49 @@ class EliminateIdentityRewrite : public DFPatternRewrite {
   DFPattern const_;
 };
 
+/*! \brief Make two consecutive add able to be constant_folded.
+ * This pattern matching supports commutative property for addition.
+ */
+class SimplifyConsecutiveAdd : public DFPatternRewrite {
+ public:
+  SimplifyConsecutiveAdd() {
+    x_ = IsWildcard();
+    const1_ = IsConstant();
+    const2_ = IsConstant();
+    DFPattern add_op = IsOp("add");
+    pattern_ = add_op({add_op({x_, const1_}), const2_});
+  }
+
+  Expr Callback(const Expr& pre, const Expr& post,
+                const Map<DFPattern, Array<Expr>>& node_map) const override {
+    const CallNode* call = pre.as<CallNode>();
+    auto x = node_map[x_][0];
+    auto c1 = node_map[const1_][0];
+    auto c2 = node_map[const2_][0];
+
+    auto pre_call = call;
+    // Find the next add call.
+    if (pre_call->args[1].as<ConstantNode>()) {
+      pre_call = pre_call->args[0].as<CallNode>();
+    } else {
+      pre_call = pre_call->args[1].as<CallNode>();
+    }
+    // Do nothing if both inputs are not constants as they will be constant folded already.
+    if (pre_call->args[0].as<ConstantNode>() && pre_call->args[1].as<ConstantNode>()) {
+      return post;
+    } else {
+      auto add_res = Call(call->op, {c1, c2});
+      return Call(call->op, {x, add_res});
+    }
+    return post;
+  }
+
+ private:
+  DFPattern x_;
+  DFPattern const1_;
+  DFPattern const2_;
+};
+
 Expr SimplifyExpr(const Expr& expr, const IRModule& mod) {
   // the rewrites will be applied in the given order, and repeated until fixed point
   DFPatternRewriteComposer composer;
@@ -599,6 +642,7 @@ Expr SimplifyExpr(const Expr& expr, const IRModule& mod) {
   composer.AddRewrite<SimplifyTranspose>();
   composer.AddRewrite<SimplifyCast>();
   composer.AddRewrite<FullElementwise>();
+  composer.AddRewrite<SimplifyConsecutiveAdd>();
   return RewritePatterns(composer.MakeCallbacks(), expr, mod);
 }
 

tests/python/relay/test_pass_simplify_expr.py

@@ -512,5 +512,47 @@ def test_concretize_multiple():
     assert tvm.ir.structural_equal(actual, expected)
 
 
+def test_simplify_consecutive_add():
+    shape = (32, 1, 1)
+    c_data = np.empty(shape).astype("float32")
+    c1 = relay.const(c_data)
+    c2 = relay.const(c_data)
+
+    def before_const_right():
+        x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32")
+        w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32")
+        y = relay.nn.conv2d(x, w, padding=(1, 1))
+        y = relay.add(y, c1)
+        y = relay.add(y, c2)
+        y = relay.nn.relu(y)
+        return relay.Function([x, w], y)
+
+    def before_const_left():
+        x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32")
+        w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32")
+        y = relay.nn.conv2d(x, w, padding=(1, 1))
+        y = relay.add(c1, y)
+        y = relay.add(c2, y)
+        y = relay.nn.relu(y)
+        return relay.Function([x, w], y)
+
+    def expected():
+        x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32")
+        w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32")
+        y = relay.nn.conv2d(x, w, padding=(1, 1))
+        c3 = relay.add(c1, c2)
+        y = relay.add(y, c3)
+        y = relay.nn.relu(y)
+        return relay.Function([x, w], y)
+
+    zr = before_const_right()
+    zl = before_const_left()
+    zzr = run_opt_pass(zr, transform.SimplifyExpr())
+    zzl = run_opt_pass(zl, transform.SimplifyExpr())
+    after = run_opt_pass(expected(), transform.InferType())
+    assert tvm.ir.structural_equal(zzr, after)
+    assert tvm.ir.structural_equal(zzl, after)
+
+
 if __name__ == "__main__":
     pytest.main([__file__])