[LLVMCPU] thread inner tile alignment hint for vector-level tiling (#24698)
Introduces a helper function that will generate the inner tile alignment hint control function introduced in https://github.com/llvm/llvm-project/pull/204007.
The control function is passed to the SCF tiling driver and is invoked per-op before tiling/fusion. It merely forwards the inner tile alignment hints set during tile size selection, added in the stack of https://github.com/iree-org/iree/pull/24806.
Also utilizes this in the `TileAndFuseProducerConsumer` pass.
Assisted-by: Claude Code
---------
Signed-off-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp b/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp
index aef1767..d62fabf 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp
@@ -10,6 +10,7 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "mlir/Analysis/TopologicalSortUtils.h"
@@ -20,7 +21,9 @@
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Dominance.h"
+#include "mlir/Interfaces/TilingInterface.h"
#include <cassert>
@@ -28,6 +31,26 @@
namespace mlir::iree_compiler {
+scf::InnerTileAlignmentFnTy
+makeInnerTileAlignmentFn(IREE::CPU::TilingLevel tilingLevel) {
+ return [tilingLevel](
+ TilingInterface op, ArrayRef<OpFoldResult>,
+ ArrayRef<Operation *>) -> SmallVector<mlir::InnerTileAlignment> {
+ // Tile-size selection already recorded this op's per-dimension inner-tile
+ // alignments per tiling level. We simply forward those along.
+ auto hint =
+ IREE::CPU::InnerTileAlignmentsAttr::getFromOp(op.getOperation());
+ if (!hint) {
+ return {};
+ }
+ auto alignments = hint.getAlignments().getAs<DenseI64ArrayAttr>(
+ IREE::CPU::getTilingLevelName(tilingLevel));
+ return alignments
+ ? mlir::convertInnerTileAlignments(alignments.asArrayRef())
+ : SmallVector<mlir::InnerTileAlignment>{};
+ };
+}
+
void fuseProducersOfSlices(RewriterBase &rewriter,
std::queue<Operation *> &worklist,
scf::SCFTileAndFuseOptions &options,
@@ -54,7 +77,9 @@
// values produced by operations that implement the `TilingInterface`.
// Add these operations to the worklist.
std::optional<scf::SCFFuseProducerOfSliceResult> fusedResult =
- scf::tileAndFuseProducerOfSlice(rewriter, candidateSlice, loops);
+ scf::tileAndFuseProducerOfSlice(
+ rewriter, candidateSlice, loops,
+ options.tilingOptions.innerTileAlignmentFn);
if (!fusedResult) {
continue;
}
@@ -109,10 +134,11 @@
};
} // namespace
-FailureOr<std::queue<Operation *>>
-fuseConsumersIntoForall(RewriterBase &rewriter, ArrayRef<Operation *> tiledOps,
- MutableArrayRef<LoopLikeOpInterface> loops,
- std::function<bool(Operation *)> filterFn) {
+FailureOr<std::queue<Operation *>> fuseConsumersIntoForall(
+ RewriterBase &rewriter, ArrayRef<Operation *> tiledOps,
+ MutableArrayRef<LoopLikeOpInterface> loops,
+ std::function<bool(Operation *)> filterFn,
+ const scf::InnerTileAlignmentFnTy &innerTileAlignmentFn) {
// Collect the candidate slices which can be potential consumers that can be
// fused. Keep them in a vector reverse-sorted by dominance: the candidate
// dominating others comes last (so it can be cheaply popped from the vector).
@@ -213,7 +239,8 @@
ConsumerFusionQueueEntry entry = candidates.pop_back_val();
FailureOr<scf::SCFFuseConsumerOfSliceResult> fusedResult =
- mlir::scf::tileAndFuseConsumer(rewriter, entry.fusableUser, loops);
+ mlir::scf::tileAndFuseConsumer(rewriter, entry.fusableUser, loops,
+ innerTileAlignmentFn);
if (failed(fusedResult)) {
return failure();
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h b/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h
index bae549c..ebd18e2 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h
+++ b/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h
@@ -7,6 +7,7 @@
#ifndef IREE_COMPILER_CODEGEN_COMMON_TILEANDFUSEUTILS_H_
#define IREE_COMPILER_CODEGEN_COMMON_TILEANDFUSEUTILS_H_
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h"
#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
@@ -31,15 +32,21 @@
void collectTiledAndFusedOps(Operation *rootOp,
llvm::SmallDenseSet<Operation *> &result);
+/// Returns an inner tile alignment control function for scalable tiling at
+/// `tilingLevel`. It forwards the per-dimension `InnerTileAlignment` hints that
+/// tile-size selection set on `linalg.pack`/`linalg.unpack` ops.
+scf::InnerTileAlignmentFnTy
+makeInnerTileAlignmentFn(IREE::CPU::TilingLevel tilingLevel);
+
/// Fuse all consumers of the given `tiledOps` into the surrounding `scf.forall`
/// unless specified otherwise by `filterFn`. Returns a list of new
/// `tensor.extract_slice` ops with new fusion opportunities.
FailureOr<std::queue<Operation *>> fuseConsumersIntoForall(
RewriterBase &rewriter, ArrayRef<Operation *> tiledOps,
MutableArrayRef<LoopLikeOpInterface> loops,
- std::function<bool(Operation *)> filterFn = [](Operation *) {
- return true;
- });
+ std::function<bool(Operation *)> filterFn =
+ [](Operation *) { return true; },
+ const scf::InnerTileAlignmentFnTy &innerTileAlignmentFn = nullptr);
/// Apply a tile and fuse transformation to all payload ops and store both the
/// tiled operation as well as the created tile loops.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuseProducerConsumer.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuseProducerConsumer.cpp
index 1295c8b..1ce9113 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuseProducerConsumer.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuseProducerConsumer.cpp
@@ -187,9 +187,16 @@
tileSizes.resize(numLoops, 0);
tileScalableFlags.resize(numLoops, false);
+ // Lets a scalable pack/unpack's tiled outer dims fold to a static single
+ // inner tile, using the alignment hints tile-size selection recorded on the
+ // op.
+ scf::InnerTileAlignmentFnTy innerTileAlignmentFn =
+ makeInnerTileAlignmentFn(tilingLevel);
+
scf::SCFTilingOptions tilingOptions;
setSCFTileSizes(tilingOptions, rootOp, std::move(tileSizes),
std::move(tileScalableFlags));
+ tilingOptions.setInnerTileAlignmentFn(innerTileAlignmentFn);
// onlyFuseProducerInputOperands implies reduction tiling.
if (!onlyFuseProducerInputOperands) {
@@ -262,7 +269,8 @@
[&tiledAndFusedOps, &unfusableOps](Operation *op) {
return tiledAndFusedOps.contains(op) &&
!unfusableOps.contains(op);
- });
+ },
+ innerTileAlignmentFn);
if (failed(newFusionOpportunities)) {
LDBG() << "failed to fuse consumers, skip";
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile_and_fuse_producer_consumer_anchoring_root_op.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile_and_fuse_producer_consumer_anchoring_root_op.mlir
index 220f9fa..bce3301 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile_and_fuse_producer_consumer_anchoring_root_op.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile_and_fuse_producer_consumer_anchoring_root_op.mlir
@@ -416,3 +416,176 @@
// CHECK: linalg.generic
// CHECK: scf.forall.in_parallel
// CHECK: linalg.pack
+
+// -----
+
+// Scalable linalg.pack with an elementwise-add producer.
+// The `Equal` hint folds the scalable inner tile's outer dim to a static 1.
+
+#config = #iree_cpu.lowering_config<distribution = [64, 64], vector_common_parallel = [[8], 1]>
+#config1 = #iree_cpu.lowering_config<vector_common_parallel = [1, 1]>
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @scalable_pack_with_producer(%arg0: tensor<384x512xf32>, %arg1: tensor<384x512xf32>, %arg2 : tensor<384x512xf32>) -> tensor<?x512x?x1xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %vscale, %c8 : index
+ %generic = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]}
+ ins(%arg0, %arg1 : tensor<384x512xf32>, tensor<384x512xf32>)
+ outs(%arg2 : tensor<384x512xf32>) attrs = {lowering_config = #config} {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %3 = arith.addf %in, %in_0 : f32
+ linalg.yield %3 : f32
+ } -> tensor<384x512xf32>
+ %mouter = affine.apply affine_map<()[s0] -> (384 ceildiv s0)>()[%c8_vscale]
+ %2 = tensor.empty(%mouter, %c8_vscale) : tensor<?x512x?x1xf32>
+ %pack = linalg.pack %generic padding_value(%cst : f32) outer_dims_perm = [0, 1] inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, 1] into %2
+ {inner_tile_alignments = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Equal, Unknown]>,
+ lowering_config = #config1} : tensor<384x512xf32> -> tensor<?x512x?x1xf32>
+ return %pack : tensor<?x512x?x1xf32>
+}
+// CHECK-LABEL: func.func @scalable_pack_with_producer
+// CHECK: %[[VSCALE:.+]] = vector.vscale
+// CHECK: %[[C8VS:.+]] = arith.muli %[[VSCALE]], %{{.+}} : index
+// CHECK: scf.forall {{.*}} = (0, 0) to (384, 512) step (%[[C8VS]], 1)
+// CHECK: %[[GEN:.+]] = linalg.generic
+// CHECK: linalg.pack %[[GEN]]
+// CHECK-SAME: inner_tiles = [%[[C8VS]], 1]
+// The outermost dimension has to equal 1 to confirm the Equal hint is taken into account.
+// CHECK-SAME: -> tensor<1x1x?x1xf32>
+// CHECK: scf.forall.in_parallel
+
+// -----
+
+// Scalable linalg.unpack fused as the producer of an elementwise-add consumer.
+// The `Equal` hint folds the scalable inner tile's outer dim to a static 1.
+
+#config = #iree_cpu.lowering_config<vector_common_parallel = [7, [8]]>
+#config1 = #iree_cpu.lowering_config<distribution = [84, 64], vector_common_parallel = [7, [8]]>
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @scalable_unpack_with_consumer(%arg0: tensor<12x?x7x?xf32>, %arg1: tensor<80x320xf32>) -> tensor<80x320xf32> {
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %vscale, %c8 : index
+ %0 = tensor.empty() : tensor<80x320xf32>
+ %unpack = linalg.unpack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1]
+ inner_tiles = [7, %c8_vscale] into %0
+ {inner_tile_alignments = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Unknown, Equal]>,
+ lowering_config = #config} : tensor<12x?x7x?xf32> -> tensor<80x320xf32>
+ %generic = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]}
+ ins(%arg1, %unpack : tensor<80x320xf32>, tensor<80x320xf32>)
+ outs(%0 : tensor<80x320xf32>) attrs = {lowering_config = #config1} {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.addf %in, %in_0 : f32
+ linalg.yield %2 : f32
+ } -> tensor<80x320xf32>
+ return %generic : tensor<80x320xf32>
+}
+// CHECK-LABEL: func.func @scalable_unpack_with_consumer
+// CHECK: %[[VSCALE:.+]] = vector.vscale
+// CHECK: %[[C8VS:.+]] = arith.muli %[[VSCALE]], %{{.+}} : index
+// CHECK: scf.forall {{.*}} = (0, 0) to (80, 320) step (7, %[[C8VS]])
+// CHECK: %[[UNPACK:.+]] = linalg.unpack
+// CHECK-SAME: inner_tiles = [7, %[[C8VS]]]
+// The outermost dimensions have to equal 1 to confirm the Equal hint is taken into account.
+// CHECK-SAME: tensor<1x1x7x?xf32> -> tensor<?x?xf32>
+// CHECK: linalg.generic
+// CHECK-SAME: ins(%{{.*}}, %[[UNPACK]]
+// CHECK: scf.forall.in_parallel
+
+// -----
+
+// Scalable linalg.mmt4d (N0 = 8*vscale) with a linalg.unpack consumer fused into
+// the mmt4d's tiling. The `Equal` hint folds the scalable inner tile's outer dim to a static 1.
+
+#config = #iree_cpu.lowering_config<distribution = [5, 8, 0, 0, 0, 0], vector_common_parallel = [1, 1, 0, 7, [8], 0]>
+#config1 = #iree_cpu.lowering_config<vector_common_parallel = [7, [8]]>
+func.func @scalable_mmt4d_with_unpack_consumer(%lhs: tensor<55x512x7x1xf32>, %rhs: tensor<?x512x?x1xf32>) -> tensor<385x?xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %vscale, %c8 : index
+ %nouter = tensor.dim %rhs, %c0 : tensor<?x512x?x1xf32>
+ %acc = tensor.empty(%nouter, %c8_vscale) : tensor<55x?x7x?xf32>
+ %fill = linalg.fill ins(%cst : f32) outs(%acc : tensor<55x?x7x?xf32>) -> tensor<55x?x7x?xf32>
+ %mmt4d = linalg.mmt4d {lowering_config = #config} ins(%lhs, %rhs : tensor<55x512x7x1xf32>, tensor<?x512x?x1xf32>) outs(%fill : tensor<55x?x7x?xf32>) -> tensor<55x?x7x?xf32>
+ %n = arith.muli %nouter, %c8_vscale : index
+ %out = tensor.empty(%n) : tensor<385x?xf32>
+ %unpack = linalg.unpack %mmt4d outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [7, %c8_vscale] into %out {inner_tile_alignments = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Unknown, Equal]>, lowering_config = #config1} : tensor<55x?x7x?xf32> -> tensor<385x?xf32>
+ return %unpack : tensor<385x?xf32>
+}
+// CHECK-LABEL: func.func @scalable_mmt4d_with_unpack_consumer
+// CHECK: %[[VSCALE:.+]] = vector.vscale
+// CHECK: %[[C8VS:.+]] = arith.muli %[[VSCALE]], %{{.+}} : index
+// CHECK: scf.forall {{.*}} step (1, 1, %[[C8VS]])
+// CHECK: %[[FILL:.+]] = linalg.fill
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: outs(%[[FILL]]
+// CHECK: linalg.unpack %[[MMT4D]]
+// CHECK-SAME: inner_tiles = [7, %[[C8VS]]]
+// The outermost dimensions have to equal 1 to confirm the Equal hint is taken into account.
+// CHECK-SAME: tensor<1x1x7x?xf32> -> tensor<7x?xf32>
+// CHECK: scf.forall.in_parallel
+
+// -----
+
+// Scalable linalg.mmt4d (M0 = N0 = 8*vscale) with fused generic, unpack and pack
+// consumers. Each scalable inner tile dim is expected to fold to a static 1.
+
+#config = #iree_cpu.lowering_config<distribution = [5, 7, 0, 0, 0, 0], vector_common_parallel = [1, 1, 0, [8], [8], 0]>
+#config1 = #iree_cpu.lowering_config<vector_common_parallel = [1, 1, [8], [8]]>
+#config2 = #iree_cpu.lowering_config<vector_common_parallel = [[8], [8]]>
+#config3 = #iree_cpu.lowering_config<vector_common_parallel = [[8], 1]>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func.func @scalable_mmt4d_generic_unpack_pack(%lhs: tensor<5x512x?x1xf32>, %rhs: tensor<7x512x?x1xf32>) -> tensor<?x?x?x1xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c5 = arith.constant 5 : index
+ %c7 = arith.constant 7 : index
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %vscale, %c8 : index
+ %acc = tensor.empty(%c8_vscale, %c8_vscale) : tensor<5x7x?x?xf32>
+ %fill = linalg.fill ins(%cst : f32) outs(%acc : tensor<5x7x?x?xf32>) -> tensor<5x7x?x?xf32>
+ %mmt4d = linalg.mmt4d {lowering_config = #config}
+ ins(%lhs, %rhs : tensor<5x512x?x1xf32>, tensor<7x512x?x1xf32>)
+ outs(%fill : tensor<5x7x?x?xf32>) -> tensor<5x7x?x?xf32>
+ %genout = tensor.empty(%c8_vscale, %c8_vscale) : tensor<5x7x?x?xf32>
+ %gen = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+ ins(%mmt4d : tensor<5x7x?x?xf32>) outs(%genout : tensor<5x7x?x?xf32>)
+ attrs = {lowering_config = #config1} {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.maximumf %in, %cst : f32
+ linalg.yield %v : f32
+ } -> tensor<5x7x?x?xf32>
+ %m = arith.muli %c5, %c8_vscale : index
+ %n = arith.muli %c7, %c8_vscale : index
+ %uout = tensor.empty(%m, %n) : tensor<?x?xf32>
+ %unpack = linalg.unpack %gen outer_dims_perm = [0, 1] inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c8_vscale] into %uout
+ {inner_tile_alignments = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Equal, Equal]>,
+ lowering_config = #config2} : tensor<5x7x?x?xf32> -> tensor<?x?xf32>
+ %m_outer = affine.apply affine_map<()[s0, s1] -> (s0 ceildiv s1)>()[%m, %c8_vscale]
+ %pout = tensor.empty(%m_outer, %n, %c8_vscale) : tensor<?x?x?x1xf32>
+ %pack = linalg.pack %unpack padding_value(%cst : f32) outer_dims_perm = [0, 1] inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, 1] into %pout
+ {inner_tile_alignments = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Equal, Multiple]>,
+ lowering_config = #config3} : tensor<?x?xf32> -> tensor<?x?x?x1xf32>
+ return %pack : tensor<?x?x?x1xf32>
+}
+// CHECK-LABEL: func.func @scalable_mmt4d_generic_unpack_pack
+// CHECK: %[[VSCALE:.+]] = vector.vscale
+// CHECK: %[[C8VS:.+]] = arith.muli %[[VSCALE]], %{{.+}} : index
+// CHECK: scf.forall {{.*}} step (1, 1, %[[C8VS]], %[[C8VS]])
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK: %[[GEN:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[MMT4D]]
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[GEN]]
+// CHECK-SAME: inner_tiles = [%[[C8VS]], %[[C8VS]]]
+// CHECK-SAME: tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+// CHECK: linalg.pack %[[UNPACK]]
+// CHECK-SAME: inner_tiles = [%[[C8VS]], 1]
+// The outermost dimension has to equal 1 to confirm the Equal hint is taken into account.
+// CHECK-SAME: tensor<?x?xf32> -> tensor<1x?x?x1xf32>
+// CHECK: scf.forall.in_parallel
