comaniac commented on a change in pull request #7731:
URL: https://github.com/apache/tvm/pull/7731#discussion_r600800776
##########
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());
Review comment:
Why we need this check?
##########
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:
You may need to comment here if the order is enforced.
##########
File path: src/relay/transforms/simplify_expr.h
##########
@@ -0,0 +1,78 @@
+/*
+ * 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; \
+ }
+
+/*! \brief A wrapper class defining a rewrite matching a specific pattern. */
+class DFPatternRewrite {
+ public:
+ /*! \brief Returns the rewritten expression. */
+ virtual Expr Callback(const Expr& pre, const Expr& post,
+ const Map<DFPattern, Array<Expr>>& node_map) const = 0;
+
+ /*! \brief Returns the pattern to be used for matching and rewriting. */
+ inline DFPattern Pattern() const { return pattern_; }
+
+ inline bool RequireType() const { return require_type_; }
+
+ inline DFPatternCallback MakeCallback() const {
+ auto func = [this](TVMArgs args, TVMRetValue* rv) {
+ Expr pre = args[0];
+ Expr post = args[1];
+ Map<DFPattern, Array<Expr>> node_map = args[2];
+ *rv = this->Callback(pre, post, node_map);
+ };
+ return DFPatternCallback(pattern_, PackedFunc(func), require_type_);
+ }
+
+ protected:
+ /*! \brief The pattern for matching and rewriting. */
+ DFPattern pattern_;
+ bool require_type_ = true;
Review comment:
docstring
--
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:
[email protected]
