[Mlir-commits] [mlir] [mlir][linalg] Update pack and unpack documentation (PR #143903)
Christopher McGirr
llvmlistbot at llvm.org
Fri Jun 27 00:20:55 PDT 2025
https://github.com/chrsmcgrr updated https://github.com/llvm/llvm-project/pull/143903
>From 02e50037004f04bc207095b0e79ac21987f4a8bf Mon Sep 17 00:00:00 2001
From: Christopher McGirr <mcgirr at roofline.ai>
Date: Thu, 12 Jun 2025 08:12:38 +0000
Subject: [PATCH 1/3] [mlir][linalg] Update pack and unpack documentation
* Clarified the `inner_dim_pos` attribute in the case of high
dimensionality tensors.
* Added a 5D examples to show-case the use-cases that triggered this
updated.
* Added a reminder for linalg.unpack that number of elements are not
required to be the same between input/output due to padding being
dropped.
I encountered some odd variations of `linalg.pack` and
`linalg.unpack` while working on some TFLite models and the
definition in the documentation did not match what I saw pass in IR
verification.
The following changes reconcile those differences.
Signed-off-by: Christopher McGirr <mcgirr at roofline.ai>
---
.../Dialect/Linalg/IR/LinalgRelayoutOps.td | 51 +++++++++++----
mlir/test/Dialect/Linalg/invalid.mlir | 11 ++++
mlir/test/Dialect/Linalg/named-ops.mlir | 63 +++++++++++++++++++
3 files changed, 114 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
index 1e48a5e3a20ee..fef1900be62ea 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
@@ -94,11 +94,13 @@ def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
and optionally transposes the tiled source tensor dimensions.
`inner_dims_pos` (mandatory) specifies `k` source tensor dimensions that are
- being tiled, where `0 < k <= n`. The order of the dimensions matters:
- - The tiled dimensions (of size `inner_tiles`) are added to the end of the result
- tensor in the order in which they appear in `inner_dims_pos`.
+ being tiled, where `0 < k <= n`.
- `inner_dims_pos[i]` specifies the source tensor dimension tiled by
- `inner_tiles[i]`.
+ `inner_tiles[i]` where `0 <= i < k`.
+ - the resulting tiled source dimension maps to an outer dimension of the
+ packed tensor in the order the non-tiled dimension appeared in the source
+ tensor, i.e. `shape(result)[inner_dims_pos[i]]` is equal to
+ `shape(source)[inner_dims_pos[i]] / inner_tiles[i]`.
`inner_tiles` (mandatory) specifies `k` tile sizes. These tile sizes
correspond to the least significant ("inner") result tensor dimension sizes,
@@ -117,6 +119,16 @@ def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
into %dest : tensor<128x256xf32> -> tensor<16x8 x 8x32 xf32>
// \ / \ /
// outer dims inner dims
+ // CHW to CHWhw
+ %0 = linalg.pack %source inner_dims_pos = [2, 1] inner_tiles = [4, 2]
+ into %dest : tensor<1x8x16xf32> -> tensor<1x2x4 x 4x2 xf32>
+ // \ / \ /
+ // outer dims inner dims
+ // HCW to HCWhw
+ %0 = linalg.pack %source inner_dims_pos = [2, 0] inner_tiles = [4, 2]
+ into %dest : tensor<20x1x12xf32> -> tensor<10x1x3 x 4x2xf32>
+ // \ / \ /
+ // Outer Dims: 10x1x3 Inner Dims: 4x2
```
`outer_dims_perm` (optional) specifies a permutation for the outer
@@ -246,12 +258,14 @@ def Linalg_UnPackOp : Linalg_RelayoutOp<"unpack"> {
The "unpack" operation converts a source tensor of rank `n` with a tiled and
packed layout to a result tensor of rank `n - k`.
- `inner_dims_pos` (mandatory) specifies `k` source tensor dimensions with
- which the last `k` source tensor dimensions are combined, where
- `0 < k <= n/2`. Each `inner_dims_pos` element must be `>= 0` and `< n - k`.
- The order of the dimensions in `inner_dims_pos` matters: dimension
- `inner_dims_pos[i]` is combined with dimension `n - k + i` (assuming that
- `outer_dims_perm` is not specified).
+ `inner_dims_pos` (mandatory) specifies `k` result tensor dimensions that
+ were tiled with the `inner_tiles` to create the packed source tensor. The
+ source tensor dimensions can be combined given `inner_dims_pos` as follows:
+ the inner tile `shape(source)[n-k+i]` is combined with
+ `shape(source)[inner_dims_pos[i]]` where `0 <= i < k` and stored at
+ `shape(result)[inner_dims_pos[i]]`. The remaining dimensions are
+ `shape(result)[j] = shape(source)[j]` where `0 <= j < n-k` and `j` is not in
+ the set of `inner_dims_pos` indices.
`inner_tiles` (mandatory) specifies `k` tile sizes. These tile sizes
correspond to the least significant ("inner") source tensor dimension sizes.
@@ -266,7 +280,11 @@ def Linalg_UnPackOp : Linalg_RelayoutOp<"unpack"> {
dimensions. If specified, it must have `n - k` elements. If specified, this
permutation is applied before combining any dimensions.
- Example:
+ Note, the amount of elements in the source (packed tensor) and the result
+ (unpacked) can be unequal, i.e. `SizeOf(source) >= SizeOf(result)`. As
+ the unpack operation may drop any padding introduced by the pack operation.
+
+ Examples:
```mlir
// NCnc to NC:
@@ -277,6 +295,17 @@ def Linalg_UnPackOp : Linalg_RelayoutOp<"unpack"> {
%0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
inner_tiles = [8, 32] into %dest
: tensor<8x16x8x32xf32> -> tensor<128x256xf32>
+
+ // CHW to CHWhw:
+ %0 = linalg.unpack %source inner_dims_pos = [2, 1] inner_tiles = [4, 2]
+ into %dest : tensor<1x3x2x4x2xf32> -> tensor<1x5x7xf32>
+ // / \
+ // Outer Dims: 1x3x2 Inner Dims: 4x2
+ // HCW to HCWhw
+ %0 = linalg.unpack %source inner_dims_pos = [2, 0] inner_tiles = [4, 2]
+ into %dest : tensor<10x1x3 x 4x2xf32> -> tensor<20x1x12xf32>
+ // / \
+ // Outer Dims: 10x1x3 Inner Dims: 4x2
```
}];
let arguments = (ins AnyRankedTensor:$source,
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index cbc863699ba9e..0f71000f7c700 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -1824,6 +1824,17 @@ 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
+func.func @pack_invalid_result_shape(%input: tensor<256x128xf32>, %output: tensor<4x16x32x16xf32>) -> tensor<4x16x32x16xf32> {
+ // expected-error at +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>
+ return %0 : tensor<4x16x32x16xf32>
+}
+
+// -----
+
func.func @pack_invalid(%input: tensor<256x128xf32>, %output: tensor<8x8x32x16xf32>) -> tensor<8x8x32x16xf32> {
// expected-error at +1 {{the shape of output is not large enough to hold the packed data. Expected at least 'tensor<8x8x16x32xf32>', got 'tensor<8x8x32x16xf32>'}}
%0 = linalg.pack %input inner_dims_pos = [1, 0] inner_tiles = [16, 32] into %output : tensor<256x128xf32> -> tensor<8x8x32x16xf32>
diff --git a/mlir/test/Dialect/Linalg/named-ops.mlir b/mlir/test/Dialect/Linalg/named-ops.mlir
index 470bc1c78640c..b21b234bc7841 100644
--- a/mlir/test/Dialect/Linalg/named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/named-ops.mlir
@@ -2771,6 +2771,69 @@ 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> {
+ %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-SAME: %[[SOURCE:.*]]: tensor<1x5x7xf32>,
+// CHECK-SAME: %[[DEST:.*]]: tensor<1x3x2x4x2xf32>,
+// CHECK-SAME: %[[PAD:.*]]: f32)
+// CHECK: %{{.*}} = linalg.pack
+// CHECK-SAME: inner_dims_pos = [2, 1]
+// CHECK-SAME: inner_tiles = [4, 2]
+// CHECK-SAME: into %[[DEST]] : tensor<1x5x7xf32> -> tensor<1x3x2x4x2xf32>
+
+// -----
+
+// The function suffix "with_padding" refers to the padding that was introduced by the pack operation. But here
+// we are dropping the padding. Creating a tensor with less elements than what we started with.
+func.func @unpack_descending_inner_dims_with_padding(%source: tensor<1x3x2x4x2xf32>, %dest: tensor<1x5x7xf32>) -> tensor<1x5x7xf32> {
+ %0 = linalg.unpack %source inner_dims_pos = [2, 1] inner_tiles = [4, 2] into %dest : tensor<1x3x2x4x2xf32> -> tensor<1x5x7xf32>
+ return %0 : tensor<1x5x7xf32>
+}
+
+// CHECK-LABEL: func.func @unpack_descending_inner_dims_with_padding(
+// CHECK-SAME: %[[SOURCE:.*]]: tensor<1x3x2x4x2xf32>,
+// CHECK-SAME: %[[DEST:.*]]: tensor<1x5x7xf32>)
+// CHECK: %{{.*}} = linalg.unpack
+// CHECK-SAME: inner_dims_pos = [2, 1]
+// CHECK-SAME: inner_tiles = [4, 2]
+// CHECK-SAME: into %[[DEST]] : tensor<1x3x2x4x2xf32> -> tensor<1x5x7xf32>
+
+// -----
+
+func.func @pack_non_adjacent_inner_dims(%source: tensor<20x1x12xf32>, %dest: tensor<10x1x3x4x2xf32>) -> tensor<10x1x3x4x2xf32> {
+ %0 = linalg.pack %source inner_dims_pos = [2, 0] inner_tiles = [4, 2] into %dest : tensor<20x1x12xf32> -> tensor<10x1x3x4x2xf32>
+ return %0 : tensor<10x1x3x4x2xf32>
+}
+
+// CHECK-LABEL: func.func @pack_non_adjacent_inner_dims(
+// CHECK-SAME: %[[SOURCE:.*]]: tensor<20x1x12xf32>,
+// CHECK-SAME: %[[DEST:.*]]: tensor<10x1x3x4x2xf32>)
+// CHECK: %{{.*}} = linalg.pack
+// CHECK-SAME: inner_dims_pos = [2, 0]
+// CHECK-SAME: inner_tiles = [4, 2]
+// CHECK-SAME: into %[[DEST]] : tensor<20x1x12xf32> -> tensor<10x1x3x4x2xf32>
+
+// -----
+
+func.func @unpack_non_adjacent_inner_dims(%source: tensor<10x1x3x4x2xf32>, %dest: tensor<20x1x12xf32>) -> tensor<20x1x12xf32> {
+ %0 = linalg.unpack %source inner_dims_pos = [2, 0] inner_tiles = [4, 2] into %dest : tensor<10x1x3x4x2xf32> -> tensor<20x1x12xf32>
+ return %0 : tensor<20x1x12xf32>
+}
+
+// CHECK-LABEL: func.func @unpack_non_adjacent_inner_dims(
+// CHECK-SAME: %[[SOURCE:.*]]: tensor<10x1x3x4x2xf32>,
+// CHECK-SAME: %[[DEST:.*]]: tensor<20x1x12xf32>)
+// CHECK: %{{.*}} = linalg.unpack
+// CHECK-SAME: inner_dims_pos = [2, 0]
+// CHECK-SAME: inner_tiles = [4, 2]
+// CHECK-SAME: into %[[DEST]] : tensor<10x1x3x4x2xf32> -> tensor<20x1x12xf32>
+
+// -----
+
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>
>From 565ee3edb13297b56dbd2cc32ddc0cb68aa50fce Mon Sep 17 00:00:00 2001
From: Christopher McGirr <mcgirr at roofline.ai>
Date: Thu, 26 Jun 2025 15:47:55 +0200
Subject: [PATCH 2/3] Update(1) [mlir][linalg] Update pack and unpack
documentation
Signed-off-by: Christopher McGirr <mcgirr at roofline.ai>
---
.../Dialect/Linalg/IR/LinalgRelayoutOps.td | 82 ++++++++++---------
1 file changed, 45 insertions(+), 37 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
index fef1900be62ea..81729fab05127 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
@@ -93,19 +93,20 @@ def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
tensor of rank `n + k` with a tiled and packed layout (maybe with padding)
and optionally transposes the tiled source tensor dimensions.
- `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`.
- - the resulting tiled source dimension maps to an outer dimension of the
- packed tensor in the order the non-tiled dimension appeared in the source
- tensor, i.e. `shape(result)[inner_dims_pos[i]]` is equal to
- `shape(source)[inner_dims_pos[i]] / inner_tiles[i]`.
-
`inner_tiles` (mandatory) specifies `k` tile sizes. These tile sizes
correspond to the least significant ("inner") result tensor dimension sizes,
in the same order. Tile sizes can be static or dynamic.
+ `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`.
+ - 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`.
+ - The following relationship for the tiled dimensions holds:
+ `shape(result)[inner_dims_pos[i]] = shape(source)[inner_dims_pos[i]] / inner_tiles[i]`.
+
Example: If `inner_tiles = [16, 32]`, the result tensor has a shape of
`...x16x32`. If `inner_dims_pos = [0, 1]`, the 0th source dimension is tiled
by 16 and the 1st source dimension is tiled by 32. Other source dimensions
@@ -118,17 +119,19 @@ def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
%0 = linalg.pack %source inner_dims_pos = [0, 1] inner_tiles = [8, 32]
into %dest : tensor<128x256xf32> -> tensor<16x8 x 8x32 xf32>
// \ / \ /
- // outer dims inner dims
+ // Outer Dims: 16x8 Inner Dims: 8x32
+
// CHW to CHWhw
%0 = linalg.pack %source inner_dims_pos = [2, 1] inner_tiles = [4, 2]
- into %dest : tensor<1x8x16xf32> -> tensor<1x2x4 x 4x2 xf32>
- // \ / \ /
- // outer dims inner dims
+ into %dest : tensor<3x20x24xf32> -> tensor<3x10x6 x 4x2 xf32>
+ // \ / \ /
+ // Outer Dims: 3x10x6 Inner Dims: 4x2
+
// HCW to HCWhw
%0 = linalg.pack %source inner_dims_pos = [2, 0] inner_tiles = [4, 2]
- into %dest : tensor<20x1x12xf32> -> tensor<10x1x3 x 4x2xf32>
- // \ / \ /
- // Outer Dims: 10x1x3 Inner Dims: 4x2
+ into %dest : tensor<18x3x32xf32> -> tensor<9x3x8 x 4x2 xf32>
+ // \ / \ /
+ // Outer Dims: 9x3x8 Inner Dims: 4x2
```
`outer_dims_perm` (optional) specifies a permutation for the outer
@@ -258,15 +261,6 @@ def Linalg_UnPackOp : Linalg_RelayoutOp<"unpack"> {
The "unpack" operation converts a source tensor of rank `n` with a tiled and
packed layout to a result tensor of rank `n - k`.
- `inner_dims_pos` (mandatory) specifies `k` result tensor dimensions that
- were tiled with the `inner_tiles` to create the packed source tensor. The
- source tensor dimensions can be combined given `inner_dims_pos` as follows:
- the inner tile `shape(source)[n-k+i]` is combined with
- `shape(source)[inner_dims_pos[i]]` where `0 <= i < k` and stored at
- `shape(result)[inner_dims_pos[i]]`. The remaining dimensions are
- `shape(result)[j] = shape(source)[j]` where `0 <= j < n-k` and `j` is not in
- the set of `inner_dims_pos` indices.
-
`inner_tiles` (mandatory) specifies `k` tile sizes. These tile sizes
correspond to the least significant ("inner") source tensor dimension sizes.
The behavior of this op is undefined if:
@@ -276,36 +270,50 @@ def Linalg_UnPackOp : Linalg_RelayoutOp<"unpack"> {
`inner_dims_pos[i]` (assuming that `outer_dims_perm` is not specified)
evenly.
+ `inner_dims_pos` (mandatory) specifies `k` result tensor (i.e. unpacked
+ tensor) dimensions that were tiled with the `inner_tiles` to create the
+ 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]]`.
+ - For `0 <= j < n-k` and `j` not in `inner_dims_pos` the following relationship holds:
+ `shape(result)[j] = shape(source)[j]`.
+
`outer_dims_perm` (optional) specifies a permutation for the outer
dimensions. If specified, it must have `n - k` elements. If specified, this
permutation is applied before combining any dimensions.
- Note, the amount of elements in the source (packed tensor) and the result
- (unpacked) can be unequal, i.e. `SizeOf(source) >= SizeOf(result)`. As
- the unpack operation may drop any padding introduced by the pack operation.
+ Note, the unpack operation may drop any padding introduced by the pack
+ operation and hence the following holds
+ `NumElementsOf(source) >= NumElementsOf(result)`.
Examples:
```mlir
// NCnc to NC:
%0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [8, 32]
- into %dest : tensor<16x8x8x32xf32> -> tensor<128x256xf32>
+ into %dest : tensor<16x8 x 8x32 xf32> -> tensor<128x256xf32>
+ // \ / \ /
+ // Outer Dims: 16x8 Inner Dims: 8x32
// CK to KCck:
%0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
- inner_tiles = [8, 32] into %dest
- : tensor<8x16x8x32xf32> -> tensor<128x256xf32>
+ inner_tiles = [8, 32]
+ into %dest : tensor<8x16 x 8x32 xf32> -> tensor<128x256xf32>
+ // \ / \ /
+ // Outer Dims: 8x16 Inner Dims: 8x32
// CHW to CHWhw:
%0 = linalg.unpack %source inner_dims_pos = [2, 1] inner_tiles = [4, 2]
- into %dest : tensor<1x3x2x4x2xf32> -> tensor<1x5x7xf32>
- // / \
- // Outer Dims: 1x3x2 Inner Dims: 4x2
+ into %dest : tensor<3x10x6 x 4x2 xf32> -> tensor<3x20x24xf32>
+ // \ / \ /
+ // Outer Dims: 3x10x6 Inner Dims: 4x2
+
// HCW to HCWhw
%0 = linalg.unpack %source inner_dims_pos = [2, 0] inner_tiles = [4, 2]
- into %dest : tensor<10x1x3 x 4x2xf32> -> tensor<20x1x12xf32>
- // / \
- // Outer Dims: 10x1x3 Inner Dims: 4x2
+ into %dest : tensor<9x3x8 x 4x2 xf32> -> tensor<18x3x32xf32>
+ // \ / \ /
+ // Outer Dims: 9x3x8 Inner Dims: 4x2
```
}];
let arguments = (ins AnyRankedTensor:$source,
>From a72fcc72f79ccd89e482eca89098f852e9c75c4c Mon Sep 17 00:00:00 2001
From: Christopher McGirr <mcgirr at roofline.ai>
Date: Fri, 27 Jun 2025 09:20:33 +0200
Subject: [PATCH 3/3] Update(2) [mlir][linalg] Update pack and unpack
documentation
Signed-off-by: Christopher McGirr <mcgirr at roofline.ai>
---
.../Dialect/Linalg/IR/LinalgRelayoutOps.td | 5 +--
mlir/test/Dialect/Linalg/invalid.mlir | 5 ++-
mlir/test/Dialect/Linalg/named-ops.mlir | 36 +++++++++++++++++--
3 files changed, 39 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
index 81729fab05127..c384e8b638382 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
+++ b/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]`.
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 0f71000f7c700..17f25a800d17e 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/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 at +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>
diff --git a/mlir/test/Dialect/Linalg/named-ops.mlir b/mlir/test/Dialect/Linalg/named-ops.mlir
index b21b234bc7841..412f40d501154 100644
--- a/mlir/test/Dialect/Linalg/named-ops.mlir
+++ b/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>
More information about the Mlir-commits
mailing list