[Mlir-commits] [mlir] [mlir][tensor] Fold producer linalg transpose with consumer tensor pack (PR #75658)
lorenzo chelini
llvmlistbot at llvm.org
Thu Dec 28 06:58:35 PST 2023
================
@@ -142,11 +143,45 @@ struct FoldProducerPackWithConsumerLinalgTransposeOp
return success();
}
};
+
+/// Fold 'transpose' -> 'pack' into 'pack' since 'pack' already has transpose
+/// semantics.
+struct FoldConsumerPackWithProducerLinalgTransposeOp
----------------
chelini wrote:
We can use the transpose permutation only if we do not move the `inner_dims_pos` loops in the transpose operation. The semantics of `outer_dim_perms` is an interchange on the tile loops *after* the packing. Consider the following example:
```
func.func @main(%arg0: tensor<56x57x1x64xf32>) -> tensor<64x57x1x7x8xf32> {
%0 = tensor.empty() : tensor<64x57x1x56xf32>
%transposed = linalg.transpose
ins(%arg0 : tensor<56x57x1x64xf32>)
outs(%0 : tensor<64x57x1x56xf32>)
permutation = [3, 1, 2, 0] // here I swap dim 3 with 0.
%1 = tensor.empty() : tensor<64x57x1x7x8xf32>
%pack = tensor.pack %transposed
inner_dims_pos = [3] // here I want to pack dimension 3 (the 56 not the 64).
inner_tiles = [8]
into %1 : tensor<64x57x1x56xf32> -> tensor<64x57x1x7x8xf32>
return %pack : tensor<64x57x1x7x8xf32>
}
```
In the example above, if we fold the transpose, we are packing dimension 64 in the input tensor `arg0`, but we want to pack the 56 dimension; this is because `outer_dims_perm` is an interchange on the tile loops after tiling.
https://github.com/llvm/llvm-project/pull/75658
More information about the Mlir-commits
mailing list