[Mlir-commits] [mlir] [MLIR] Folding unpack and pack sequence in data layout propagation from padded domain (PR #138332)

[llvmlistbot at llvm.org]

================
@@ -455,13 +455,10 @@ func.func @unpack_on_output(%arg0: tensor<12x2x56x56x32xf32>) -> tensor<12x56x56
 // CHECK:         %[[UNPACKED_ARG0:.+]] = linalg.unpack %[[ARG0]]
 // CHECK-SAME:      outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]
 // CHECK-SAME:      into %[[ARG0_EMPTY_UNPACK]]
-// CHECK:         %[[ARG0_EMPTY_PACK:.+]] = tensor.empty() : tensor<12x2x56x56x32xf32>
-// CHECK:         %[[PACKED_ARG0:.+]] = linalg.pack %[[UNPACKED_ARG0]]
-// CHECK-SAME:      outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]
-// CHECK-SAME:      into %[[ARG0_EMPTY_PACK]]
+// CHECK:         %[[EMPTY:.+]] = tensor.empty() : tensor<12x2x56x56x32xf32>
 // CHECK:         %[[RES:.+]] = linalg.generic
 // CHECK-SAME:      indexing_maps = [#[[$MAP]]]
-// CHECK-SAME:      outs(%[[PACKED_ARG0]]
+// CHECK-SAME:      outs(%[[EMPTY]]
----------------
Max191 wrote:

Thanks for catching this! It looks like the issue is probably with: https://github.com/llvm/llvm-project/blob/540cf25a6df56fa1810a7411477dca9896aeed20/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp#L331-L335

I think the intention may have been to do something similar to above, where we look for an unpack producer of the pack, and then using the source of the unpack (`%arg0` in this case):
https://github.com/llvm/llvm-project/blob/540cf25a6df56fa1810a7411477dca9896aeed20/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp#L318-L322

I think the original author of this PR is on vacation right now, but I can send out a quick patch to fix it.

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