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`? This way `TupleGetItem` is always
defined so you can clean up `GetFunctionOutput` as well.
----------------------------------------------------------------
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]
