[Codegen][LLVMCPU] fix tile size selection for consumer unpack ops (#24709)
In the case of consumer pack operations, e.g. `mmt4d -> unpack`
dispatches, the tile sizes set for the unpack operation did not reflect
unpacked destination dimension, but rather the outer source dimensions.
This PR fixes that. Correct tile sizes are necessary for vectorization
and tile size alignment hints in case of scalable vectors.
Also, refactors the `(undo)scaleAndPermutateTilingForPack/UnpackOp`
helpers were a bit confusing, I switched `scaling` and `undoScaling` and
introduced a new helper method for code-sharing.
---------
Signed-off-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index c590d8f..ede7559 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -3163,17 +3163,18 @@
}
/// Transforms tiling sizes from the unpacked domain to the packed domain
-/// for a `PackOp` by scaling inner dimensions and applying outer dimension
-/// permutations.
+/// for a `PackOp` by undoing the scaling for inner dimensions and applying
+/// outer dimension permutations.
///
/// Steps:
/// 1. Divide the tile sizes of inner dimensions by the corresponding inner
/// tile factors. Handles static and scalable sizes but ignores dynamic
/// sizes.
/// 2. Apply the outer dimension permutation, if present.
-static void scaleAndPermutateTilingForPackOp(linalg::PackOp packOp,
- SmallVector<int64_t> &tileSizes,
- SmallVector<bool> &scalableFlags) {
+static void
+undoScaleAndPermutateTilingForPackOp(linalg::PackOp packOp,
+ SmallVector<int64_t> &tileSizes,
+ SmallVector<bool> &scalableFlags) {
SmallVector<int64_t> innerTiles(packOp.getStaticInnerTiles());
SmallVector<bool> innerTileScalableFlags(innerTiles.size(), false);
// Infer scalable tile sizes and flags if present.
@@ -3204,36 +3205,35 @@
}
}
-/// Transforms tiling sizes from the packed domain back to the unpacked
-/// domain for a `PackOp` by undoing the scaling of inner dimensions and
-/// reversing outer dimension permutations.
-///
-/// Steps:
+/// Helper function to map `linalg.pack/unpack` operation tile sizes from the
+/// packed dimension to the unpacked dimension.
/// 1. Undo the outer dimension permutation, if present, by applying the
/// inverted permutation.
/// 2. Multiply the inner dimension tile sizes by the corresponding inner
/// tile factors. Handles static and scalable tile sizes but ignores dynamic
/// sizes.
-static void
-undoScaleAndPermutateTilingForPackOp(linalg::PackOp packOp,
- SmallVector<int64_t> &tileSizes,
- SmallVector<bool> &scalableFlags) {
- SmallVector<int64_t> innerTiles(packOp.getStaticInnerTiles());
+template <typename PackUnpackType>
+static void undoPermutationAndScaleTileSizes(PackUnpackType op,
+ SmallVector<int64_t> &tileSizes,
+ SmallVector<bool> &scalableFlags) {
+ SmallVector<int64_t> innerTiles(op.getStaticInnerTiles());
SmallVector<bool> innerTileScalableFlags(innerTiles.size(), false);
// Infer scalable tile sizes and flags if present.
if (auto sizesAndScalableFlags =
- getScalableTileSizesAndFlags(packOp.getMixedTiles())) {
+ getScalableTileSizesAndFlags(op.getMixedTiles())) {
innerTiles = sizesAndScalableFlags->first;
innerTileScalableFlags = sizesAndScalableFlags->second;
}
- ArrayRef<int64_t> innerDimPos = packOp.getInnerDimsPos();
- ArrayRef<int64_t> outerDimsPerm = packOp.getOuterDimsPerm();
+ ArrayRef<int64_t> innerDimPos = op.getInnerDimsPos();
+ ArrayRef<int64_t> outerDimsPerm = op.getOuterDimsPerm();
// First undo dimension permutation if present.
if (!outerDimsPerm.empty()) {
auto invertedPerm = invertPermutationVector(outerDimsPerm);
applyPermutationToVector(tileSizes, invertedPerm);
applyPermutationToVector(scalableFlags, invertedPerm);
}
+ assert(llvm::none_of(scalableFlags, [](bool scalable) { return scalable; }) &&
+ "Tile sizes for outer dims should not be scalable!");
// Then unscale tile sizes by multiplying the inner tile sizes and setting the
// corresponding scalable flags to true.
for (auto [pos, size, scalable] :
@@ -3247,6 +3247,27 @@
}
}
+/// Transforms tiling sizes from the packed domain back to the unpacked
+/// domain for a `PackOp` by scaling the inner dimensions and
+/// reversing outer dimension permutations.
+static void scaleAndPermutateTilingForPackOp(linalg::PackOp packOp,
+ SmallVector<int64_t> &tileSizes,
+ SmallVector<bool> &scalableFlags) {
+ undoPermutationAndScaleTileSizes<linalg::PackOp>(packOp, tileSizes,
+ scalableFlags);
+}
+
+/// Transforms tiling sizes from the packed domain back to the unpacked
+/// domain for a `UnPackOp` by scaling the inner dimensions and
+/// reversing outer dimension permutations.
+static void
+undoScaleAndPermutateTilingForUnpackOp(linalg::UnPackOp unpackOp,
+ SmallVector<int64_t> &tileSizes,
+ SmallVector<bool> &scalableFlags) {
+ undoPermutationAndScaleTileSizes<linalg::UnPackOp>(unpackOp, tileSizes,
+ scalableFlags);
+}
+
/// A helper class that propagates and sets lowering configurations for multiple
/// compute operations.
///
@@ -3324,7 +3345,7 @@
/// For a Matmul RHS Pack op with `outer_dims_perm = [1, 0]` and `inner_tiles
/// = [16, 1]`, `getPackVectorTileSize` initially returns `[1, 1]`. Since the
/// `MultiLoweringConfigGenerator` propagates tiling of the outer (unpacked)
- /// dimensions, `undoScaleAndPermutateTilingForPackOp` translates the
+ /// dimensions, `scaleAndPermutateTilingForPackOp` translates the
/// tile sizes from `[1, 1]` to `[1, 16]`.
///
/// As a result, the Pack op expects its producer (potentially the root op) to
@@ -3427,11 +3448,11 @@
// `MultiLoweringConfigGenerator` propagates tiling on the unpacked
// dimensions, while `rootLoweringConfig` defines tiling on the packed
// inner dimensions. Therefore, use
- // `undoScaleAndPermutateTilingForPackOp` to translate tiling information
+ // `scaleAndPermutateTilingForPackOp` to translate tiling information
// from the packed back to the unpacked dimensions.
if (auto packOp = dyn_cast<linalg::PackOp>(rootOperation);
packOp && !sizes.empty()) {
- undoScaleAndPermutateTilingForPackOp(packOp, sizes, flags);
+ scaleAndPermutateTilingForPackOp(packOp, sizes, flags);
}
// Map the tiling information from the op-level local dimension indices
@@ -3723,9 +3744,20 @@
// `MultiLoweringConfigGenerator` propagates tiling on the
// unpacked dimensions, while for a pack operation, `LoweringConfig`
// defines tiling on the packed inner dimensions. Therefore, use
- // `scaleAndPermutateTilingForPackOp` to translate the tiling
+ // `undoScaleAndPermutateTilingForPackOp` to translate the tiling
// information from the unpacked to the packed dimensions.
- scaleAndPermutateTilingForPackOp(packOp, tileSizes, scalableFlags);
+ undoScaleAndPermutateTilingForPackOp(packOp, tileSizes, scalableFlags);
+ } else if (auto unpackOp = dyn_cast<linalg::UnPackOp>(op);
+ unpackOp &&
+ level == IREE::CPU::TilingLevel::VectorCommonParallelTiles &&
+ unpackOp.getSource().getDefiningOp<linalg::LinalgOp>()) {
+ // The `IterationDimTracker` ties an unpack's destination loops only to
+ // the *outer* dims of its packed source operand, so `tileSizes`
+ // holds the source's outer tiling for consumer unpack operations. Map
+ // these onto the destination unpacked domain by scaling and permuting
+ // with the inner tile sizes.
+ undoScaleAndPermutateTilingForUnpackOp(unpackOp, tileSizes,
+ scalableFlags);
}
// Append tiling info.
@@ -3748,7 +3780,7 @@
// Invert the Pack op's `outer_dims_perm` on `vecTileSizes` and
// `scalableFlags`, then multiply `vecTileSizes` by the Pack op's
// `inner_tiles` and set the corresponding `scalableFlags`.
- undoScaleAndPermutateTilingForPackOp(packOp, vecTileSizes, scalableFlags);
+ scaleAndPermutateTilingForPackOp(packOp, vecTileSizes, scalableFlags);
return {vecTileSizes, scalableFlags};
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir
index 7f79c40..9a00d51 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir
@@ -456,7 +456,7 @@
}
// CHECK-DAG: #[[$CONFIG0:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 1, 16, [16]]>
// CHECK-DAG: #[[$CONFIG1:.+]] = #iree_cpu.lowering_config<distribution = [1, 1, 0, 0, 0, 0], vector_common_parallel = [1, 1, 0, 16, [16], 0], vector_reduction = [0, 0, 1, 0, 0, 1]>
-// CHECK-DAG: #[[$CONFIG2:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 1]>
+// CHECK-DAG: #[[$CONFIG2:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [16, [16]]>
// CHECK-LABEL: func.func @mmt4d_generic_unpack_pack(
// CHECK: linalg.fill
// CHECK-SAME: {lowering_config = #[[$CONFIG0]]}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
index a551c3c..3317074 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
@@ -1744,7 +1744,7 @@
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK: #[[MMT_FILL_CONFIG:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 1, 16, 16]>
// CHECK: #[[MMT_MMT4D_CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [1, 1, 0, 0, 0, 0], vector_common_parallel = [1, 1, 0, 16, 16, 0], vector_reduction = [0, 0, 1, 0, 0, 1]>
-// CHECK: #[[MMT_UNPACK_CONFIG:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 1]>
+// CHECK: #[[MMT_UNPACK_CONFIG:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [16, 16]>
func.func @mmt4d_generic_unpack_pack(%arg0: tensor<5x4096x16x1xf16>, %arg1: tensor<640x4096x16x1xf16>) -> tensor<5x10240x16x1xf16> attributes {hal.executable.target = #executable_target_embedded_elf_x86_64} {
%cst = arith.constant 0.000000e+00 : f16
%cst_0 = arith.constant 0.000000e+00 : f32
