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

[Jerry Shih]

================
@@ -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> {
----------------
JerryShih wrote:

@banach-space 

case1:
```
linalg.pack %arg0 inner_dims_pos = [1, 2] inner_tiles = [32, 64] into %empty : tensor<3x32x64xf32> -> tensor<3x1x1x32x64xf32>
```
steps:
```
expand for packed dims:
3x32x64
=>
3x1x32x1x64

transpose to make inner tile to the last dims:
3x1x32x1x64
=>
3x1x1x32x64
since the outer dims are all 1, we could do nothing for transpose
```

case2:
```
linalg.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [2, 2] into %empty : tensor<4x4xf32> -> tensor<2x2x2x2xf32>
```
steps:
```
expand for packed dims:
4x4
=>
2x2x2x2

transpose to make inner tile to the last dims:
2_outer0*2_inner0*2_outer1*2_inner1*
=>
2_outer0*2_outer1*2_inner0*2_inner1*
We actually need transpose here.
The expand_shape op can't handle the transpose. So, this is the negative case.
```

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