Look for TupleType instead of TupleNode in LayoutRewriter

[trevor-m]

When getting the input shapes from the args LayoutRewriter checked for tuples only by seeing if the node is an instance of TupleNode. With that code, the following relay program would encounter an error during LayoutRewriter because the arg to the function is a relay.Var but has a TupleTypeNode instead of a TensorTypeNode as it expected.

def @main(%input_0: Tensor[(1, 2, 6, 6), float32], %input_1: Tensor[(1, 3, 6, 6), float32]) -> Tensor[(1, 5, 6, 6), float32] {
  %0 = (%input_0, %input_1);
  %1 = fn (%x: (Tensor[(1, 2, 6, 6), float32], Tensor[(1, 3, 6, 6), float32])) -> Tensor[(1, 5, 6, 6), float32] {
    concatenate(%x, axis=1) /* ty=Tensor[(1, 5, 6, 6), float32] */
  };
  %1(%0) /* ty=Tensor[(1, 5, 6, 6), float32] */
}

@zhiics @anijain2305 Could you please review?

[anijain2305]

Trevor, can you please add a test case? The implementation is good, but we need a test case.

[zhiics]

IMHO, I don't think we need to test this functionality in the partitioning pass as it is a unit test of layout itself. Instead, You can directly make the "partitioned" graph and run Alterlayout pass and assert whatever IR it should emit.

[trevor-m]

IMHO, I don't think we need to test this functionality in the partitioning pass as it is a unit test of layout itself. Instead, You can directly make the "partitioned" graph and run Alterlayout pass and assert whatever IR it should emit.

Hi Zhi, I don't think it is possible to create the partitioned graph without actually using partitoning. I recreated the subgraph below:

input_type = relay.TensorType((1, 5, 6, 6), "float32")
x = relay.var("x", relay.TupleType([input_type, input_type]))
out = relay.concatenate(x, axis=1)
func = relay.Function([x], out)

However, the relay.concatenate python API does not accept a Var node as the input:

  File "test_pass_alter_op_layout.py", line 1053, in test_concatenate
    out = relay.concatenate(x, axis=1)

  File "/data/neo-ai-tvm/python/tvm/relay/op/tensor.py", line 888, in concatenate
    data = list(data)

TypeError: 'Var' object is not iterable

Edit: I was able to bypass this by calling tvm.relay.op._make.concatenate. Updating PR...

[anijain2305]

#5066

Above PR have same line changes. Might be a good idea to try that first.

[tqchen]

@trevor-m @anijain2305 @zhiics please followup

[trevor-m]

Closing since #5066 also solves this issue.