sunggg commented on code in PR #14282:
URL: https://github.com/apache/tvm/pull/14282#discussion_r1143852335
##########
src/relax/transform/decompose_composite_ops.cc:
##########
@@ -110,21 +111,60 @@ class NormInferenceSimplifier : public ExprMutator {
Map<Expr, Expr> batch_norm_map_;
};
+class OpDecomposer : public ExprMutator {
+ public:
+ static Expr Decompose(Expr expr) { return OpDecomposer()(expr); }
+
+ private:
+ using ExprMutator::VisitExpr_;
+ Expr TensorToShape(const Call& call_node) {
+ ICHECK(call_node->args.size() == 1);
+ ICHECK(call_node->struct_info_.defined());
+ Expr expr = call_node->args[0];
+ const ShapeStructInfoNode* sinfo =
GetStructInfoAs<ShapeStructInfoNode>(call_node);
+ ICHECK(sinfo);
+ // call builtin function that converts tensor to shape tuple
+ Var call = builder_->Emit(Call(ExternFunc("vm.builtin.tensor_to_shape"),
{expr}, {},
Review Comment:
Realized that I haven't introduced the relax op for this PackedFunc
"vm.builtin.tensor_to_shape" yet. May I revisit this when I get the e2e
pipeline working and refactor? I left a comment regarding this.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
