[Graph memory plan] Support nested tuples

[masahi]

The fix for the issue https://discuss.tvm.apache.org/t/graph-plan-memory-doesnt-support-nested-tuples/8278

A token is created for each tensor in a nested tuple. Graph runtime doesn't deal with tuples, so nested tuples are all flattened (runtime.get_num_outputs() and runtime.get_output(i) don't take tuples or nesting into account).

please review @tqchen @zhiics (cc @mbaret @manupa-arm )
I didn't look into memory planning in detail, but I hope this is correct.

[jroesch]

Does this handle 2 levels of nesting?

[masahi]

yes, there is VisitExpr(field) inside GetToken, where nesting is recursively visited and token generated for each tensor in the nest.

[jroesch]

For example what if x2 is also a tuple.

[masahi]

just updated this test to verify two level nesting works.

[mbaret]

This change lgtm, but I have a question relating to the original discuss post. I've been looking at the compile_engine recently and noticed a number of TODOs there with the comment 'Allow recursive tuple', see for instance: https://github.com/apache/incubator-tvm/blob/3222cad9464cadafcab29dbfb5cf149bf083913f/src/relay/backend/compile_engine.cc#L118

It also seems to me that there's logic in graph_runtime_codegen.cc that flattens tuples assuming they're not recursive. Therefore, I'm a bit surprised it was only graph_plan_memory that was blocking you.

[masahi]

@mbaret yes, so this patch suffices for my simple use case of returning nested tuples from the main function. But for use cases like BYOC where the checked_type of CallNode->op can be a nested tuple, it seems I need further fixes (one in graph_plan_memory.cc and another in graph_runtime_codegen.cc).

For example, I've just tried compiling this graph containing a sub function with nested tuple output, and it doesn't compile. I'll work on the fix, that should enable removing the tuple flatting in the BYOC partitioning pass @manupa-arm

#[version = "0.0.5"]
def @func0(%x: Tensor[(10, 10), uint8]) -> (Tensor[(5, 10), uint8], (Tensor[(5, 10), uint8], Tensor[(5, 10), uint8])) {
  %0 = split(%x, indices_or_sections=2) /* ty=(Tensor[(5, 10), uint8], Tensor[(5, 10), uint8]) */;
  %1 = %0.0;
  %2 = %0.1;
  %3 = %0.0;
  %4 = abs(%3) /* ty=Tensor[(5, 10), uint8] */;
  %5 = (%2, %4);
  (%1, %5)
}

def @main(%a: Tensor[(10, 10), uint8]) -> (Tensor[(5, 10), uint8], (Tensor[(5, 10), uint8], Tensor[(5, 10), uint8])) {
  @func0(%a) /* ty=(Tensor[(5, 10), uint8], (Tensor[(5, 10), uint8], Tensor[(5, 10), uint8])) */
}

I've been looking at the compile_engine recently and noticed a number of TODOs there with the comment 'Allow recursive tuple'

Those TODOs look unrelated to graph memory planning. They refer to Function params and shape func return type not supporting nested tuples, and that's definitely still the case.

[masahi]

@tqchen @zhiics @jroesch @mbaret Supporting nested tuples in CallNode->op return type required making non trivial changes to multiple files in the commit d384062

Maybe it would help BYOC use cases, but otherwise I don't know if we want to support this in practice. So I'm happy to revert that change and only support returning nested tuples from main function.

[zhiics]

I am actually okay without the new change

[mbaret]

I'm fine to leave the further changes out of this PR (especially as it's titled with graph plan memory). If it looks more complicated, that may need to be the subject of an RFC.

[masahi]

ok reverted

[masahi]

Thanks @zhiics @jroesch @mbaret