[Mlir-commits] [mlir] [MLIR][Linalg] Fix empty tensor assumptions for linalg.pack decomposition (PR #160246)

[Andrzej Warzyński]

banach-space wrote:

Thanks for the fix!

> The original code seemed to assume that the tiling dimensions for the tensor.empty op before applying the transposing were always the last dimensions. 

I think that it's a bit more subtle:
* If [tileSizes](https://github.com/llvm/llvm-project/blob/0d989b2aefe8d8f66981f1d39f56ce0214b5de86/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp#L1159) **are not transposed**, then we need to transpose them.
* If [tileSizes](https://github.com/llvm/llvm-project/blob/0d989b2aefe8d8f66981f1d39f56ce0214b5de86/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp#L1159) **are  transposed**, then we do not need to transpose them.

In your example, the inner tiles are not transposed:
```mlir
func.func @pack_with_zero_pos_tile_size(%arg0: tensor<8x1x1x1xf32>, %arg1:tensor<1x1x1x1x8x1xf32>) -> tensor<1x1x1x1x8x1xf32> {
  %pack = linalg.pack %arg0 outer_dims_perm = [0, 1, 2, 3] inner_dims_pos = [0, 3] inner_tiles = [8, 1] into %arg1: tensor<8x1x1x1xf32> -> tensor<1x1x1x1x8x1xf32>
  return %pack : tensor<1x1x1x1x8x1xf32>
}
```

However, in [this example,](https://github.com/llvm/llvm-project/blob/0d989b2aefe8d8f66981f1d39f56ce0214b5de86/mlir/test/Dialect/Linalg/decompose-pack.mlir#L160-L181), we do transpose the inner tiles:
```mlir
func.func @simple_pad_and_pack_dynamic_tile_not_all_dims_tiled(%input: tensor<1x1x5x1xf32>, %output: tensor<1x1x1x1x2x?xf32>, %pad: f32, %high: index) -> tensor<1x1x1x1x2x?xf32> {
  %0 = linalg.pack %input padding_value(%pad : f32) outer_dims_perm = [1, 0, 2, 3] inner_dims_pos = [3, 2] inner_tiles = [2, %high] into %output : tensor<1x1x5x1xf32> -> tensor<1x1x1x1x2x?xf32>
  return %0 : tensor<1x1x1x1x2x?xf32>
}
``` 

I think that tracking _what_ and _when_ to transpose is fragile. Let me proposed a simpler fix:
```cpp
SmallVector<OpFoldResult> transShapeForEmptyOp(srcRank - numTiles, oneIdxAttr);
transShapeForEmptyOp.append(packOp.getMixedTiles());
```

WDYT? Could this work?

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