[Mlir-commits] [mlir] [mlir][vector] Fold transpose(broadcast(shape_cast)) to broadcast (PR #215940)

[Adam Siemieniuk]

adam-smnk wrote:

> applicable to broadcast: the existing transpose -> shape_cast folders already drop information that cannot be recovered (!). Two transposes that permute unit dims differently fold to the exact same shape cast. For example, both of these:
> 
> ```
> %0 = vector.shape_cast %arg : vector<4xf32> to vector<4x1x1xf32>
> %1 = vector.transpose %0, [2, 0, 1] : vector<4x1x1xf32> to vector<1x4x1xf32>
> 
> %0 = vector.shape_cast %arg : vector<4xf32> to vector<4x1x1xf32>
> %1 = vector.transpose %0, [1, 0, 2] : vector<4x1x1xf32> to vector<1x4x1xf32>
> ```
>
> canonicalize to the same:
> 
> ```
> %0 = vector.shape_cast %arg : vector<4xf32> to vector<1x4x1xf32>
> ```

To preface, I think your example of dimension/layout tracking is a valid use case. However, isn't that a higher meta-level interpretation of IR?
I feel there are some implicit assumptions here that are not captured by IR itself.

I'm not sure I see what information is exactly lost.
I don't think mathematically there is any difference between these layouts. If we just focus on the input and final output, the mapping is performed from 1D to the same <1x4x1> shape. Shuffling only unit dimensions should be rather invisible.

There could be implications in the underlying hardware but I don't think any of these ops are capable of capturing these extended semantics. Alternatively, more pessimistic approach would mean that a transpose op can never be rewritten due to hardware side-effects.


https://github.com/llvm/llvm-project/pull/215940