mbaret commented on a change in pull request #5616:
URL: https://github.com/apache/incubator-tvm/pull/5616#discussion_r427170802
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -404,18 +404,85 @@ IRModule RemoveDefaultAnnotations(IRModule module) {
return module;
}
+/*! \brief There can be regions with multiple outputs where each output
+ * could be a tuple output. Such tuple outputs needs to be flattened
+ * otherwise the function would create tuples of tuples.
+ */
+
+// New annotations would be required to be added for each flattened output
+const PackedFunc* make_end_op =
runtime::Registry::Get("relay.op.annotation._make.compiler_end");
+
+IRModule FlattenTupleOutputs(IRModule module) {
+ class TupleOutFlattener : public ExprRewriter {
+ public:
+ TupleOutFlattener() = default;
+
+ Expr InsertAnnotation(const Expr& expr, const std::string& target, const
PackedFunc* ann_op) {
+ Expr new_op = (*ann_op)(expr, target);
+ new_op->checked_type_ = expr->checked_type_;
+ return new_op;
+ }
+
+ Expr Rewrite_(const CallNode* call, const Expr& post) final {
+ if (call->op == compiler_end_op) {
+ std::string target = call->attrs.as<CompilerAttrs>()->compiler;
+ // Arguments of annotation ops should be 1
+ CHECK_EQ(call->args.size(), 1U);
+ auto annotated_op = Downcast<Call>(post)->args[0];
+ if (annotated_op->IsInstance<TupleNode>()) {
+ auto tn = annotated_op.as<TupleNode>();
+ Array<Expr> new_fields;
+
+ // Here each input of the tuple will be annotated with compiler_ends
+ for (auto& tn_arg : tn->fields) {
+ auto nf = InsertAnnotation(tn_arg, target, make_end_op);
+ new_fields.push_back(nf);
+ }
+
+ // Return a tuple of compiler_ends in the place of the tuple that was
+ // annotated with a compiler_end.
+ auto out = Tuple(new_fields);
+ return std::move(out);
+ }
+ }
+ return post;
+ }
+ };
+
+ auto glob_funcs = module->functions;
+ // module is mutable, hence, we make a copy of it.
+ module.CopyOnWrite();
+ for (const auto& pair : glob_funcs) {
+ if (auto* fn = pair.second.as<FunctionNode>()) {
+ auto func = GetRef<Function>(fn);
+ TupleOutFlattener to_flattener;
+ auto removed = PostOrderRewrite(func->body, &to_flattener);
+ func = Function(func->params, removed, func->ret_type,
func->type_params, func->attrs);
+ module->Update(pair.first, func);
+ }
+ }
+ return module;
+}
+
} // namespace partitioning
namespace transform {
Pass PartitionGraph() {
runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> part_func =
[=](IRModule m,
PassContext pc) {
+ // There could be compiler_end annotations on tuples
+ // If the corresponding region is having multiple compiler_ends,
+ // this would lead to creation of tuples of tuples.
+ // Thus, we flatten the tuples by transfering the compiler_end to
+ // the tuple inputs.
+ auto _m = partitioning::FlattenTupleOutputs(m);
Review comment:
I'm in favour of keeping this a separate pass. We could technically
incorporate `RemoveDefaultAnnotations` directly into the logic of `Partition`
but we don't because it's less clear and would require introducing
special-cased behaviour. I think a similar argument applies here. The logic of
`Partition` is already reasonably challenging to follow, so anything to reduce
code complexity is, I think, for the best.
If `Partition` was actually broken without this pass being run, then I think
there would be stronger argument for incorporating it. But as it stands, this
is required because the runtime doesn't (currently) support nested tuples.
----------------------------------------------------------------
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]
