altanh commented on a change in pull request #7731:
URL: https://github.com/apache/tvm/pull/7731#discussion_r600826676
##########
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:
(in any case it shouldn't matter now that I've added support for
eliminating constants)
--
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]
