altanh commented on a change in pull request #7731:
URL: https://github.com/apache/tvm/pull/7731#discussion_r600858326
##########
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:
I can add x * 0 = 0 although to me it semantically doesn't exactly fit
and would need slightly different logic, perhaps a new rewrite `ZeroMultiply`
or something?
And I wasn't aware of the commutative matching, that's helpful (although I
wonder if it should be more visibly defined somewhere?)
--
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]
