[Relay][Pass] Simplify consecutive transpose/layout_transform

[comaniac]

This PR adds a simplify pattern to the SImplifyExpr pass. The pattern simplifies back-to-back transpose and layout_transform ops, which can be introduced by Relay frontends or ConvertLayout pass.

cc @mbrookhart @icemelon9 @anijain2305

[ANSHUMAN87]

Thanks @comaniac for the PR! I was also working on similar approach. This solves my case too.

[ANSHUMAN87]

Overall pretty good for me.
I just have 2 points which needs clarification:
1:> When the axes has negative values, we should handle it.
2:> Now we are just looking at the start of the chain of transpose, i suggest we look at intermediate residue as well.

Please let me know in case my points are not clear. Thanks!

[comaniac]

Thanks for the comments.

1:> When the axes has negative values, we should handle it.

Good point. I'll add the support for it later.

2:> Now we are just looking at the start of the chain of transpose, i suggest we look at intermediate residue as well.

Could you provide an example of this case for better illustration? Thanks.

[ANSHUMAN87]

2:> Now we are just looking at the start of the chain of transpose, i suggest we look at intermediate residue as well.

Could you provide an example of this case for better illustration? Thanks.

For example we have below case:
Expression:
y = nn.conv(x)
y = transpose(y, axes=0, 3, 1, 2)
y = transpose(y)
y = transpose(y)

Expected transformation:
y = nn.conv(x)
y = transpose(y, axes=0, 3, 1, 2)

Above such cases can be multiple in a sequence as well. So we need to have a solution which can handle multiple simplification in one chain. Please let me know, in case i have not clearly expressed. Thanks!

I understand we can have these points covered in follow up PRs as well. If we are not planning to do it in this PR. Suggest to add TODOs for these points in the source code, which will help in future.

[comaniac]

It seems to me that the current pattern already covers this case? You can refer to the test cases I added, which have a ReLU followed by transposes. Is this what you want?

[ANSHUMAN87]

It seems to me that the current pattern already covers this case? You can refer to the test cases I added, which have a ReLU followed by transposes. Is this what you want?

Thanks @comaniac for clarification. It is indeed supported.
Can we add 1 more test case as below, it will help reviewer to understand the flow is supported:

    def before4():
        x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32")  # NCHW
        y = relay.nn.relu(x)
        y = relay.transpose(y)  # Reverse
        y = relay.transpose(y)  # Reverse
        y = relay.transpose(y, axes=[0, 2, 3, 1])
        y = relay.transpose(y)  # Reverse
        y = relay.transpose(y)  # Reverse
        return relay.Function([x], y)

    def expected4():
        x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32")  # NCHW
        y = relay.nn.relu(x)
        y = relay.transpose(y, axes=[0, 2, 3, 1])
        return relay.Function([x], y)

[comaniac]

@ANSHUMAN87 comments are addressed.
Also cc @mbrookhart @icemelon9 @anijain2305

[mbrookhart]

LGTM. Sorry for the delay, I was out the end of last week.

One possible extension would be to try to combine this with the simplify reshape above and target patterns like reshape->transpose->reshape or transpose->reshape->transpose

[ANSHUMAN87]

Thanks @comaniac!

[comaniac]

One possible extension would be to try to combine this with the simplify reshape above and target patterns like reshape->transpose->reshape or transpose->reshape->transpose

That's an interesting point and I just thought about it. However, reshape and transpose cannot be equivalent in any cases, because reshape just changes the definition of accessing the tensor; while transpose reorganizes the element order in memory. Even we combine them to the same patterm, we still need to deal with each case separately and result in no code reuse. Please let me know if I missed something. Thanks.

[mbrookhart]

You're right. I think a set of modular patterns is probably best

[comaniac]

Thanks @mbrookhart @ANSHUMAN87