[Mlir-commits] [mlir] [mlir][linalg] Fix and Refactor DecomposeOuterUnitDimsUnPackOpPattern (PR #119379)

[Andrzej Warzyński]

================
@@ -35,15 +35,15 @@ func.func @simple_unpack_static_tiles(%input: tensor<1x1x8x2xf32>, %output: tens
 
 /// Same as example above, but with 1 dynamic tile size.
 
-func.func @simple_unpack_dynamic_tile(%input: tensor<1x1x?x2xf32>, %output: tensor<5x1xf32>, %tile_dim_0: index) -> tensor<5x1xf32> {
-  %0 = tensor.unpack %input inner_dims_pos = [0, 1] inner_tiles = [%tile_dim_0, 2] into %output : tensor<1x1x?x2xf32> -> tensor<5x1xf32>
+func.func @simple_unpack_dynamic_tile(%input: tensor<1x1x?x2xf32>, %output: tensor<5x1xf32>, %tile_dim: index) -> tensor<5x1xf32> {
+  %0 = tensor.unpack %input inner_dims_pos = [0, 1] inner_tiles = [%tile_dim, 2] into %output : tensor<1x1x?x2xf32> -> tensor<5x1xf32>
   return %0 : tensor<5x1xf32>
 }
 // CHECK-LABEL: func.func @simple_unpack_dynamic_tile
 // CHECK-SAME:    %[[SRC:[a-zA-Z0-9]+]]
 // CHECK-SAME:    %[[DEST:[a-zA-Z0-9]+]]
-// CHECK-SAME:    %[[TILE_DIM_1:[a-zA-Z0-9]+]]
-// CHECK:         %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0] [1, 1, %[[TILE_DIM_1]], 2] [1, 1, 1, 1]
+// CHECK-SAME:    %[[TILE_DIM:[a-zA-Z0-9]+]]
----------------
banach-space wrote:

After the pass, there's no longer my carefully selected `tile_dim`. Instead, there's `arg2` 😂 
```mlir
  func.func @simple_unpack_dynamic_tile(%arg0: tensor<1x1x?x2xf32>, %arg1: tensor<5x1xf32>, %arg2: index) -> tensor<5x1xf32> {
    %extracted_slice = tensor.extract_slice %arg0[0, 0, 0, 0] [1, 1, %arg2, 2] [1, 1, 1, 1] : tensor<1x1x?x2xf32> to tensor<?x2xf32>
    %extracted_slice_0 = tensor.extract_slice %extracted_slice[0, 0] [5, 1] [1, 1] : tensor<?x2xf32> to tensor<5x1xf32>
    return %extracted_slice_0 : tensor<5x1xf32>
  }
```

Good catch though! Something to pay attention to.

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