[GitHub] [tvm] mbrookhart commented on a change in pull request #7731: [Relay][Pass] ConcretizeLike and EliminateIdentity rewrites for SimplifyExpr

[GitBox <gi...@apache.org>]

mbrookhart commented on a change in pull request #7731:
URL: https://github.com/apache/tvm/pull/7731#discussion_r600836110



##########
File path: src/relay/transforms/simplify_expr.cc
##########
@@ -249,36 +249,214 @@ class FullElementwise : public SimplifyPattern {
 };
 
 /*!
- * \brief ExprSimplifier simplifies the Relay expression.
+ * \brief Converts `*_like` operators to their explicit shape equivalent (e.g. `zeros_like(x, y)` to
+ * `zeros(x, y.shape)`), when the target shape is concrete. This removes unnecessary dependencies
+ * and can enable more opportunities for operator fusion.
  */
-class ExprSimplifier {
+class ConcretizeLikeRewrite : public DFPatternRewrite {
  public:
-  explicit ExprSimplifier(IRModule mod) : mod_(mod) {
-    CreateCallback(SimplifyReshape());
-    CreateCallback(SimplifyTranspose());
-    CreateCallback(FullElementwise());
+  explicit ConcretizeLikeRewrite(const Op& op) {
+    ICHECK(op->num_inputs == 1 || op->num_inputs == 2)
+        << "ConcretizeLike does not handle operators that aren't unary or binary, got: " << op;
+    like_pat_ = IsWildcard();
+    data_pat_ = IsWildcard();
+    if (op->num_inputs == 1) {
+      pattern_ = IsExpr(op)({like_pat_});
+    } else {
+      pattern_ = IsExpr(op)({data_pat_, like_pat_});
+    }
   }
-  template <typename T>
-  void CreateCallback(const T& pattern) {
-    auto func = [pattern](TVMArgs args, TVMRetValue* rv) {
-      Expr pre = args[0];
-      Expr post = args[1];
-      Map<DFPattern, Array<Expr>> node_map = args[2];
-      *rv = pattern.callback(pre, post, node_map);
-    };
-    callbacks_.push_back(DFPatternCallback(pattern.pattern(), PackedFunc(func), true));
+
+  virtual bool Check(const Expr& pre, const Expr& post,
+                     const Map<DFPattern, Array<Expr>>& node_map) const {
+    const CallNode* call_node = pre.as<CallNode>();
+    ICHECK(call_node);
+
+    if (!call_node->checked_type().as<TensorTypeNode>()) {
+      return false;
+    }
+
+    return true;
+  }
+
+  virtual Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                          DataType dtype) const = 0;
+
+  Expr Callback(const Expr& pre, const Expr& post,
+                const Map<DFPattern, Array<Expr>>& node_map) const override {
+    if (!Check(pre, post, node_map)) {
+      return post;
+    }
+
+    const TensorTypeNode* like_ty = pre->checked_type().as<TensorTypeNode>();
+    Array<Integer> cshape;
+    for (const auto& dim : like_ty->shape) {
+      if (const auto* imm = dim.as<IntImmNode>()) {
+        cshape.push_back(Integer(GetRef<IntImm>(imm)));
+      } else {
+        // shape is not static, don't concretize
+        return post;
+      }
+    }
+
+    return Concretize(node_map, cshape, like_ty->dtype);
+  }
+
+ protected:
+  DFPattern data_pat_;
+  DFPattern like_pat_;
+};
+
+class ConcretizeZerosLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeZerosLikeRewrite() : ConcretizeLikeRewrite(Op::Get("zeros_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    return MakeZeros(shape, dtype);
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeZerosLikeRewrite);
+};
+
+class ConcretizeOnesLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeOnesLikeRewrite() : ConcretizeLikeRewrite(Op::Get("ones_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    return MakeOnes(shape, dtype);
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeOnesLikeRewrite);
+};
+
+class ConcretizeReshapeLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeReshapeLikeRewrite() : ConcretizeLikeRewrite(Op::Get("reshape_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    return MakeReshape(node_map[data_pat_][0], shape);
   }
 
-  Expr Simplify(const Expr& expr) { return RewritePatterns(callbacks_, expr, mod_); }
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeReshapeLikeRewrite);
+};
+
+class ConcretizeCollapseSumLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeCollapseSumLikeRewrite() : ConcretizeLikeRewrite(Op::Get("collapse_sum_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    ICHECK_LE(shape.size(), std::numeric_limits<int64_t>::max());
+    static const Op& op = Op::Get("collapse_sum_to");
+    auto attrs = make_object<InitOpAttrs>();
+    attrs->shape = shape;
+    auto cshape =
+        MakeConstantTensor(DataType::Int(32), {static_cast<int64_t>(shape.size())}, shape);
+    return Call(op, {node_map[data_pat_][0], cshape}, Attrs(attrs));
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeCollapseSumLikeRewrite);
+};
+
+class ConcretizeBroadcastToLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeBroadcastToLikeRewrite() : ConcretizeLikeRewrite(Op::Get("broadcast_to_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    return MakeBroadCastTo(node_map[data_pat_][0], shape);
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeBroadcastToLikeRewrite);
+};
+
+/*! \brief Eliminates expressions that are equivalent to identity. */
+class EliminateIdentityRewrite : public DFPatternRewrite {
+ public:
+  EliminateIdentityRewrite() {
+    x_ = IsWildcard();
+    const_ = IsConstant();
+
+    DFPattern add_op = IsOp("add");
+    DFPattern mul_op = IsOp("multiply");
+    DFPattern zeros_expr = IsOp("zeros")({}) || IsOp("zeros_like")({IsWildcard()}) || const_;
+    DFPattern ones_expr = IsOp("ones")({}) || IsOp("ones_like")({IsWildcard()}) || const_;
+
+    DFPattern add_id = add_op({x_, zeros_expr}) || add_op({zeros_expr, x_});
+    DFPattern mul_id = mul_op({x_, ones_expr}) || mul_op({ones_expr, x_});

Review comment:
       Can we also simplify `x * 0 = 0` and `x + 0 = x`?
   
   The pattern matcher should be able to match these ops irrespective of order:
   https://github.com/apache/tvm/blob/cfe2e288a331b10e72e10c7e465df375b44e6ae9/src/relay/ir/dataflow_matcher.cc#L275-L281
   
   You can probably get away without the AltPattern here?

##########
File path: src/relay/transforms/simplify_expr.cc
##########
@@ -249,36 +248,214 @@ class FullElementwise : public SimplifyPattern {
 };
 
 /*!
- * \brief ExprSimplifier simplifies the Relay expression.
+ * \brief Converts `*_like` operators to their explicit shape equivalent (e.g. `zeros_like(x, y)` to
+ * `zeros(x, y.shape)`), when the target shape is concrete. This removes unnecessary dependencies
+ * and can enable more opportunities for operator fusion.
  */
-class ExprSimplifier {
+class ConcretizeLikeRewrite : public DFPatternRewrite {
  public:
-  explicit ExprSimplifier(IRModule mod) : mod_(mod) {
-    CreateCallback(SimplifyReshape());
-    CreateCallback(SimplifyTranspose());
-    CreateCallback(FullElementwise());
+  explicit ConcretizeLikeRewrite(const Op& op) {
+    ICHECK(op->num_inputs == 1 || op->num_inputs == 2)
+        << "ConcretizeLike does not handle operators that aren't unary or binary, got: " << op;
+    like_pat_ = IsWildcard();
+    data_pat_ = IsWildcard();
+    if (op->num_inputs == 1) {
+      pattern_ = IsExpr(op)({like_pat_});
+    } else {
+      pattern_ = IsExpr(op)({data_pat_, like_pat_});
+    }
   }
-  template <typename T>
-  void CreateCallback(const T& pattern) {
-    auto func = [pattern](TVMArgs args, TVMRetValue* rv) {
-      Expr pre = args[0];
-      Expr post = args[1];
-      Map<DFPattern, Array<Expr>> node_map = args[2];
-      *rv = pattern.callback(pre, post, node_map);
-    };
-    callbacks_.push_back(DFPatternCallback(pattern.pattern(), PackedFunc(func), true));
+
+  virtual bool Check(const Expr& pre, const Expr& post,
+                     const Map<DFPattern, Array<Expr>>& node_map) const {
+    const CallNode* call_node = pre.as<CallNode>();
+    ICHECK(call_node);
+
+    if (!call_node->checked_type().as<TensorTypeNode>()) {
+      return false;
+    }
+
+    return true;
+  }
+
+  virtual Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                          DataType dtype) const = 0;
+
+  Expr Callback(const Expr& pre, const Expr& post,
+                const Map<DFPattern, Array<Expr>>& node_map) const override {
+    if (!Check(pre, post, node_map)) {
+      return post;
+    }
+
+    const TensorTypeNode* like_ty = pre->checked_type().as<TensorTypeNode>();
+    Array<Integer> cshape;
+    for (const auto& dim : like_ty->shape) {
+      if (const auto* imm = dim.as<IntImmNode>()) {
+        cshape.push_back(Integer(GetRef<IntImm>(imm)));
+      } else {
+        // shape is not static, don't concretize
+        return post;
+      }
+    }
+
+    return Concretize(node_map, cshape, like_ty->dtype);
+  }
+
+ protected:
+  DFPattern data_pat_;
+  DFPattern like_pat_;
+};
+
+class ConcretizeZerosLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeZerosLikeRewrite() : ConcretizeLikeRewrite(Op::Get("zeros_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    return MakeZeros(shape, dtype);
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeZerosLikeRewrite);
+};
+
+class ConcretizeOnesLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeOnesLikeRewrite() : ConcretizeLikeRewrite(Op::Get("ones_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    return MakeOnes(shape, dtype);
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeOnesLikeRewrite);
+};
+
+class ConcretizeReshapeLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeReshapeLikeRewrite() : ConcretizeLikeRewrite(Op::Get("reshape_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    return MakeReshape(node_map[data_pat_][0], shape);
   }
 
-  Expr Simplify(const Expr& expr) { return RewritePatterns(callbacks_, expr, mod_); }
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeReshapeLikeRewrite);
+};
+
+class ConcretizeCollapseSumLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeCollapseSumLikeRewrite() : ConcretizeLikeRewrite(Op::Get("collapse_sum_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    ICHECK_LE(shape.size(), std::numeric_limits<int64_t>::max());
+    static const Op& op = Op::Get("collapse_sum_to");
+    auto attrs = make_object<InitOpAttrs>();
+    attrs->shape = shape;
+    auto cshape =
+        MakeConstantTensor(DataType::Int(32), {static_cast<int64_t>(shape.size())}, shape);
+    return Call(op, {node_map[data_pat_][0], cshape}, Attrs(attrs));
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeCollapseSumLikeRewrite);
+};
+
+class ConcretizeBroadcastToLikeRewrite : public ConcretizeLikeRewrite {
+ public:
+  ConcretizeBroadcastToLikeRewrite() : ConcretizeLikeRewrite(Op::Get("broadcast_to_like")) {}
+
+  Expr Concretize(const Map<DFPattern, Array<Expr>>& node_map, Array<Integer> shape,
+                  DataType dtype) const override {
+    return MakeBroadCastTo(node_map[data_pat_][0], shape);
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(ConcretizeBroadcastToLikeRewrite);
+};
+
+/*! \brief Eliminates expressions that are equivalent to identity. */
+class EliminateIdentityRewrite : public DFPatternRewrite {
+ public:
+  EliminateIdentityRewrite() {
+    x_ = IsWildcard();
+    const_ = IsConstant();
+
+    DFPattern add_op = IsOp("add");
+    DFPattern mul_op = IsOp("multiply");
+    DFPattern zeros_expr = IsOp("zeros")({}) || IsOp("zeros_like")({IsWildcard()}) || const_;
+    DFPattern ones_expr = IsOp("ones")({}) || IsOp("ones_like")({IsWildcard()}) || const_;
+
+    DFPattern add_id = add_op({x_, zeros_expr}) || add_op({zeros_expr, x_});
+    DFPattern mul_id = mul_op({x_, ones_expr}) || mul_op({ones_expr, x_});
+    DFPattern sub_id = IsOp("subtract")({x_, zeros_expr});
+    DFPattern div_id = IsOp("divide")({x_, ones_expr});
+
+    pattern_ = add_id || mul_id || sub_id || div_id;
+  }
+
+  bool CheckConstant(const OpNode* op, const ConstantNode* constant) const {
+    if (!IsScalar(GetRef<Expr>(constant))) {
+      return false;
+    }
+    long double value = ToScalar(constant->data);
+    if (op->name == "add" || op->name == "subtract") {
+      return value == 0.0;
+    } else if (op->name == "multiply" || op->name == "divide") {
+      return value == 1.0;
+    }
+    return false;
+  }
+
+  Expr Callback(const Expr& pre, const Expr& post,
+                const Map<DFPattern, Array<Expr>>& node_map) const override {
+    const CallNode* call = pre.as<CallNode>();
+    ICHECK(call);
+    Type pre_type = pre->checked_type_;
+    ICHECK(pre_type.as<TensorTypeNode>());
+    auto x = node_map[x_][0];
+    bool is_left = post.as<CallNode>()->args[1] == x;
+    Type x_type;
+    if (is_left) {
+      x_type = call->args[1]->checked_type_;
+    } else {
+      x_type = call->args[0]->checked_type_;
+    }
+
+    if (node_map.count(const_)) {
+      // the other argument is a Constant in this case
+      const ConstantNode* constant = node_map[const_][0].as<ConstantNode>();
+      const OpNode* op = call->op.as<OpNode>();
+      ICHECK(constant);
+      ICHECK(op);
+      if (!CheckConstant(op, constant)) {
+        return post;
+      }
+    }
+
+    if (StructuralEqual()(x_type, pre_type)) {
+      return x;
+    }
+
+    return post;
+  }
+
+  TVM_DF_PATTERN_REWRITE_GETTER(EliminateIdentityRewrite);
 
  private:
-  IRModule mod_;
-  /*! \brief Callbacks for expr simplification */
-  Array<DFPatternCallback> callbacks_;
+  DFPattern x_;
+  DFPattern const_;
 };
 
 Expr SimplifyExpr(const Expr& expr, const IRModule& mod) {
-  return ExprSimplifier(mod).Simplify(expr);
+  static Array<DFPatternCallback> callbacks = {ConcretizeZerosLikeRewrite::GetCallback(),

Review comment:
       Correct, they will do as much as they can to the graph in this order, and then loop back and try again. That being said, that's the same thing you'd get if you ran them as separate passes, you'd rewrite everything you could before the next pass, but the next pass might open opportunities to do more with the current pass if you ran it again.

##########
File path: src/relay/transforms/simplify_expr.h
##########
@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file src/relay/transforms/simplify_expr.h
+ * \brief Utility data structures for simplifying Relay expressions.
+ */
+#ifndef TVM_RELAY_TRANSFORMS_SIMPLIFY_EXPR_H_
+#define TVM_RELAY_TRANSFORMS_SIMPLIFY_EXPR_H_
+
+#include <tvm/relay/dataflow_matcher.h>
+#include <tvm/relay/expr.h>
+
+#include <vector>
+
+namespace tvm {
+namespace relay {
+
+/*! \brief Defines a static function `RewriteType::Get()` that returns a statically initialized
+ * instance of RewriteType. */
+#define TVM_DF_PATTERN_REWRITE_GETTER(RewriteType)                    \
+  static DFPatternRewrite* Get() {                                    \
+    static RewriteType rw;                                            \
+    return &rw;                                                       \
+  }                                                                   \
+  static DFPatternCallback GetCallback() {                            \
+    static DFPatternCallback cb = RewriteType::Get()->MakeCallback(); \
+    return cb;                                                        \
+  }

Review comment:
       1) I don't like Macros
   2) I don't like static initialization
   
   Why not just initialize the object and call the method?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org