comaniac commented on a change in pull request #5493:
URL: https://github.com/apache/incubator-tvm/pull/5493#discussion_r418668212
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -374,19 +270,16 @@ class Partitioner : public ExprMutator {
* node to create the function for the subgraph.
*/
void CreateFunction(AnnotatedRegion region, const CallNode* call) {
- // Create fields which is a unique list of outputs. Also populate
- // region_return_indices_ map which maps parent of compiler_end node to
- // corresponding index in fields.
+ // Create fields which is a unique list of outputs.
Array<Expr> fields;
int i = 0;
for (auto ret : region->GetOutputs()) {
auto ret_node = Downcast<Call>(ret)->args[0];
// Don't duplicate outputs.
- if (!region_return_indices_.count(region) ||
- !region_return_indices_[region].count(ret_node)) {
- auto ret_expr = VisitExpr(ret_node);
+ if (!region_func_meta_[region].out_expr_indices.count(ret_node)) {
+ auto ret_expr = MixedModeMutator::VisitExpr(ret_node);
fields.push_back(ret_expr);
- region_return_indices_[region][ret_node] = i;
+ region_func_meta_[region].out_expr_indices[ret_node] = {i,
TupleGetItem()};
Review comment:
Good point. Modified accordingly.
----------------------------------------------------------------
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]
