Update(2) [mlir][linalg] Update pack and unpack documentation

Signed-off-by: Christopher McGirr <[email protected]>

Author: chrsmcgrr

mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td

@@ -100,7 +100,8 @@ def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
     `inner_dims_pos` (mandatory) specifies `k` source tensor dimensions that are
     being tiled, where `0 <= k <= n`.
      - `inner_dims_pos[i]` specifies the source tensor dimension tiled by
-    `inner_tiles[i]` where `0 <= i < k`.
+    `inner_tiles[i]` where `0 <= i < k`. All the values in `inner_dims_pos` are
+    within [0, n).
      - The tiled dimensions (of size `inner_tiles`) are added to the end of the
      result tensor in the order in which they appear, i.e.
      `shape(result)[rank(result) + i] = inner_tiles[i]` for `0 <= i < k`.
@@ -275,7 +276,7 @@ def Linalg_UnPackOp : Linalg_RelayoutOp<"unpack"> {
     packed source tensor. The source tensor (i.e. packed tensor) dimensions can
     be unpacked given `inner_dims_pos` as follows.
     - For `0 <= i < k` the following relationship holds:
-    `shape(result)[inner_dims_pos[i]] = shape(source)[n-k+i] + shape(source)[inner_dims_pos[i]]`.
+    `shape(result)[inner_dims_pos[i]] <= shape(source)[n-k+i] * shape(source)[inner_dims_pos[i]]`.
     - For `0 <= j < n-k` and `j` not in `inner_dims_pos` the following relationship holds:
     `shape(result)[j] = shape(source)[j]`.
 

mlir/test/Dialect/Linalg/invalid.mlir

@@ -1824,9 +1824,8 @@ func.func @unpack_invalid_outer_dims_perm(%source: tensor<128x256xf32>, %dest: t
 
 // -----
 
-// Here we have the source tensor being tiled as: `source[1] / 32` and `source[0] / 16` but the inner_dims_pos does not imply
-// a transpose of the outer dimensions for the result tensor. The tiled dimensions appear in the result tensor in the order
-// they appear in the source tensor, i.e. 16x4x32x16
+// The outer dims in the output tensor are incorrectly/unexpectedly transposed.
+// This could be fixed by adding `outer_dims_perm = [1, 0]` (the default value assumes no transpose).
 func.func @pack_invalid_result_shape(%input: tensor<256x128xf32>, %output: tensor<4x16x32x16xf32>) -> tensor<4x16x32x16xf32> {
   // expected-error@+1 {{the shape of output is not large enough to hold the packed data. Expected at least 'tensor<16x4x32x16xf32>', got 'tensor<4x16x32x16xf32>'}}
   %0 = linalg.pack %input inner_dims_pos = [1, 0] inner_tiles = [32, 16] into %output : tensor<256x128xf32> -> tensor<4x16x32x16xf32>

mlir/test/Dialect/Linalg/named-ops.mlir

@@ -2771,12 +2771,12 @@ func.func @pad_and_pack_partially_dynamic(%source: tensor<?x?xf32>, %dest: tenso
 
 // -----
 
-func.func @pack_descending_inner_dims_with_padding(%source: tensor<1x5x7xf32>, %dest: tensor<1x3x2x4x2xf32>, %pad: f32) -> tensor<1x3x2x4x2xf32> {
+func.func @pack_transposed_inner_dims_with_padding(%source: tensor<1x5x7xf32>, %dest: tensor<1x3x2x4x2xf32>, %pad: f32) -> tensor<1x3x2x4x2xf32> {
   %0 = linalg.pack %source padding_value(%pad : f32) inner_dims_pos = [2, 1] inner_tiles = [4, 2] into %dest : tensor<1x5x7xf32> -> tensor<1x3x2x4x2xf32>
   return %0 : tensor<1x3x2x4x2xf32>
 }
 
-// CHECK-LABEL: func.func @pack_descending_inner_dims_with_padding(
+// CHECK-LABEL: func.func @pack_transposed_inner_dims_with_padding(
 // CHECK-SAME:  %[[SOURCE:.*]]: tensor<1x5x7xf32>,
 // CHECK-SAME:  %[[DEST:.*]]: tensor<1x3x2x4x2xf32>,
 // CHECK-SAME:  %[[PAD:.*]]: f32)
@@ -2834,6 +2834,38 @@ func.func @unpack_non_adjacent_inner_dims(%source: tensor<10x1x3x4x2xf32>, %dest
 
 // -----
 
+func.func @pack_implementing_transpose(%source: tensor<3x5x7xf32>, %dest: tensor<3x7x5xf32>) -> tensor<3x7x5xf32> {
+  %0 = linalg.pack %source outer_dims_perm = [0, 2, 1] inner_dims_pos = [] inner_tiles = [] into %dest : tensor<3x5x7xf32> -> tensor<3x7x5xf32>
+  return %0 : tensor<3x7x5xf32>
+}
+
+// CHECK-LABEL: func.func @pack_implementing_transpose(
+// CHECK-SAME:  %[[SOURCE:.*]]: tensor<3x5x7xf32>,
+// CHECK-SAME:  %[[DEST:.*]]: tensor<3x7x5xf32>)
+// CHECK:       %{{.*}} = linalg.pack
+// CHECK-SAME:      outer_dims_perm = [0, 2, 1]
+// CHECK-SAME:      inner_dims_pos = []
+// CHECK-SAME:      inner_tiles = []
+// CHECK-SAME:      into %[[DEST]] : tensor<3x5x7xf32> -> tensor<3x7x5xf32>
+
+// -----
+
+func.func @unpack_implementing_transpose(%source: tensor<3x7x5xf32>, %dest: tensor<3x5x7xf32>) -> tensor<3x5x7xf32> {
+  %0 = linalg.unpack %source outer_dims_perm = [0, 2, 1] inner_dims_pos = [] inner_tiles = [] into %dest : tensor<3x7x5xf32> -> tensor<3x5x7xf32>
+  return %0 : tensor<3x5x7xf32>
+}
+
+// CHECK-LABEL: func.func @unpack_implementing_transpose(
+// CHECK-SAME:  %[[SOURCE:.*]]: tensor<3x7x5xf32>,
+// CHECK-SAME:  %[[DEST:.*]]: tensor<3x5x7xf32>)
+// CHECK:       %{{.*}} = linalg.unpack
+// CHECK-SAME:      outer_dims_perm = [0, 2, 1]
+// CHECK-SAME:      inner_dims_pos = []
+// CHECK-SAME:      inner_tiles = []
+// CHECK-SAME:      into %[[DEST]] : tensor<3x7x5xf32> -> tensor<3x5x7xf32>
+
+// -----
+
 func.func @unpack_fully_dynamic(%source: tensor<?x?x?x?xf32>, %dest: tensor<?x?xf32>, %tile_n : index, %tile_m : index) -> tensor<?x?xf32> {
   %0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [%tile_n, %tile_m] into %dest : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>