[Mlir-commits] [mlir] [mlir][linalg] Reimplement SimplifyPackToExpandShape and SimplifyUnPackToCollapseShape for more cases. (PR #204971)

[Artem Gindinson]

================
@@ -394,3 +397,196 @@ func.func @unpad_like_unpack_with_transpose(%arg0: tensor<32x1x16x64xf32>) -> te
   %0 = linalg.unpack %arg0 inner_dims_pos = [1] inner_tiles = [64] into %empty : tensor<32x1x16x64xf32> -> tensor<32x64x16xf32>
   return %0 : tensor<32x64x16xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func.func @pack_3d_to_5d(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<3x32x64xf32>)
+// CHECK:         %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2, 3], [4]] output_shape [3, 1, 1, 32, 64] : tensor<3x32x64xf32> into tensor<3x1x1x32x64xf32>
+// CHECK:         return %[[EXPANDED]] : tensor<3x1x1x32x64xf32>
+func.func @pack_3d_to_5d(%arg0: tensor<3x32x64xf32>) -> tensor<3x1x1x32x64xf32> {
+  %empty = tensor.empty() : tensor<3x1x1x32x64xf32>
+  %0 = linalg.pack %arg0 inner_dims_pos = [1, 2] inner_tiles = [32, 64] into %empty : tensor<3x32x64xf32> -> tensor<3x1x1x32x64xf32>
+  return %0 : tensor<3x1x1x32x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_3d_to_5d_with_outer_dims_perm(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<3x32x64xf32>)
+// CHECK:         %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2, 3], [4]] output_shape [3, 1, 1, 32, 64] : tensor<3x32x64xf32> into tensor<3x1x1x32x64xf32>
+// CHECK:         return %[[EXPANDED]] : tensor<3x1x1x32x64xf32>
+func.func @pack_3d_to_5d_with_outer_dims_perm(%arg0: tensor<3x32x64xf32>) -> tensor<3x1x1x32x64xf32> {
+  %empty = tensor.empty() : tensor<3x1x1x32x64xf32>
+  %0 = linalg.pack %arg0 outer_dims_perm = [0, 2, 1] inner_dims_pos = [1, 2] inner_tiles = [32, 64] into %empty : tensor<3x32x64xf32> -> tensor<3x1x1x32x64xf32>
+  return %0 : tensor<3x1x1x32x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_3d_to_5d_dynamic_shape(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<32x?x64xf32>)
+// CHECK:         %[[C1:.+]] = arith.constant 1 : index
+// CHECK:         %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
+// CHECK:         %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2], [3, 4]] output_shape [32, 1, %[[DIM1]], 1, 64] : tensor<32x?x64xf32> into tensor<32x1x?x1x64xf32>
+// CHECK:         return %[[EXPANDED]] : tensor<32x1x?x1x64xf32>
+func.func @pack_3d_to_5d_dynamic_shape(%arg0: tensor<32x?x64xf32>) -> tensor<32x1x?x1x64xf32> {
+  %c1 = arith.constant 1 : index
+  %dim1 = tensor.dim %arg0, %c1 : tensor<32x?x64xf32>
+  %empty = tensor.empty(%dim1) : tensor<32x1x?x1x64xf32>
+  %0 = linalg.pack %arg0 outer_dims_perm = [0, 2, 1] inner_dims_pos = [1, 2] inner_tiles = [1, 64] into %empty : tensor<32x?x64xf32> -> tensor<32x1x?x1x64xf32>
+  return %0 : tensor<32x1x?x1x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_nd_with_non_unit_outer_tile_dims_perm(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<3x3x32x64xf32>)
+// CHECK-NOT:     tensor.expand_shape
+// CHECK:         linalg.pack
+func.func @pack_nd_with_non_unit_outer_tile_dims_perm(%arg0: tensor<3x3x32x64xf32>) -> tensor<3x3x1x1x32x64xf32> {
+  %empty = tensor.empty() : tensor<3x3x1x1x32x64xf32>
+  %0 = linalg.pack %arg0 outer_dims_perm = [1, 0, 2, 3] inner_dims_pos = [2, 3] inner_tiles = [32, 64] into %empty : tensor<3x3x32x64xf32> -> tensor<3x3x1x1x32x64xf32>
+  return %0 : tensor<3x3x1x1x32x64xf32>
+
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_with_non_unit_packed_dims(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<4x4xf32>)
+// CHECK-NOT:     tensor.expand_shape
+// CHECK:         linalg.pack
+func.func @pack_with_non_unit_packed_dims(%arg0: tensor<4x4xf32>) -> tensor<2x2x2x2xf32> {
+  %empty = tensor.empty() : tensor<2x2x2x2xf32>
+  %0 = linalg.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [2, 2] into %empty : tensor<4x4xf32> -> tensor<2x2x2x2xf32>
+  return %0 : tensor<2x2x2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_with_non_unit_inner_tile_dims_perm(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<32x32xf32>)
+// CHECK-NOT:     tensor.expand_shape
+// CHECK:         linalg.pack
+func.func @pack_with_non_unit_inner_tile_dims_perm(%arg0: tensor<32x32xf32>) -> tensor<1x1x32x32xf32> {
+  %empty = tensor.empty() : tensor<1x1x32x32xf32>
+  %0 = linalg.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [32, 32] into %empty : tensor<32x32xf32> -> tensor<1x1x32x32xf32>
+  return %0 : tensor<1x1x32x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_5d_to_3d(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<3x1x1x32x64xf32>)
+// CHECK:         %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2, 3], [4]] : tensor<3x1x1x32x64xf32> into tensor<3x32x64xf32>
+// CHECK:         return %[[COLLAPSED]] : tensor<3x32x64xf32>
+func.func @unpack_5d_to_3d(%arg0: tensor<3x1x1x32x64xf32>) -> tensor<3x32x64xf32> {
+  %empty = tensor.empty() : tensor<3x32x64xf32>
+  %0 = linalg.unpack %arg0 inner_dims_pos = [1, 2] inner_tiles = [32, 64] into %empty : tensor<3x1x1x32x64xf32> -> tensor<3x32x64xf32>
+  return %0 : tensor<3x32x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_5d_to_3d_with_outer_dims_perm(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<3x1x1x32x64xf32>)
+// CHECK:         %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2, 3], [4]] : tensor<3x1x1x32x64xf32> into tensor<3x32x64xf32>
+// CHECK:         return %[[COLLAPSED]] : tensor<3x32x64xf32>
+func.func @unpack_5d_to_3d_with_outer_dims_perm(%arg0: tensor<3x1x1x32x64xf32>) -> tensor<3x32x64xf32> {
+  %empty = tensor.empty() : tensor<3x32x64xf32>
+  %0 = linalg.unpack %arg0 outer_dims_perm = [0, 2, 1] inner_dims_pos = [1, 2] inner_tiles = [32, 64] into %empty : tensor<3x1x1x32x64xf32> -> tensor<3x32x64xf32>
+  return %0 : tensor<3x32x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_5d_to_3d_dynamic_shape(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<32x1x?x1x64xf32>)
+// CHECK:         %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2], [3, 4]] : tensor<32x1x?x1x64xf32> into tensor<32x?x64xf32>
+// CHECK:         return %[[COLLAPSED]] : tensor<32x?x64xf32>
+func.func @unpack_5d_to_3d_dynamic_shape(%arg0: tensor<32x1x?x1x64xf32>) -> tensor<32x?x64xf32> {
+  %c2 = arith.constant 2 : index
+  %dim2 = tensor.dim %arg0, %c2 : tensor<32x1x?x1x64xf32>
+  %empty = tensor.empty(%dim2) : tensor<32x?x64xf32>
+  %0 = linalg.unpack %arg0 outer_dims_perm = [0, 2, 1] inner_dims_pos = [1, 2] inner_tiles = [1, 64] into %empty : tensor<32x1x?x1x64xf32> -> tensor<32x?x64xf32>
+  return %0 : tensor<32x?x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_nd_with_non_unit_outer_tile_dims_perm(
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<3x3x1x1x32x64xf32>)
+// CHECK-NOT:     tensor.collapse_shape
----------------
AGindinson wrote:

I'd suggest renaming this and all ensuing negative test functions into `no_simplify_*` / `no_fold_*`, for clarity

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