tvm-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From GitBox <...@apache.org>
Subject [GitHub] [incubator-tvm] manupa-arm commented on a change in pull request #5616: [RELAY][BYOC] Fix the creation of tuple of tuples in PartitionGraph
Date Tue, 19 May 2020 09:27:17 GMT

manupa-arm commented on a change in pull request #5616:
URL: https://github.com/apache/incubator-tvm/pull/5616#discussion_r427160275



##########
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:
       > Invalid naming
   
   Not sure I get you, any suggestion ?




----------------------------------------------------------------
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:
users@infra.apache.org



Mime
View raw message