[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