altanh commented on a change in pull request #7731:
URL: https://github.com/apache/tvm/pull/7731#discussion_r600823616
##########
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:
while extremely unlikely, in theory it's possible for `shape.size()`
(which is `size_t`) to be greater than INT_MAX, so I just added a check here.
(I have to cast the size later to int64_t to pass the dimension to
`MakeConstantTensor`)
--
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]
