[Mlir-commits] [mlir] [llvm] [mlir][tensor] Fold consumer linalg transpose with producer tensor pack (PR #74206)

[Han-Chung Wang]

================
@@ -114,3 +114,238 @@ func.func @pad_pack_different_padding_value(%src: tensor<16641x16xf32>) -> tenso
 // CHECK-LABEL: func.func @pad_pack_different_padding_value
 // CHECK:         tensor.pad
 // CHECK:         tensor.pack
+
+// -----
+
+func.func @tensor_pack_linalg_transpose_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x57x56x2x32xf32> {
+  %0 = tensor.empty() : tensor<56x2x1x57x32xf32>
+  %pack = tensor.pack %arg0
+    outer_dims_perm = [0, 3, 2, 1]
+    inner_dims_pos = [3]
+    inner_tiles = [32]
+    into %0 : tensor<56x57x1x64xf32> -> tensor<56x2x1x57x32xf32>
+
+  %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
+  %transposed = linalg.transpose
+    ins(%pack : tensor<56x2x1x57x32xf32>)
+    outs(%1 : tensor<1x57x56x2x32xf32>)
+    permutation = [2, 3, 0, 1, 4]
+  return %transposed : tensor<1x57x56x2x32xf32>
+}
+//      CHECK: func @tensor_pack_linalg_transpose_fold(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
+//      CHECK:   %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
+//      CHECK:   %[[PACK:.+]] = tensor.pack %[[ARG0]]
+// CHECK-SAME:      outer_dims_perm = [2, 1, 0, 3]
+// CHECK-SAME:      inner_dims_pos = [3] inner_tiles = [32] 
+// CHECK-SAME:       into %[[INIT]]
+//      CHECK:   return %[[PACK]]
+
+// -----
+
+func.func @tensor_pack_linalg_transpose_fold_with_padding(%arg0: tensor<56x57x1x55xf32>, %padding: f32) -> tensor<1x57x56x2x32xf32> {
+  %0 = tensor.empty() : tensor<56x2x1x57x32xf32>
+  %pack = tensor.pack %arg0 padding_value(%padding : f32)
+    outer_dims_perm = [0, 3, 2, 1]
+    inner_dims_pos = [3]
+    inner_tiles = [32]
+    into %0 : tensor<56x57x1x55xf32> -> tensor<56x2x1x57x32xf32>
+
+  %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
+  %transposed = linalg.transpose
+    ins(%pack : tensor<56x2x1x57x32xf32>)
+    outs(%1 : tensor<1x57x56x2x32xf32>)
+    permutation = [2, 3, 0, 1, 4]
+  return %transposed : tensor<1x57x56x2x32xf32>
+}
+//      CHECK: func @tensor_pack_linalg_transpose_fold_with_padding(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<56x57x1x55xf32>, %[[PADDING:.+]]: f32)
+//      CHECK:   %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
+//      CHECK:   %[[PACK:.+]] = tensor.pack %[[ARG0]] padding_value(%[[PADDING]] : f32)
+// CHECK-SAME:      outer_dims_perm = [2, 1, 0, 3]
+// CHECK-SAME:      inner_dims_pos = [3] inner_tiles = [32] 
+// CHECK-SAME:       into %[[INIT]]
+//      CHECK:   return %[[PACK]]
+
+// -----
+
+func.func @tensor_pack_linalg_transpose_fold_no_outer_dims_perm(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x2x56x57x32xf32> {
+  %0 = tensor.empty() : tensor<56x57x1x2x32xf32>
+  %pack = tensor.pack %arg0
+    inner_dims_pos = [3]
+    inner_tiles = [32]
+    into %0 : tensor<56x57x1x64xf32> -> tensor<56x57x1x2x32xf32>
+
+  %1 = tensor.empty() : tensor<1x2x56x57x32xf32>
+  %transposed = linalg.transpose
+    ins(%pack : tensor<56x57x1x2x32xf32>)
+    outs(%1 : tensor<1x2x56x57x32xf32>)
+    permutation = [2, 3, 0, 1, 4]
+  return %transposed : tensor<1x2x56x57x32xf32>
+}
+//      CHECK: func @tensor_pack_linalg_transpose_fold_no_outer_dims_perm(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
+//      CHECK:   %[[INIT:.+]] = tensor.empty() : tensor<1x2x56x57x32xf32>
+//      CHECK:   %[[PACK:.+]] = tensor.pack %[[ARG0]]
+// CHECK-SAME:      outer_dims_perm = [2, 3, 0, 1]
+// CHECK-SAME:      inner_dims_pos = [3] inner_tiles = [32] 
+// CHECK-SAME:       into %[[INIT]]
+//      CHECK:   return %[[PACK]]
+
+// -----
+
+func.func @tensor_pack_linalg_transpose_fold_tile_dims_transpose(%arg0: tensor<56x64x4x64xf32>) -> tensor<2x2x56x2x32x32x2xf32> {
+  %0 = tensor.empty() : tensor<56x2x2x2x32x2x32xf32>
+  %pack = tensor.pack %arg0
+    outer_dims_perm = [0, 1, 2, 3]
+    inner_dims_pos = [1, 2, 3]
+    inner_tiles = [32, 2, 32]
+    into %0 : tensor<56x64x4x64xf32> -> tensor<56x2x2x2x32x2x32xf32>
----------------
hanhanW wrote:

It is useful if all the dim sizes are different. Because we can rely on verifier to catch some potential bugs. Can you update the tests?

Also, can you either remove `outer_dims_perm = [0, 1, 2, 3]` or update it to non-identity permutation? My suggestion is removing some of them and updating some of them to non-identity permutation in the below tests.

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