masahi commented on a change in pull request #5493:
URL: https://github.com/apache/incubator-tvm/pull/5493#discussion_r418541385
##########
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:
Can we move this to the end of this function so that we can properly
initialize `TupleGetItem` with `func_call` and index? This way you can remove
index from `out_expr_indices`. Since `TupleGetItem` is always defined, you can
also clean up `GetFunctionOutput` (always return the `TupleGetItem` already
created).
----------------------------------------------------------------
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]
