Enhance shape inference for ONNX Reshape

[tungld]

This patch is to enhance the shape inference for ONNX Reshape in the case where we can infer a static value for the output dimension at position of -1.

In Reshape, when the second input shape is a 1D tensor containing the dimensions of the output. shape can accept a value of -1, meaning that the output dimension is computed from the remaining dimensions.
In some situations, the remaining dimensions have unknown dimensions but these unknown dimensions are the same as the the unknown dimensions in the input. So we can remove these unknown dimensions when computing the dimension at -1, which increases the chance of getting a static value for position -1.

For example, in the following program, we can infer the dimension at position -1 (the second dim from the last of tensor<?x?x?x64xf32) to be 32 instead of unknown.

  %1 = onnx.Constant dense<64> : tensor<1xi64>
  %2 = onnx.Constant dense<-1> : tensor<1xi64>
  %3 = "onnx.Dim"(%arg0) {axis = 0 : si64} : (tensor<?x?x2048xf32>) -> tensor<1xi64>
  %4 = "onnx.Dim"(%arg0) {axis = 1 : si64} : (tensor<?x?x2048xf32>) -> tensor<1xi64>
  %5 = "onnx.Concat"(%3, %4, %2, %1) {axis = 0 : si64} : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<4xi64>
  %6 = "onnx.Reshape"(%arg0, %5) {allowzero = 0 : si64} : (tensor<?x?x2048xf32>, tensor<4xi64>) -> tensor<?x?x?x64xf32>

In this example, the output shape is [%3, %4, -1, 64]. Because %3 and %4 are the same as the 1st and 2nd dims of %arg0 respectively, we can ignore them in computing dim at position -1. So the dim at position -1 is simple 2048/64 or 32.

[AlexandreEichenberger]

Smart! Added two questions to make sure the corner cases are covered.

[AlexandreEichenberger]

Is this on purpose at the top? If the last i iteration results in isBijective is set to false, then this code is never executed. If that is what you want, maybe a small comments to explain the corner case where bijective is false but the sets are not cleared.

[tungld]

Good catch! I fixed this case by moving the check outside the for loop. Thanks!

[AlexandreEichenberger]

A comment on the condition being tested would be useful.
Are you also testing the case where (though unlikely)

(%1, %2, 64) -> (%1, %1, %2, -1)

namely where %1 appears twice, which would be clearly not bijective.

[tungld]

Added a lit test for this case.

[chentong319]

For readability, can we propagate 'onnx.dim_params' attribute introduced by onnx model along onnx.dim op?

[chentong319]

It seems to me that the code works on the restricted pattern with the data of reshape coming from MatMul. Could we relax the condition a little bit for a general case for data?
We need onnx.dim(%data, i) and value of %shape[i] to compute the size of unspecified dimension.
For the value of shape[I], we track back to the defining op, onnx.ConcatOp, and then to its inputs of constant or onnx.dim op. That's the pattern appears in the model. Anyway, we may define a utility function with a general name in case we find more pattern to locate the definition from onnx.dim.
For the onnx.dim(%data, i), do we have to track back to its defining op or simply rely on DimAnalysis to check it with the onnx.dim from the value of shape?
Anyway, in the process of cancel out the same onnx.dim, I think that the canceled one should be removed from the list after matched. For example, in 'onnx.Reshape(%data, %reshape): (tensor<?x?x?x256xf32>, tensor<4xi64>)`, there may be two symbols, one of which appears twice in the dimension.

[tungld]

@chentong319 trying to understand your comment.

For the onnx.dim(%data, i), do we have to track back to its defining op or simply rely on DimAnalysis to check it with the onnx.dim from the value of shape?

I think using DimAnalysis would be definitely better, so that we can check dim equality for two arbitrary dims without tracing back. Unfortunately, it is non-trivial to use DimAnalysis inside shape inference at this moment because 1) DimAnalysis uses ShapeHelper here so we may have a cyclic dependency, 2) DimAnalysis depends on the IR at a time, whenever the IR changes DimAnalysis needs to be updated.

Anyway, in the process of cancel out the same onnx.dim, I think that the canceled one should be removed from the list after matched. For example, in 'onnx.Reshape(%data, %reshape): (tensor<?x?x?x256xf32>, tensor<4xi64>)`, there may be two symbols, one of which appears twice in the dimension.

Yes, it looks similar to Alex's example, I added a lit test for this case.

[tungld]

For readability, can we propagate 'onnx.dim_params' attribute introduced by onnx model along onnx.dim op?

It sounds a good thing. Will think about that.

[chentong319]

1. Agree with complexity of using DimAnalysis.

2. The example I was thinking of is more like this:
   %0 = onnx.Dim(%D1, 0): (tensor<?xf32> {onnx.dim_params = "0:dimA"} -> tensor<1xi64>
   %1 = onnx.Dim(%D2, 0): (tensor<?xf32> {onnx.dim_params = "0:dimA"} -> tensor<1xi64>
   %2 = onnx.Dim(%D3, 0): (tensor<?xf32> {onnx.dim_params = "0:dimB"} -> tensor<1xi64>
   %3 = onnx.Constant<-1>
   %4 = onnx.Constant<64>
   %shape = onnx.Concat(%0, %1, %2, %3, %4, %5)
   onnx.Reshape(%data, %shape) : (tensor<?x?x?x256xf32> {onnx.dim_params = "0:dimA,1:dimB, 2:dimA"}, tensor<5xi64>) -> tensor<*xf32>

I assume that the dim_params is propagated so that it is easier for me to show the relationship of the symbolic dimensions.
This result should be tensor<?x?x?x4x64xf32>
It is a general case, which we may not need to worry about now.

By the way, from the example I constructed, it is very useful to propagate the dim_params. It is a simplified symbolic DimAnalysis. We should do it. Schedule a discussion?

[chentong319]

LGTM!

[tungld]

tensor<?xf32> {onnx.dim_params = "0:dimB"}

Is it possible to define that attribute {onnx.dim_params = "0:dimB"} for a tensor? I know it's possible for function arguments, but not sure about a single tensor type.

[jenkins-droid]

Jenkins Linux amd64 Build #16539 [push] null... failed after 1 hr 26 min

[jenkins-droid]

Jenkins Linux s390x Build #16541 [push] null... failed after 1 hr 44 min