[Mlir-commits] [llvm] [mlir] [mlir][linalg/scf/transform] scalable tiling and fusion for pack/unpack ops (PR #204007)
Ege Beysel
llvmlistbot at llvm.org
Fri Jul 3 02:42:52 PDT 2026
https://github.com/egebeysel updated https://github.com/llvm/llvm-project/pull/204007
>From c456cae1e6861224d54b901f740293935655e01e Mon Sep 17 00:00:00 2001
From: Ege Beysel <beyselege at gmail.com>
Date: Wed, 10 Jun 2026 17:07:05 +0000
Subject: [PATCH 1/4] feat(mlir/TilingInterface): add InnerTileAlignment hint
as overloaded tiling methods
Add an optional, caller-supplied per-dimension InnerTileAlignment
{Unknown,Multiple,Equal} hint describing how each loop tile size relates to a
linalg.pack/linalg.unpack inner tile, so the pack/unpack TilingInterface impls
need not re-derive it from the materialized IR (which fails for scalable inner
tiles). The caller performed the tiling and knows both the tile sizes and the
op's inner tiles.
Rather than changing the existing method signatures, the hint is added as
same-name OVERLOADS of getTiledImplementation, generateResultTileValue,
getTiledImplementationFromOperandTiles and getIterationDomainTileFromOperandTiles.
The hint-bearing overloads default to forwarding to the hint-less ones, so only
linalg.pack/unpack consult them; every other implementer inherits the defaulted
forwarder and need not reason about the hint. This is a caller assertion and is
always honored by pack/unpack; when the relationship is statically decidable,
it is asserted that the hint agrees with them.
Downstream integration: because the hint methods are overloads of existing
interface methods, an out-of-tree TilingInterface external model that overrides
only a hint-less overload hides the inherited defaulted hint-bearing overload
(C++ name hiding) and fails to compile when the interface forwards to it. Such
implementers must add a `using Base::<method>;` declaration (or define the
overload explicitly to forward) for each affected method; all in-tree
implementers are updated here.
Signed-off-by: Ege Beysel <beyselege at gmail.com>
---
.../include/mlir/Interfaces/TilingInterface.h | 33 +++
.../mlir/Interfaces/TilingInterface.td | 104 ++++++++
mlir/lib/Dialect/Linalg/IR/CMakeLists.txt | 1 +
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 36 +++
.../Linalg/Transforms/TilingInterfaceImpl.cpp | 232 +++++++++++++++---
.../Tensor/IR/TensorTilingInterfaceImpl.cpp | 5 +
mlir/lib/Interfaces/TilingInterface.cpp | 22 ++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 6 +
.../llvm-project-overlay/mlir/BUILD.bazel | 1 +
9 files changed, 408 insertions(+), 32 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/TilingInterface.h b/mlir/include/mlir/Interfaces/TilingInterface.h
index 8693cbea7f0b0..8b685a30868e8 100644
--- a/mlir/include/mlir/Interfaces/TilingInterface.h
+++ b/mlir/include/mlir/Interfaces/TilingInterface.h
@@ -66,6 +66,39 @@ struct MergeResult {
SmallVector<Value> replacements;
};
+/// Per-dimension alignment of a loop tile size to a `linalg.pack` /
+/// `linalg.unpack` inner tile size, supplied by the caller (which performed the
+/// tiling and knows both the tile sizes and the inner tiles) so that
+/// pack/unpack TilingInterface implementations need not re-derive it from the
+/// materialized IR. An absent entry (or `Unknown`) means "no information": the
+/// implementation must fall back to its prior behavior for that dimension.
+/// - `Multiple`: the loop tile size is an integer multiple of the pack/unpack
+/// inner tile.
+/// - `Equal`: the loop tile size equals the pack/unpack inner tile size.
+enum class InnerTileAlignment : int64_t { Unknown = 0, Multiple, Equal };
+
+/// Returns true iff `value` is a valid `InnerTileAlignment` enumerator.
+inline bool isValidInnerTileAlignment(int64_t value) {
+ switch (static_cast<InnerTileAlignment>(value)) {
+ case InnerTileAlignment::Unknown:
+ case InnerTileAlignment::Multiple:
+ case InnerTileAlignment::Equal:
+ return true;
+ }
+ return false;
+}
+
+/// Verifies that every entry of a raw `inner_tile_alignments` integer array is
+/// a valid `InnerTileAlignment`, emitting the standard op error on `op`
+/// otherwise.
+LogicalResult verifyInnerTileAlignments(Operation *op,
+ ArrayRef<int64_t> alignments);
+
+/// Maps a validated `inner_tile_alignments` integer array onto the
+/// per-dimension `InnerTileAlignment` hints consumed by the tiling driver.
+SmallVector<InnerTileAlignment>
+convertInnerTileAlignments(ArrayRef<int64_t> alignments);
+
} // namespace mlir
/// Include the ODS generated interface header files.
diff --git a/mlir/include/mlir/Interfaces/TilingInterface.td b/mlir/include/mlir/Interfaces/TilingInterface.td
index e14ea926f4d00..5845545ce686a 100644
--- a/mlir/include/mlir/Interfaces/TilingInterface.td
+++ b/mlir/include/mlir/Interfaces/TilingInterface.td
@@ -59,6 +59,16 @@ def TilingInterface : OpInterface<"TilingInterface"> {
description below):
- `getTiledImplementationFromOperandTiles`
- `getIterationDomainTileFromOperandTiles`.
+
+ `getTiledImplementation`, `generateResultTileValue`,
+ `getTiledImplementationFromOperandTiles` and
+ `getIterationDomainTileFromOperandTiles` each have a second, overloaded form
+ taking an extra `ArrayRef<InnerTileAlignment>` caller hint (see
+ `InnerTileAlignment` in TilingInterface.h). The hint-bearing overload
+ defaults to forwarding to the hint-less one, so only ops that opt in consult
+ it. Because the two forms share a source name, an implementer that defines
+ only the hint-less overload hides the inherited defaulted one (C++ name
+ hiding) and fails to compile when the interface forwards to it.
}];
let cppNamespace = "::mlir";
let methods = [
@@ -115,6 +125,28 @@ def TilingInterface : OpInterface<"TilingInterface"> {
return {};
}]
>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Variant of `getTiledImplementation` that additionally takes a
+ per-iteration-domain `InnerTileAlignment` hint (see
+ `InnerTileAlignment`) asserting how each loop tile size relates to a
+ pack/unpack inner tile. The hint is consulted only by pack/unpack
+ implementations and may be empty (all `Unknown`). The default
+ implementation ignores it and forwards to the hint-less overload.
+ }],
+ /*retType=*/"::mlir::FailureOr<::mlir::TilingResult>",
+ /*methodName=*/"getTiledImplementation",
+ /*args=*/(ins
+ "::mlir::OpBuilder &":$b,
+ "::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets,
+ "::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes,
+ "::mlir::ArrayRef<::mlir::InnerTileAlignment> ":$innerTileAlignments),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::cast<::mlir::TilingInterface>($_op.getOperation())
+ .getTiledImplementation(b, offsets, sizes);
+ }]
+ >,
InterfaceMethod<
/*desc=*/[{
Method to return the position of the result tile computed by the
@@ -199,6 +231,28 @@ def TilingInterface : OpInterface<"TilingInterface"> {
return failure();
}]
>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Variant of `generateResultTileValue` that additionally takes a
+ per-iteration-domain `InnerTileAlignment` hint (see
+ `InnerTileAlignment`). The hint is consulted only by pack/unpack
+ implementations and may be empty (all `Unknown`). The default
+ implementation ignores it and forwards to the hint-less overload.
+ }],
+ /*retType=*/"::mlir::FailureOr<::mlir::TilingResult>",
+ /*methodName=*/"generateResultTileValue",
+ /*args=*/(ins
+ "::mlir::OpBuilder &":$b,
+ "unsigned":$resultNumber,
+ "::mlir::ArrayRef<::mlir::OpFoldResult>":$offsets,
+ "::mlir::ArrayRef<::mlir::OpFoldResult>":$sizes,
+ "::mlir::ArrayRef<::mlir::InnerTileAlignment>":$innerTileAlignments),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::cast<::mlir::TilingInterface>($_op.getOperation())
+ .generateResultTileValue(b, resultNumber, offsets, sizes);
+ }]
+ >,
InterfaceMethod<
/*desc=*/[{
Method to generate the tiled implementation of an operation that uses
@@ -229,6 +283,29 @@ def TilingInterface : OpInterface<"TilingInterface"> {
return failure();
}]
>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Variant of `getTiledImplementationFromOperandTiles` that additionally
+ takes a per-iteration-domain `InnerTileAlignment` hint (see
+ `InnerTileAlignment`). The hint is consulted only by pack/unpack
+ implementations and may be empty (all `Unknown`). The default
+ implementation ignores it and forwards to the hint-less overload.
+ }],
+ /*retType=*/"::mlir::FailureOr<::mlir::TilingResult>",
+ /*methodName=*/"getTiledImplementationFromOperandTiles",
+ /*args=*/(ins
+ "::mlir::OpBuilder &":$b,
+ "::mlir::ArrayRef<unsigned>":$operandNumbers,
+ "::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>>":$allOffsets,
+ "::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>>":$allSizes,
+ "::mlir::ArrayRef<::mlir::InnerTileAlignment>":$innerTileAlignments),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::cast<::mlir::TilingInterface>($_op.getOperation())
+ .getTiledImplementationFromOperandTiles(b, operandNumbers,
+ allOffsets, allSizes);
+ }]
+ >,
InterfaceMethod<
/*desc=*/[{
Method to return the tile of the iteration domain that uses a given
@@ -304,6 +381,33 @@ def TilingInterface : OpInterface<"TilingInterface"> {
return failure();
}]
>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Variant of `getIterationDomainTileFromOperandTiles` that additionally
+ takes a per-iteration-domain `InnerTileAlignment` hint (see
+ `InnerTileAlignment`). The hint is consulted only by pack/unpack
+ implementations and may be empty (all `Unknown`). The default
+ implementation ignores it and forwards to the hint-less overload.
+ }],
+ /*retType=*/"::llvm::LogicalResult",
+ /*methodName=*/"getIterationDomainTileFromOperandTiles",
+ /*args=*/(ins
+ "::mlir::OpBuilder &":$b,
+ "::mlir::ArrayRef<unsigned>":$operandNumbers,
+ "::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>> ":$allOffsets,
+ "::mlir::ArrayRef<::mlir::SmallVector<::mlir::OpFoldResult>> ":$allSizes,
+ "::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainOffsets,
+ "::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainSizes,
+ "::mlir::ArrayRef<::mlir::InnerTileAlignment> ":$innerTileAlignments),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::cast<::mlir::TilingInterface>($_op.getOperation())
+ .getIterationDomainTileFromOperandTiles(b, operandNumbers,
+ allOffsets, allSizes,
+ iterDomainOffsets,
+ iterDomainSizes);
+ }]
+ >,
InterfaceMethod<
/*desc=*/[{
Method to return the tile of the iteration domain based
diff --git a/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt b/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
index ec433284e17ad..ec4fd82468dbb 100644
--- a/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
@@ -34,6 +34,7 @@ add_mlir_dialect_library(MLIRLinalgDialect
MLIRMathDialect
MLIRMemRefDialect
MLIRTensorDialect
+ MLIRTilingInterface
MLIRValueBoundsOpInterface
MLIRViewLikeInterface
)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 1a56c5a483e73..8b2f5064e78bd 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2940,6 +2940,15 @@ SmallVector<utils::IteratorType> SoftmaxOp::getLoopIteratorTypes() {
return iteratorTypes;
}
+/// The inner tile alignment hint is only used by `linalg.pack` and
+/// `linalg.unpack` operations. Therefore, this is forwarded to the hint-less
+/// overload.
+FailureOr<TilingResult> SoftmaxOp::getTiledImplementation(
+ OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
+ return getTiledImplementation(builder, offsets, sizes);
+}
+
FailureOr<TilingResult>
SoftmaxOp::getTiledImplementation(OpBuilder &builder,
ArrayRef<OpFoldResult> offsets,
@@ -3293,6 +3302,15 @@ LogicalResult WinogradFilterTransformOp::getResultTilePosition(
return success();
}
+/// The inner tile alignment hint is only used by `linalg.pack` and
+/// `linalg.unpack` operations. Therefore, this is forwarded to the hint-less
+/// overload.
+FailureOr<TilingResult> WinogradFilterTransformOp::getTiledImplementation(
+ OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
+ return getTiledImplementation(builder, offsets, sizes);
+}
+
/// Implement tiling for winograd_filter_transform
/// The input of winograd_filter_transform is (F, KH, KW, C).
/// The output of winograd_filter_transform is (alphaH, alphaW, C, F)
@@ -3446,6 +3464,15 @@ LogicalResult WinogradInputTransformOp::getResultTilePosition(
return success();
}
+/// The inner tile alignment hint is only used by `linalg.pack` and
+/// `linalg.unpack` operations. Therefore, this is forwarded to the hint-less
+/// overload.
+FailureOr<TilingResult> WinogradInputTransformOp::getTiledImplementation(
+ OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
+ return getTiledImplementation(builder, offsets, sizes);
+}
+
/// Implement tiling for winograd_input_transform
/// The input of winograd_input_transform is (N, H, W, C).
/// The output of winograd_input_transform is (alphaH, alphaW, tileH, tileW, N,
@@ -3641,6 +3668,15 @@ LogicalResult WinogradOutputTransformOp::getResultTilePosition(
return success();
}
+/// The inner tile alignment hint is only used by `linalg.pack` and
+/// `linalg.unpack` operations. Therefore, this is forwarded to the hint-less
+/// overload.
+FailureOr<TilingResult> WinogradOutputTransformOp::getTiledImplementation(
+ OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
+ return getTiledImplementation(builder, offsets, sizes);
+}
+
/// Implement tiling for winograd_output_transform
/// The input of winograd_output_transform is (alphaH, alphaW, tileH, tileW, N,
/// F). The output of winograd_output_transform is (N, H, W, F) Users can
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index 4eaa7bf0233c6..13b959fc7b0cc 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -96,6 +96,16 @@ template <typename LinalgOpTy>
struct LinalgOpTilingInterface
: public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
LinalgOpTy> {
+ using Base =
+ TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
+ LinalgOpTy>;
+ // Inherit the defaulted hint-bearing overloads; these ops do not require the
+ // hint (no inner tiles).
+ using Base::generateResultTileValue;
+ using Base::getIterationDomainTileFromOperandTiles;
+ using Base::getTiledImplementation;
+ using Base::getTiledImplementationFromOperandTiles;
+
/// Return the loop iterator type.
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
@@ -924,6 +934,8 @@ static void generatePackOpScalarImplementationBody(PackOp packOp,
struct PackOpTiling
: public TilingInterface::ExternalModel<PackOpTiling, linalg::PackOp> {
+ using Base = TilingInterface::ExternalModel<PackOpTiling, linalg::PackOp>;
+ using Base::getTiledImplementation;
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
// Note that here we only consider untiled dimensions and outer tiled data
@@ -1055,20 +1067,39 @@ struct PackOpTiling
generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
+ return generateResultTileValue(op, b, resultNumber, offsets, sizes,
+ /*innerTileAlignments=*/{});
+ }
+
+ FailureOr<TilingResult> generateResultTileValue(
+ Operation *op, OpBuilder &b, unsigned resultNumber,
+ ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) const {
auto packOp = cast<PackOp>(op);
int64_t numTiles = packOp.getInnerDimsPos().size();
- // tensor.pack op is fusible (as a producer) only if full inner tiles are
+ // linalg.pack op is fusible (as a producer) only if full inner tiles are
// iterated or inner dims are not tiled. Otherwise, it will generate a
// sequence of non-trivial ops (for partial tiles).
for (auto offset : offsets.take_back(numTiles))
if (!isZeroInteger(offset))
return failure();
- for (auto iter :
- llvm::zip_equal(packOp.getMixedTiles(), sizes.take_back(numTiles)))
- if (!isEqualConstantIntOrValue(std::get<0>(iter), std::get<1>(iter)))
+ // Each requested inner-dim size must cover a full inner tile. A caller may
+ // instead assert this via an `Equal` alignment hint. The hint is indexed by
+ // source dim, matching the consumer-fusion path.
+ ArrayRef<int64_t> innerDimsPos = packOp.getInnerDimsPos();
+ SmallVector<OpFoldResult> mixedTiles = packOp.getMixedTiles();
+ ArrayRef<OpFoldResult> innerSizes = sizes.take_back(numTiles);
+ for (auto [i, pos] : llvm::enumerate(innerDimsPos)) {
+ InnerTileAlignment alignment =
+ pos < static_cast<int64_t>(innerTileAlignments.size())
+ ? innerTileAlignments[pos]
+ : InnerTileAlignment::Unknown;
+ if (alignment != InnerTileAlignment::Equal &&
+ !isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i]))
return failure();
+ }
FailureOr<TilingResult> tilingResult = getTiledImplementation(
op, b, offsets.drop_back(numTiles), sizes.drop_back(numTiles));
@@ -1123,15 +1154,27 @@ struct PackOpTiling
return success();
}
+ LogicalResult getIterationDomainTileFromOperandTiles(
+ Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
+ ArrayRef<SmallVector<OpFoldResult>> allOffsets,
+ ArrayRef<SmallVector<OpFoldResult>> allSizes,
+ SmallVectorImpl<OpFoldResult> &resultOffsets,
+ SmallVectorImpl<OpFoldResult> &resultSizes) const {
+ return getIterationDomainTileFromOperandTiles(
+ op, b, operandNumbers, allOffsets, allSizes, resultOffsets, resultSizes,
+ /*innerTileAlignments=*/{});
+ }
+
/// Method to return the position of iteration domain tile computed by the
- /// tiled operation. In current `tensor.pack` context, the `resultOffsets` and
+ /// tiled operation. In current `linalg.pack` context, the `resultOffsets` and
/// `resultSizes` only cover outer dimensions.
LogicalResult getIterationDomainTileFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
SmallVectorImpl<OpFoldResult> &resultOffsets,
- SmallVectorImpl<OpFoldResult> &resultSizes) const {
+ SmallVectorImpl<OpFoldResult> &resultSizes,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) const {
if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
LLVM_DEBUG(
{ llvm::dbgs() << "unsupported operands for consumer fusion"; });
@@ -1162,6 +1205,14 @@ struct PackOpTiling
std::optional<int64_t> cstInnerSize =
getConstantIntValue(dimAndTileMapping[dim]);
+ // A caller-supplied alignment hint (see InnerTileAlignment) asserts
+ // that this packed dimension is tiled and how its loop tile size
+ // relates to the pack op inner tile size.
+ InnerTileAlignment innerTileAlignment =
+ dim < static_cast<int64_t>(innerTileAlignments.size())
+ ? innerTileAlignments[dim]
+ : InnerTileAlignment::Unknown;
+
// If a dimension is not tiled, it is always valid to fuse the pack op,
// even if the op has padding semantics. Because it always generates a
// full slice along the dimension. The tile sizes are for unpacked
@@ -1169,9 +1220,14 @@ struct PackOpTiling
// dimension is tiled.
// TODO: It could be untiled if the `srcDimSize` is dynamic. It is a
// hard check to determine if a dimension is tiled or not.
+ // A non-`Unknown` hint also means the caller asserts the dimension is
+ // tiled: `cstTileSize` is an upper bound, so a scalable/`min`-shaped
+ // tile (whose bound equals `srcDimSize`) would otherwise be mistaken
+ // for untiled and bypass the hint below.
int64_t srcDimSize = packOp.getSourceType().getDimSize(dim);
int64_t destDimSize = outerShapeWithoutTranspose[dim];
- bool isTiled = failed(cstTileSize) ||
+ bool isTiled = innerTileAlignment != InnerTileAlignment::Unknown ||
+ failed(cstTileSize) ||
ShapedType::isDynamic(srcDimSize) ||
cstTileSize.value() < srcDimSize;
if (!isTiled) {
@@ -1199,9 +1255,30 @@ struct PackOpTiling
// another word, we can only support tiling with consumer if the tile
// size for the producer is a multiple of the inner tile size for the
// packed dimensions at this moment.
- if ((failed(cstTileSize) || !cstInnerSize ||
- *cstTileSize % *cstInnerSize != 0))
- return failure();
+
+ // The caller may assert how this packed dimension's loop tile size
+ // relates to the inner tile size via `innerTileAlignments` (see
+ // InnerTileAlignment). The hint is the source of truth and is honored
+ // when present. When both sizes are also statically known we assert the
+ // hint agrees with them (a contradicting hint is a caller bug). When
+ // the hint is `Unknown`, fall back to requiring a statically-provable
+ // multiple.
+ bool assumeInnerTileSizesMatchTiles =
+ innerTileAlignment == InnerTileAlignment::Equal;
+ bool staticallyDecidable =
+ !failed(cstTileSize) && cstInnerSize.has_value();
+ if (innerTileAlignment == InnerTileAlignment::Unknown) {
+ if (!staticallyDecidable || *cstTileSize % *cstInnerSize != 0)
+ return failure();
+ } else if (staticallyDecidable) {
+ assert(*cstTileSize % *cstInnerSize == 0 &&
+ "InnerTileAlignment hint contradicts statically known tile "
+ "sizes");
+ assert((innerTileAlignment != InnerTileAlignment::Equal ||
+ *cstTileSize == *cstInnerSize) &&
+ "InnerTileAlignment::Equal contradicts statically known tile "
+ "sizes");
+ }
using AV = affine::AffineValueExpr;
affine::AffineBuilder ab(b, loc);
@@ -1212,7 +1289,11 @@ struct PackOpTiling
auto avSize = AV(dim0).bind(sizes[dim]);
auto avTileSize = AV(sym).bind(dimAndTileMapping[dim]);
outerDimOffsets.push_back(ab.floor(avOffset, avTileSize));
- outerDimSizes.push_back(ab.ceil(avSize, avTileSize));
+ // If the tile size equals the inner tile size, the outer dims are
+ // always 1.
+ outerDimSizes.push_back(assumeInnerTileSizesMatchTiles
+ ? b.getIndexAttr(1)
+ : ab.ceil(avSize, avTileSize));
} else {
outerDimOffsets.push_back(offsets[dim]);
outerDimSizes.push_back(sizes[dim]);
@@ -1224,14 +1305,23 @@ struct PackOpTiling
return success();
}
- /// Method to return the tiled implementation of tensor.pack as a consumer.
FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
+ return getTiledImplementationFromOperandTiles(op, b, operandNumbers,
+ allOffsets, allSizes,
+ /*innerTileAlignments=*/{});
+ }
+
+ /// Method to return the tiled implementation of linalg.pack as a consumer.
+ FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
+ Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
+ ArrayRef<SmallVector<OpFoldResult>> allOffsets,
+ ArrayRef<SmallVector<OpFoldResult>> allSizes,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) const {
if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
- LLVM_DEBUG(
- { llvm ::dbgs() << "unhandled operands for consumer fusion"; });
+ LLVM_DEBUG({ llvm::dbgs() << "unhandled operands for consumer fusion"; });
return failure();
}
@@ -1257,7 +1347,7 @@ struct PackOpTiling
SmallVector<OpFoldResult> outerDimOffsets, outerDimSizes;
if (failed(getIterationDomainTileFromOperandTiles(
op, b, operandNumbers, allOffsets, allSizes, outerDimOffsets,
- outerDimSizes)))
+ outerDimSizes, innerTileAlignments)))
return failure();
SmallVector<OpFoldResult> outputOffsets, outputSizes;
@@ -1296,10 +1386,10 @@ struct UnpackTileDimInfo {
/// Returns the needed information for tiling unpack op on `tileDim` with given
/// `tileOffset` and `tileSize`. For more details, see the comment of the
/// `getTiledImplementation`.
-static UnpackTileDimInfo getUnpackTileDimInfo(OpBuilder &b, UnPackOp unpackOp,
- int64_t tileDim,
- OpFoldResult tileOffset,
- OpFoldResult tileSize) {
+static UnpackTileDimInfo
+getUnpackTileDimInfo(OpBuilder &b, UnPackOp unpackOp, int64_t tileDim,
+ OpFoldResult tileOffset, OpFoldResult tileSize,
+ InnerTileAlignment innerTileAlignment) {
UnpackTileDimInfo info;
Attribute zeroAttr = b.getIndexAttr(0);
Attribute oneAttr = b.getIndexAttr(1);
@@ -1329,13 +1419,36 @@ static UnpackTileDimInfo getUnpackTileDimInfo(OpBuilder &b, UnPackOp unpackOp,
presburger::BoundType::UB, tileSize,
/*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
std::optional<int64_t> cstInnerSize = getConstantIntValue(innerTileSize);
- if (!failed(cstSize) && cstInnerSize) {
- if (*cstSize % *cstInnerSize == 0)
+ // The caller may assert how this dimension's loop tile size relates to the
+ // op's inner tile size via `innerTileAlignment` (see InnerTileAlignment). The
+ // hint is the source of truth and is honored when present: `Equal`/`Multiple`
+ // both mean the tile is aligned to (a multiple of) the inner tile, and
+ // `Equal` additionally collapses the source slice to a single inner tile.
+ // When both sizes are also statically known we assert the hint agrees with
+ // them (a contradicting hint is a caller bug). When `Unknown`, fall back to
+ // the static upper-bound path below.
+ bool assumeInnerTileSizesMatchTiles =
+ innerTileAlignment == InnerTileAlignment::Equal;
+ bool staticallyDecidable = !failed(cstSize) && cstInnerSize.has_value();
+ if (innerTileAlignment != InnerTileAlignment::Unknown) {
+ info.isAlignedToInnerTileSize = true;
+ if (staticallyDecidable) {
+ assert(*cstSize % *cstInnerSize == 0 &&
+ "InnerTileAlignment hint contradicts statically known tile sizes");
+ assert((innerTileAlignment != InnerTileAlignment::Equal ||
+ *cstSize == *cstInnerSize) &&
+ "InnerTileAlignment::Equal contradicts statically known tile "
+ "sizes");
+ }
+ }
+ if (info.isAlignedToInnerTileSize || (!failed(cstSize) && cstInnerSize)) {
+ if (!info.isAlignedToInnerTileSize && *cstSize % *cstInnerSize == 0)
info.isAlignedToInnerTileSize = true;
// If the tiling size equals to the inner tiling size, the outer dims are
// always 1.
- if (*cstInnerSize == *cstSize) {
+ if (assumeInnerTileSizesMatchTiles ||
+ (cstInnerSize && !failed(cstSize) && *cstInnerSize == *cstSize)) {
auto lhs = AV(dim0).bind(tileOffset);
auto rhs = AV(dim1).bind(innerTileSize);
info.sourceOffset = ab.floor(lhs, rhs);
@@ -1391,6 +1504,8 @@ static UnpackTileDimInfo getUnpackTileDimInfo(OpBuilder &b, UnPackOp unpackOp,
struct UnPackOpTiling
: public TilingInterface::ExternalModel<UnPackOpTiling, linalg::UnPackOp> {
+ using Base = TilingInterface::ExternalModel<UnPackOpTiling, linalg::UnPackOp>;
+ using Base::getIterationDomainTileFromOperandTiles;
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
auto unpackOp = cast<UnPackOp>(op);
@@ -1421,6 +1536,14 @@ struct UnPackOpTiling
getTiledImplementation(Operation *op, OpBuilder &b,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
+ return getTiledImplementation(op, b, offsets, sizes,
+ /*innerTileAlignments=*/{});
+ }
+
+ FailureOr<TilingResult> getTiledImplementation(
+ Operation *op, OpBuilder &b, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) const {
auto unpackOp = cast<UnPackOp>(op);
// TODO: Support Memref UnPackOp. Temporarily return failure.
if (!unpackOp.hasPureTensorSemantics())
@@ -1440,8 +1563,11 @@ struct UnPackOpTiling
SmallVector<OpFoldResult> sliceSrcIndices, sliceSrcSizes;
SmallVector<OpFoldResult> destExpandedSizes, resultOffsetsFromDest;
for (auto dim : llvm::seq<int64_t>(0, destRank)) {
- UnpackTileDimInfo info =
- getUnpackTileDimInfo(b, unpackOp, dim, offsets[dim], sizes[dim]);
+ UnpackTileDimInfo info = getUnpackTileDimInfo(
+ b, unpackOp, dim, offsets[dim], sizes[dim],
+ dim < static_cast<int64_t>(innerTileAlignments.size())
+ ? innerTileAlignments[dim]
+ : InnerTileAlignment::Unknown);
if (!info.isAlignedToInnerTileSize)
isPerfectTilingCase = false;
sliceSrcIndices.push_back(info.sourceOffset);
@@ -1510,8 +1636,16 @@ struct UnPackOpTiling
generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
+ return generateResultTileValue(op, b, resultNumber, offsets, sizes,
+ /*innerTileAlignments=*/{});
+ }
+
+ FailureOr<TilingResult> generateResultTileValue(
+ Operation *op, OpBuilder &b, unsigned resultNumber,
+ ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) const {
FailureOr<TilingResult> tilingResult =
- getTiledImplementation(op, b, offsets, sizes);
+ getTiledImplementation(op, b, offsets, sizes, innerTileAlignments);
if (failed(tilingResult))
return failure();
return tilingResult.value();
@@ -1642,11 +1776,21 @@ struct UnPackOpTiling
return success();
}
- /// Method to return the tiled implementation of tensor.unpack as a consumer.
FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
+ return getTiledImplementationFromOperandTiles(op, b, operandNumbers,
+ allOffsets, allSizes,
+ /*innerTileAlignments=*/{});
+ }
+
+ /// Method to return the tiled implementation of linalg.unpack as a consumer.
+ FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
+ Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
+ ArrayRef<SmallVector<OpFoldResult>> allOffsets,
+ ArrayRef<SmallVector<OpFoldResult>> allSizes,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) const {
if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
LLVM_DEBUG({ llvm::dbgs() << "unhandled operands for consumer fusion"; });
return failure();
@@ -1659,13 +1803,37 @@ struct UnPackOpTiling
ArrayRef<OpFoldResult> offsets(allOffsets[0]);
ArrayRef<OpFoldResult> sizes(allSizes[0]);
- // tensor.unpack op is fusible (as a consumer) only if inner dims are not
- // tiled.
+ // linalg.unpack op is fusible (as a consumer) only if the inner dims are
+ // not tiled, i.e. each inner-dim loop tile size equals the inner tile size.
+ // The caller may assert this per inner dim via InnerTileAlignment::Equal;
+ // otherwise we require a statically-provable equality.
int64_t numTiles = unPackOp.getInnerDimsPos().size();
- for (auto iter :
- llvm::zip_equal(unPackOp.getMixedTiles(), sizes.take_back(numTiles))) {
- if (!isEqualConstantIntOrValue(std::get<0>(iter), std::get<1>(iter)))
- return failure();
+ ArrayRef<int64_t> innerDimsPos = unPackOp.getInnerDimsPos();
+ SmallVector<OpFoldResult> mixedTiles = unPackOp.getMixedTiles();
+ ArrayRef<OpFoldResult> innerSizes = sizes.take_back(numTiles);
+ for (int64_t i = 0; i < numTiles; ++i) {
+ // `innerTileAlignments` is indexed by the unpack iteration domain (the
+ // dest dims); the i-th inner tile lives on dest dim `innerDimsPos[i]`.
+ int64_t destDim = innerDimsPos[i];
+ bool hintedEqual =
+ destDim < static_cast<int64_t>(innerTileAlignments.size()) &&
+ innerTileAlignments[destDim] == InnerTileAlignment::Equal;
+ // The hint is the source of truth: honor a caller `Equal` assertion. When
+ // both sizes are also statically known, assert the hint agrees with them
+ // (a contradicting hint is a caller bug) rather than silently ignoring
+ // it. Without an `Equal` hint, require a statically-provable equality
+ // (the inner dim must not be tiled).
+ if (hintedEqual) {
+ assert((!getConstantIntValue(mixedTiles[i]) ||
+ !getConstantIntValue(innerSizes[i]) ||
+ isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i])) &&
+ "InnerTileAlignment::Equal contradicts statically known tile "
+ "sizes");
+ continue;
+ }
+ if (isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i]))
+ continue;
+ return failure();
}
Location loc = unPackOp.getLoc();
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
index 124a63281a37c..782412557fe27 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
@@ -22,6 +22,11 @@ using namespace mlir::tensor;
namespace {
struct PadOpTiling : public TilingInterface::ExternalModel<PadOpTiling, PadOp> {
+ using Base = TilingInterface::ExternalModel<PadOpTiling, PadOp>;
+ // Inherit the defaulted hint-bearing overloads; this op does not require the
+ // hint.
+ using Base::generateResultTileValue;
+ using Base::getTiledImplementation;
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
auto padOp = cast<PadOp>(op);
diff --git a/mlir/lib/Interfaces/TilingInterface.cpp b/mlir/lib/Interfaces/TilingInterface.cpp
index 67ddb5b1c1c50..4183d4fd59cd6 100644
--- a/mlir/lib/Interfaces/TilingInterface.cpp
+++ b/mlir/lib/Interfaces/TilingInterface.cpp
@@ -12,6 +12,28 @@
#include "mlir/Interfaces/TilingInterface.h"
+#include "llvm/ADT/SmallVectorExtras.h"
+
using namespace mlir;
+LogicalResult mlir::verifyInnerTileAlignments(Operation *op,
+ ArrayRef<int64_t> alignments) {
+ for (int64_t a : alignments)
+ if (!isValidInnerTileAlignment(a))
+ return op->emitOpError()
+ << "expected inner_tile_alignments entries to be one of 0 "
+ "(Unknown), 1 (Multiple) or 2 (Equal), but got "
+ << a;
+ return success();
+}
+
+SmallVector<InnerTileAlignment>
+mlir::convertInnerTileAlignments(ArrayRef<int64_t> alignments) {
+ return llvm::map_to_vector(alignments, [](int64_t v) {
+ assert(isValidInnerTileAlignment(v) &&
+ "invalid InnerTileAlignment; should be rejected by the verifier");
+ return static_cast<InnerTileAlignment>(v);
+ });
+}
+
#include "mlir/Interfaces/TilingInterface.cpp.inc"
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 20d72a2fd09ca..f241e6f05185b 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1237,6 +1237,12 @@ SmallVector<utils::IteratorType> TilingNoDpsOp::getLoopIteratorTypes() {
return {};
}
+FailureOr<TilingResult> TilingNoDpsOp::getTiledImplementation(
+ OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
+ return getTiledImplementation(builder, offsets, sizes);
+}
+
FailureOr<TilingResult>
TilingNoDpsOp::getTiledImplementation(OpBuilder &builder,
ArrayRef<OpFoldResult> offsets,
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index f6c845fefc915..0ce66a086186d 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -11731,6 +11731,7 @@ cc_library(
":Support",
":TilingInterfaceIncGen",
":ViewLikeInterface",
+ "//llvm:Support",
],
)
>From 998af51e748cac344d4c88c28baecf8b1d268ea4 Mon Sep 17 00:00:00 2001
From: Ege Beysel <beyselege at gmail.com>
Date: Wed, 10 Jun 2026 17:07:06 +0000
Subject: [PATCH 2/4] feat(mlir/scf): thread InnerTileAlignment through the SCF
tiling/fusion driver
Forward the optional InnerTileAlignment hint to the hint-bearing TilingInterface
overloads from the SCF tiling/fusion driver. It is passed as a defaulted
argument to tileUsingSCF and tileConsumerAndFuseProducersUsingSCF (and the
existing tileAndFuseProducerOfSlice / tileAndFuseConsumer entry points). As a
pack/unpack-specific caller assertion the hint is kept off the general
SCFTilingOptions struct and passed explicitly where pack/unpack tiling/fusion
happens.
Signed-off-by: Ege Beysel <beyselege at gmail.com>
---
.../SCF/Transforms/TileUsingInterface.h | 60 ++++++++++++++++---
.../Dialect/Tensor/Transforms/Transforms.h | 12 ++--
.../SCF/Transforms/TileUsingInterface.cpp | 60 ++++++++++++++-----
.../SwapExtractSliceWithProducerPatterns.cpp | 11 ++--
4 files changed, 111 insertions(+), 32 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
index 0005fad3d5c01..0bd5d19b136d0 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
@@ -31,6 +31,17 @@ namespace scf {
using SCFTileSizeComputationFunction =
std::function<SmallVector<OpFoldResult>(OpBuilder &, Operation *)>;
+/// Computation function returning, for the op currently being tiled or fused,
+/// the per-iteration-domain-dimension `InnerTileAlignment` array.
+///`tileSizes` holds the given tile sizes for pure tiling
+/// and is empty for fusion; `slices` holds the `tensor.extract_slice` (producer
+/// fusion) or `tensor.insert_slice`/`tensor.parallel_insert_slice`s (consumer
+/// fusion) being fused, and is empty for pure tiling. Only pack/unpack
+/// implementations consult the result; every other op ignores it.
+using InnerTileAlignmentFnTy = std::function<SmallVector<InnerTileAlignment>(
+ TilingInterface op, ArrayRef<OpFoldResult> tileSizes,
+ ArrayRef<Operation *> slices)>;
+
/// Options to use to control tiling.
struct SCFTilingOptions {
/// Specify which loop construct to use for tile and fuse.
@@ -65,6 +76,26 @@ struct SCFTilingOptions {
return *this;
}
+ /// Optional control function returning, per tiled/fused op, the inner tile
+ /// alignment hints.
+ InnerTileAlignmentFnTy innerTileAlignmentFn = nullptr;
+ SCFTilingOptions &setInnerTileAlignmentFn(InnerTileAlignmentFnTy fn) {
+ innerTileAlignmentFn = std::move(fn);
+ return *this;
+ }
+
+ /// Sets a constant `innerTileAlignmentFn` that returns `alignments` for every
+ /// op, ignoring which op is being tiled or fused. Use when a single fixed
+ /// array is correct for the whole tiling/fusion (e.g. tiling a single op).
+ SCFTilingOptions &
+ setInnerTileAlignments(ArrayRef<InnerTileAlignment> alignments) {
+ SmallVector<InnerTileAlignment> fixed = llvm::to_vector(alignments);
+ innerTileAlignmentFn =
+ [fixed = std::move(fixed)](TilingInterface, ArrayRef<OpFoldResult>,
+ ArrayRef<Operation *>) { return fixed; };
+ return *this;
+ }
+
//-------------------------------------------------------------------------//
// Options related to tiling using `scf.forall`.
//-------------------------------------------------------------------------//
@@ -287,20 +318,26 @@ struct SCFTileAndFuseOptions {
std::optional<FrozenRewritePatternSet> cleanupPatterns = std::nullopt;
};
-/// Fuse the producer of the source of `candidateSliceOp` by computing the
-/// required slice of the producer in-place. Note that the method
-/// replaces the uses of `candidateSliceOp` with the tiled and fused producer
-/// value but does not delete the slice operation.
+/// Result of fusing the producer of the source of a `tensor.extract_slice`.
struct SCFFuseProducerOfSliceResult {
OpResult origProducer; // Original untiled producer.
Value tiledAndFusedProducer; // Tile and fused producer value.
SmallVector<Operation *> tiledOps;
SmallVector<Operation *> generatedSlices;
};
+/// Fuse the producer of the source of `candidateSliceOp` by computing the
+/// required slice of the producer in-place. Note that the method
+/// replaces the uses of `candidateSliceOp` with the tiled and fused producer
+/// value but does not delete the slice operation.
+///
+/// When the fused producer is a `linalg.pack`/`linalg.unpack`, `fn` (if
+/// non-null) is invoked with the producer and `candidateSliceOp` to obtain the
+/// inner-tile alignment hint in the producer's own iteration domain.
std::optional<SCFFuseProducerOfSliceResult>
tileAndFuseProducerOfSlice(RewriterBase &rewriter,
tensor::ExtractSliceOp candidateSliceOp,
- MutableArrayRef<LoopLikeOpInterface> loops);
+ MutableArrayRef<LoopLikeOpInterface> loops,
+ const InnerTileAlignmentFnTy &fn = nullptr);
/// Reconstruct the fused producer from within the tiled-and-fused code. Based
/// on the slice of the producer computed in place it is possible that within
@@ -426,18 +463,27 @@ struct SCFFuseConsumerOfSliceResult {
SmallVector<OpOperand *> tiledAndFusedConsumerOperands;
SmallVector<Operation *> tiledOps;
};
+/// When the consumer is a `linalg.pack`/`linalg.unpack`, `fn` (if non-null) is
+/// invoked with the consumer and `candidateSlices` to obtain the inner-tile
+/// alignment hint in the consumer's own iteration domain.
FailureOr<scf::SCFFuseConsumerOfSliceResult>
tileAndFuseConsumerOfSlices(RewriterBase &rewriter,
ArrayRef<Operation *> candidateSlices,
- MutableArrayRef<LoopLikeOpInterface> loops);
+ MutableArrayRef<LoopLikeOpInterface> loops,
+ const InnerTileAlignmentFnTy &fn = nullptr);
/// Fuse the `consumer` operation into the loop nest provided by `loops`.
/// The transformation looks for operands in the `consumer` that are defined
/// by the outermost loop of the loop nest in `loops`. The nested loop is
/// expected to have the structure of the loops generated through tiling.
+///
+/// When the consumer is a `linalg.pack`/`linalg.unpack`, `fn` (if non-null) is
+/// invoked with the `consumer` and the fused slices to obtain the inner-tile
+/// alignment hint in the consumer's own iteration domain.
FailureOr<scf::SCFFuseConsumerOfSliceResult>
tileAndFuseConsumer(RewriterBase &rewriter, Operation *consumer,
- MutableArrayRef<LoopLikeOpInterface> loops);
+ MutableArrayRef<LoopLikeOpInterface> loops,
+ const InnerTileAlignmentFnTy &fn = nullptr);
/// Method to lower an `op` that implements the `TilingInterface` to
/// loops/scalars.
diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
index 093393eca7436..341b3b1dc3472 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
@@ -11,6 +11,7 @@
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
namespace mlir {
@@ -29,7 +30,8 @@ namespace tensor {
/// transform dialect that control is done within the transform dialect. Other
/// use cases can inherit from this pattern and add necessary controls.
FailureOr<TilingResult> replaceExtractSliceWithTiledProducer(
- OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp);
+ OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp,
+ ArrayRef<InnerTileAlignment> innerTileAlignments = {});
/// Method to swap `tensor.insert_slice`s with their consumers when the
/// consumer implements the `TilingInterface`. The size of `sliceOps` and
@@ -37,10 +39,10 @@ FailureOr<TilingResult> replaceExtractSliceWithTiledProducer(
/// `consumerOperands` represents a use of the the corresponding
/// entry in `sliceOps` in the consumer. All entries of `consumerOperands` is
/// expected to be uses in the same consumer.
-FailureOr<TilingResult>
-replaceInsertSlicesWithTiledConsumer(OpBuilder &builder,
- ArrayRef<tensor::InsertSliceOp> sliceOps,
- ArrayRef<OpOperand *> consumerOperands);
+FailureOr<TilingResult> replaceInsertSlicesWithTiledConsumer(
+ OpBuilder &builder, ArrayRef<tensor::InsertSliceOp> sliceOps,
+ ArrayRef<OpOperand *> consumerOperands,
+ ArrayRef<InnerTileAlignment> innerTileAlignments = {});
//===----------------------------------------------------------------------===//
// Populate functions.
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index 963e2d6d4f37b..ad4aff893a03a 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -824,9 +824,11 @@ getTiledImplementation(RewriterBase &rewriter, TilingInterface op,
ArrayRef<OpFoldResult> sizes, ValueRange ivs,
ArrayRef<OpFoldResult> numThreads,
ArrayRef<OpFoldResult> givenTileSizes,
+ ArrayRef<InnerTileAlignment> innerTileAlignments,
const SetVector<unsigned> &reductionDims) {
if (reductionStrategy == ReductionTilingStrategy::FullReduction) {
- return op.getTiledImplementation(rewriter, offsets, sizes);
+ return op.getTiledImplementation(rewriter, offsets, sizes,
+ innerTileAlignments);
}
auto redOp = dyn_cast<PartialReductionOpInterface>(op.getOperation());
@@ -1190,10 +1192,15 @@ mlir::scf::tileUsingSCF(RewriterBase &rewriter, TilingInterface op,
}
// 5c. Tile the cloned operation.
- tilingResult =
- getTiledImplementation(rewriter, clonedOp, options.reductionStrategy,
- regionIterArgs, tileOffsetsVec, tileSizesVec,
- ivs, numThreads, givenTileSizes, reductionDims);
+ SmallVector<InnerTileAlignment> innerTileAlignments =
+ options.innerTileAlignmentFn
+ ? options.innerTileAlignmentFn(clonedOp, givenTileSizes,
+ /*slices=*/{})
+ : SmallVector<InnerTileAlignment>{};
+ tilingResult = getTiledImplementation(
+ rewriter, clonedOp, options.reductionStrategy, regionIterArgs,
+ tileOffsetsVec, tileSizesVec, ivs, numThreads, givenTileSizes,
+ innerTileAlignments, reductionDims);
if (failed(tilingResult)) {
rewriter.eraseOp(clonedOp);
return op.emitOpError("failed to tile operation");
@@ -1346,7 +1353,8 @@ getUntiledProducerFromSliceSource(OpOperand *source,
std::optional<scf::SCFFuseProducerOfSliceResult>
mlir::scf::tileAndFuseProducerOfSlice(
RewriterBase &rewriter, tensor::ExtractSliceOp candidateSliceOp,
- MutableArrayRef<LoopLikeOpInterface> loops) {
+ MutableArrayRef<LoopLikeOpInterface> loops,
+ const InnerTileAlignmentFnTy &fn) {
// 1. Get the producer of the source (potentially walking through
// `iter_args` of nested `scf.for`)
auto [fusableProducer, destinationInitArg] =
@@ -1356,6 +1364,14 @@ mlir::scf::tileAndFuseProducerOfSlice(
return std::nullopt;
unsigned resultNumber = fusableProducer.getResultNumber();
+ // Resolve the inner-tile alignment hint for the producer in its own iteration
+ // domain via the control function (consulted only by pack/unpack).
+ SmallVector<InnerTileAlignment> innerTileAlignments;
+ if (fn)
+ if (auto producer = dyn_cast<TilingInterface>(fusableProducer.getOwner()))
+ innerTileAlignments =
+ fn(producer, /*tileSizes=*/{}, {candidateSliceOp.getOperation()});
+
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(candidateSliceOp);
@@ -1394,7 +1410,7 @@ mlir::scf::tileAndFuseProducerOfSlice(
FailureOr<TilingResult> tileAndFuseResult =
tensor::replaceExtractSliceWithTiledProducer(
rewriter, clonedCandidateSliceOp,
- clonedProducerOp->getResult(resultNumber));
+ clonedProducerOp->getResult(resultNumber), innerTileAlignments);
if (failed(tileAndFuseResult))
return std::nullopt;
// Note: Do not delete the candidateSliceOp, since its passed in from the
@@ -1819,8 +1835,8 @@ mlir::scf::tileConsumerAndFuseProducersUsingSCF(
// values produced by operations that implement the `TilingInterface`.
// Add these operations to the worklist.
std::optional<scf::SCFFuseProducerOfSliceResult> fusedResult =
- tileAndFuseProducerOfSlice(rewriter, worklistItem.candidateSlice,
- loops);
+ tileAndFuseProducerOfSlice(rewriter, worklistItem.candidateSlice, loops,
+ options.tilingOptions.innerTileAlignmentFn);
if (!fusedResult)
continue;
@@ -2206,9 +2222,17 @@ static FailureOr<scf::SCFFuseConsumerOfSliceResult>
tileAndFuseConsumerOfSlicesImpl(RewriterBase &rewriter, Operation *consumerOp,
ArrayRef<OpOperand *> consumerOpOperands,
ArrayRef<Operation *> candidateSlices,
- MutableArrayRef<LoopLikeOpInterface> loops) {
+ MutableArrayRef<LoopLikeOpInterface> loops,
+ const mlir::scf::InnerTileAlignmentFnTy &fn) {
assert(!loops.empty() && "expected loops to be not empty");
+ // Resolve the inner-tile alignment hint for the consumer in its own iteration
+ // domain via the control function (consulted only by pack/unpack).
+ SmallVector<InnerTileAlignment> innerTileAlignments;
+ if (fn)
+ if (auto consumer = dyn_cast<TilingInterface>(consumerOp))
+ innerTileAlignments = fn(consumer, /*tileSizes=*/{}, candidateSlices);
+
// 1. Check assumption for loop with `reorderOperations` disabled.
if (failed(checkAssumptionForLoop(loops.front(), consumerOp, false))) {
return rewriter.notifyMatchFailure(
@@ -2282,7 +2306,8 @@ tileAndFuseConsumerOfSlicesImpl(RewriterBase &rewriter, Operation *consumerOp,
// `operandNumber` with the source of the cloned tensor.insert_slice op.
FailureOr<TilingResult> tileAndFuseResult =
tensor::replaceInsertSlicesWithTiledConsumer(rewriter, clonedInsertSlices,
- clonedOpFusedOperandsList);
+ clonedOpFusedOperandsList,
+ innerTileAlignments);
if (failed(tileAndFuseResult)) {
return failure();
}
@@ -2329,7 +2354,7 @@ tileAndFuseConsumerOfSlicesImpl(RewriterBase &rewriter, Operation *consumerOp,
SmallVector<OpFoldResult> iterDomainOffsets, iterDomainSizes;
if (failed(clonedConsumerOp.getIterationDomainTileFromOperandTiles(
rewriter, operandNumbers, allOffsets, allSizes, iterDomainOffsets,
- iterDomainSizes))) {
+ iterDomainSizes, innerTileAlignments))) {
return rewriter.notifyMatchFailure(
clonedConsumerOp,
"can't get iter domain position from input position");
@@ -2420,7 +2445,8 @@ tileAndFuseConsumerOfSlicesImpl(RewriterBase &rewriter, Operation *consumerOp,
FailureOr<scf::SCFFuseConsumerOfSliceResult>
mlir::scf::tileAndFuseConsumerOfSlices(
RewriterBase &rewriter, ArrayRef<Operation *> candidateSlices,
- MutableArrayRef<LoopLikeOpInterface> loops) {
+ MutableArrayRef<LoopLikeOpInterface> loops,
+ const InnerTileAlignmentFnTy &fn) {
if (candidateSlices.empty()) {
return rewriter.notifyMatchFailure(
rewriter.getUnknownLoc(),
@@ -2455,7 +2481,7 @@ mlir::scf::tileAndFuseConsumerOfSlices(
return tileAndFuseConsumerOfSlicesImpl(rewriter, consumerOp,
maybeConsumerOpOperands.value(),
- candidateSlices, loops);
+ candidateSlices, loops, fn);
}
/// For a given `result` of a `forallOp` return the
@@ -2520,7 +2546,8 @@ getProducingInsertSliceLikeOp(OpResult result,
FailureOr<scf::SCFFuseConsumerOfSliceResult>
mlir::scf::tileAndFuseConsumer(RewriterBase &rewriter, Operation *consumer,
- MutableArrayRef<LoopLikeOpInterface> loops) {
+ MutableArrayRef<LoopLikeOpInterface> loops,
+ const InnerTileAlignmentFnTy &fn) {
if (!isa<TilingInterface>(consumer)) {
return rewriter.notifyMatchFailure(
consumer, "unhandled consumer that does not implement TilingInterface");
@@ -2565,8 +2592,9 @@ mlir::scf::tileAndFuseConsumer(RewriterBase &rewriter, Operation *consumer,
}
candidateSlices.push_back(slice.value());
}
+
return tileAndFuseConsumerOfSlicesImpl(
- rewriter, consumer, consumerFusableOperands, candidateSlices, loops);
+ rewriter, consumer, consumerFusableOperands, candidateSlices, loops, fn);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Tensor/Transforms/SwapExtractSliceWithProducerPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/SwapExtractSliceWithProducerPatterns.cpp
index 7903f3c51b73b..0e0262037f08d 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/SwapExtractSliceWithProducerPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/SwapExtractSliceWithProducerPatterns.cpp
@@ -23,7 +23,8 @@
using namespace mlir;
FailureOr<TilingResult> tensor::replaceExtractSliceWithTiledProducer(
- OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producer) {
+ OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producer,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) {
auto producerOp = dyn_cast<TilingInterface>(producer.getOwner());
if (!producerOp)
return failure();
@@ -34,7 +35,7 @@ FailureOr<TilingResult> tensor::replaceExtractSliceWithTiledProducer(
FailureOr<TilingResult> tiledResult = producerOp.generateResultTileValue(
builder, producer.getResultNumber(), sliceOp.getMixedOffsets(),
- sliceOp.getMixedSizes());
+ sliceOp.getMixedSizes(), innerTileAlignments);
if (failed(tiledResult))
return failure();
@@ -61,7 +62,8 @@ FailureOr<TilingResult> tensor::replaceExtractSliceWithTiledProducer(
FailureOr<TilingResult> tensor::replaceInsertSlicesWithTiledConsumer(
OpBuilder &builder, ArrayRef<tensor::InsertSliceOp> sliceOps,
- ArrayRef<OpOperand *> consumerOperands) {
+ ArrayRef<OpOperand *> consumerOperands,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) {
if (sliceOps.empty()) {
LLVM_DEBUG(
{ llvm::dbgs() << "expected candidate slices list to be non-empty"; });
@@ -107,7 +109,8 @@ FailureOr<TilingResult> tensor::replaceInsertSlicesWithTiledConsumer(
}
FailureOr<TilingResult> tiledResult =
consumerOp.getTiledImplementationFromOperandTiles(
- builder, consumerOperandNums, allOffsets, allSizes);
+ builder, consumerOperandNums, allOffsets, allSizes,
+ innerTileAlignments);
if (failed(tiledResult))
return failure();
>From 7ec7f9d0a5c9b3746f2f2a8026dd98f7127e81c8 Mon Sep 17 00:00:00 2001
From: Ege Beysel <beyselege at gmail.com>
Date: Wed, 10 Jun 2026 17:07:06 +0000
Subject: [PATCH 3/4] feat(mlir/transform): expose inner_tile_alignments on
tiling/fusion transform ops
Expose the per-dimension InnerTileAlignment hint as an optional
inner_tile_alignments attribute on transform.structured.tile_using_for,
transform.structured.fuse and transform.structured.fuse_into_containing_op (and
the test consumer-fusion op), forwarding it to the SCF driver.
Signed-off-by: Ege Beysel <beyselege at gmail.com>
---
.../Linalg/TransformOps/LinalgTransformOps.td | 36 ++++++++++--
.../include/mlir/Interfaces/TilingInterface.h | 19 +++++++
.../TransformOps/LinalgTransformOps.cpp | 36 +++++++++---
mlir/lib/Interfaces/TilingInterface.cpp | 56 +++++++++++++++++++
.../Dialect/Linalg/transform-ops-invalid.mlir | 36 ++++++++++++
mlir/test/Dialect/Linalg/transform-ops.mlir | 22 ++++++++
.../TestTilingInterfaceTransformOps.cpp | 21 ++++++-
.../TestTilingInterfaceTransformOps.td | 12 +++-
8 files changed, 220 insertions(+), 18 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 7f155c7f75d75..09a6c4c4d652f 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -432,6 +432,12 @@ def FuseOp : Op<Transform_Dialect, "structured.fuse",
If `apply_cleanup` is true then slice canonicalization is applied between
fusion steps. If `use_forall` is true then tiling method generates a
`scf.forall` loop instead of `scf.for` loops.
+
+ The optional `inner_tile_alignments` attribute forwards a per-dimension
+ `InnerTileAlignment` hint to the tiling driver for pack/unpack tiling.
+ These are only used to hint alignment or equality between the pack/unpack
+ operations inner tile sizes, and the loop tile sizes. No other operation
+ needs them.
}];
let arguments =
@@ -441,6 +447,7 @@ def FuseOp : Op<Transform_Dialect, "structured.fuse",
Optional<TransformAnyParamTypeOrAnyHandle> : $packed_tile_sizes,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_tile_sizes,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_tile_interchange,
+ DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$inner_tile_alignments,
UnitAttr:$apply_cleanup,
UnitAttr:$use_forall);
let results = (outs TransformHandleTypeInterface:$transformed,
@@ -477,6 +484,8 @@ def FuseOp : Op<Transform_Dialect, "structured.fuse",
$static_tile_sizes) |
`interchange` custom<DynamicIndexList>($tile_interchange, $static_tile_interchange)
)
+ (`inner_tile_alignments` `=`
+ custom<InnerTileAlignmentArray>($inner_tile_alignments)^)?
attr-dict `:` functional-type(operands, results)
}];
let hasVerifier = 1;
@@ -536,14 +545,24 @@ def FuseIntoContainingOp :
This operation consumes the producer handle.
This operation only reads the containing op handle.
+
+ The optional `inner_tile_alignments` attribute forwards a per-dimension
+ `InnerTileAlignment` hint to the tiling driver for pack/unpack tiling.
+ These are only used to hint alignment or equality between the pack/unpack
+ operations inner tile sizes, and the loop tile sizes. No other operation
+ needs them.
}];
let arguments = (ins TransformHandleTypeInterface:$producer_op,
- TransformHandleTypeInterface:$containing_op);
+ TransformHandleTypeInterface:$containing_op,
+ DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$inner_tile_alignments);
let results = (outs TransformHandleTypeInterface:$fused_op,
TransformHandleTypeInterface:$new_containing_op);
- let assemblyFormat = "$producer_op `into` $containing_op attr-dict "
- " `:` functional-type(operands, results)";
+ let assemblyFormat = "$producer_op `into` $containing_op "
+ "(`inner_tile_alignments` `=` "
+ "custom<InnerTileAlignmentArray>($inner_tile_alignments)^)? "
+ "attr-dict `:` functional-type(operands, results)";
+ let hasVerifier = 1;
let builders = [
OpBuilder<(ins "Value":$producerOp, "Value":$containingOp)>
@@ -2266,13 +2285,20 @@ def TileUsingForOp : Op<Transform_Dialect, "structured.tile_using_for",
If the internal implementation of tiling for any of the operations fails,
produces a definite failure.
+
+ The optional `inner_tile_alignments` attribute forwards a per-dimension
+ `InnerTileAlignment` hint to the tiling driver for pack/unpack tiling.
+ These are only used to hint alignment or equality between the pack/unpack
+ operations inner tile sizes, and the loop tile sizes. No other operation
+ needs them.
}];
let arguments = (ins TransformHandleTypeInterface:$target,
Variadic<TransformAnyParamTypeOrAnyHandle>:$dynamic_sizes,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sizes,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$interchange,
- DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:$scalable_sizes);
+ DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:$scalable_sizes,
+ DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$inner_tile_alignments);
let results = (outs TransformHandleTypeInterface:$tiled_linalg_op,
Variadic<TransformHandleTypeInterface>:$loops);
let builders = [
@@ -2307,6 +2333,8 @@ def TileUsingForOp : Op<Transform_Dialect, "structured.tile_using_for",
$static_sizes,
$scalable_sizes)
(`interchange` `=` $interchange^)?
+ (`inner_tile_alignments` `=`
+ custom<InnerTileAlignmentArray>($inner_tile_alignments)^)?
attr-dict
`:` functional-type(operands, results)
}];
diff --git a/mlir/include/mlir/Interfaces/TilingInterface.h b/mlir/include/mlir/Interfaces/TilingInterface.h
index 8b685a30868e8..d432d926647a2 100644
--- a/mlir/include/mlir/Interfaces/TilingInterface.h
+++ b/mlir/include/mlir/Interfaces/TilingInterface.h
@@ -17,6 +17,7 @@
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "mlir/Support/LLVM.h"
@@ -99,6 +100,24 @@ LogicalResult verifyInnerTileAlignments(Operation *op,
SmallVector<InnerTileAlignment>
convertInnerTileAlignments(ArrayRef<int64_t> alignments);
+/// Returns the keyword spelling of an `InnerTileAlignment` (`Unknown`,
+/// `Multiple` or `Equal`) used by the `inner_tile_alignments` assembly syntax.
+StringRef stringifyInnerTileAlignment(InnerTileAlignment alignment);
+
+/// Returns the `InnerTileAlignment` for a keyword spelling, or `std::nullopt`
+/// if `keyword` is not one of `Unknown`, `Multiple` or `Equal`.
+std::optional<InnerTileAlignment>
+symbolizeInnerTileAlignment(StringRef keyword);
+
+/// Custom directive parser/printer for an `inner_tile_alignments` attribute,
+/// rendering the `DenseI64ArrayAttr` as a keyword list, e.g.
+/// `[Equal, Multiple, Unknown]` (see `InnerTileAlignment`). Shared by the
+/// transform ops that carry the hint.
+ParseResult parseInnerTileAlignmentArray(OpAsmParser &parser,
+ DenseI64ArrayAttr &alignments);
+void printInnerTileAlignmentArray(OpAsmPrinter &printer, Operation *op,
+ DenseI64ArrayAttr alignments);
+
} // namespace mlir
/// Include the ODS generated interface header files.
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index f44693096b26b..61b81db1e6a74 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -657,6 +657,7 @@ void transform::FuseOp::build(OpBuilder &builder, OperationState &result,
/*packed_tile_sizes=*/Value(),
/*static_tile_sizes=*/staticTileSizesAttr,
/*static_tile_interchange=*/staticTileInterchangeAttr,
+ /*inner_tile_alignments=*/ArrayRef<int64_t>{},
/*apply_cleanup=*/applyCleanup,
/*use_forall=*/useForall);
}
@@ -757,6 +758,10 @@ transform::FuseOp::apply(transform::TransformRewriter &rewriter,
tilingOptions = tilingOptions.setTileSizes(mixedTileSizes);
scf::SCFTileAndFuseOptions tileAndFuseOptions;
tileAndFuseOptions.tilingOptions = tilingOptions;
+ // Optional caller-asserted pack/unpack inner-tile alignment (see
+ // InnerTileAlignment).
+ tileAndFuseOptions.tilingOptions.setInnerTileAlignments(
+ convertInnerTileAlignments(getInnerTileAlignments()));
if (getApplyCleanup()) {
MLIRContext *context = rewriter.getContext();
@@ -821,7 +826,7 @@ LogicalResult transform::FuseOp::verify() {
if (numExpectedLoops != getNumResults() - 1)
return emitOpError() << "expects " << numExpectedLoops << " loop results";
- return success();
+ return verifyInnerTileAlignments(getOperation(), getInnerTileAlignments());
}
SmallVector<OpFoldResult> transform::FuseOp::getMixedTileSizes() {
@@ -991,7 +996,8 @@ static bool sameOrEquivalentIterArg(Value src, Value dst) {
/// results of the `containingOp` or nullptr if there are no dominated uses.
static std::tuple<SmallVector<Operation *>, Operation *>
tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
- Operation *producerOp, Operation *containingOp) {
+ Operation *producerOp, Operation *containingOp,
+ ArrayRef<InnerTileAlignment> innerTileAlignments) {
LDBG() << "Try to fuse a direct extract use";
auto tileableProducer = dyn_cast<TilingInterface>(producerOp);
if (!tileableProducer) {
@@ -1067,7 +1073,7 @@ tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
FailureOr<TilingResult> tileAndFuseResult =
tileableProducer.generateResultTileValue(rewriter, resultNumber, offsets,
- sizes);
+ sizes, innerTileAlignments);
if (failed(tileAndFuseResult)) {
diag.attachNote(tileableProducer->getLoc())
@@ -1115,7 +1121,7 @@ tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
static SmallVector<Operation *>
tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
RewriterBase &rewriter, Diagnostic &diag, Operation *producerOp,
- Operation *containingOp) {
+ Operation *containingOp, ArrayRef<InnerTileAlignment> innerTileAlignments) {
LDBG() << "Try to fuse an extract use through block argument";
auto tileableProducer = dyn_cast<TilingInterface>(producerOp);
@@ -1192,7 +1198,7 @@ tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
FailureOr<TilingResult> tileAndFuseResult =
tileableProducerClone.generateResultTileValue(
rewriter, resultNumber, sliceOpToTile.getMixedOffsets(),
- sliceOpToTile.getMixedSizes());
+ sliceOpToTile.getMixedSizes(), innerTileAlignments);
if (failed(tileAndFuseResult)) {
diag.attachNote(tileableProducer->getLoc())
<< "failed to tile producer op: " << *tileableProducer;
@@ -1268,6 +1274,10 @@ bool transform::FuseIntoContainingOp::allowsRepeatedHandleOperands() {
return true;
}
+LogicalResult transform::FuseIntoContainingOp::verify() {
+ return verifyInnerTileAlignments(getOperation(), getInnerTileAlignments());
+}
+
DiagnosedSilenceableFailure
transform::FuseIntoContainingOp::apply(transform::TransformRewriter &rewriter,
transform::TransformResults &results,
@@ -1282,6 +1292,12 @@ transform::FuseIntoContainingOp::apply(transform::TransformRewriter &rewriter,
}
Operation *containingOp = *containingOps.begin();
+ // Forward the optional, caller-asserted alignment of a fused pack/unpack op's
+ // inner tiles relative to the loop tile sizes (see InnerTileAlignment) to
+ // each fused producer.
+ SmallVector<InnerTileAlignment> innerTileAlignments =
+ convertInnerTileAlignments(getInnerTileAlignments());
+
// If nothing to fuse, propagate success.
if (std::empty(producerOps)) {
results.set(cast<OpResult>(getFusedOp()), SmallVector<mlir::Operation *>{});
@@ -1332,8 +1348,8 @@ transform::FuseIntoContainingOp::apply(transform::TransformRewriter &rewriter,
// cases, we can tile/clone once and reuse the value for each use.
// Futhermore, producers should then be traversed according to a
// topological sorting.
- auto [tiledOps, newContainingOp] =
- tileAndFuseFirstExtractUse(rewriter, diag, producerOp, containingOp);
+ auto [tiledOps, newContainingOp] = tileAndFuseFirstExtractUse(
+ rewriter, diag, producerOp, containingOp, innerTileAlignments);
if (!tiledOps.empty()) {
LDBG() << "\nFused a direct extract use\n" << *containingOp;
fusedOps.append(tiledOps);
@@ -1359,7 +1375,7 @@ transform::FuseIntoContainingOp::apply(transform::TransformRewriter &rewriter,
SmallVector<Operation *> tiledContainingOpOperand =
tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
- rewriter, diag, producerOp, containingOp);
+ rewriter, diag, producerOp, containingOp, innerTileAlignments);
if (!tiledContainingOpOperand.empty()) {
LDBG() << "\nFused an extract use through block argument\n"
<< *containingOp;
@@ -3549,7 +3565,7 @@ LogicalResult transform::TileUsingForOp::verify() {
return emitOpError("expected number of loops to tile (")
<< numExpectedLoops << ") to match number of `loops` results ("
<< getLoops().size() << ")";
- return success();
+ return verifyInnerTileAlignments(getOperation(), getInnerTileAlignments());
}
DiagnosedSilenceableFailure
@@ -3679,6 +3695,8 @@ transform::TileUsingForOp::apply(transform::TransformRewriter &rewriter,
}
tilingOptions.setInterchange(getInterchange());
+ tilingOptions.setInnerTileAlignments(
+ convertInnerTileAlignments(getInnerTileAlignments()));
FailureOr<scf::SCFTilingResult> maybeTilingResult =
tileUsingSCF(rewriter, tilingInterface, tilingOptions);
if (failed(maybeTilingResult))
diff --git a/mlir/lib/Interfaces/TilingInterface.cpp b/mlir/lib/Interfaces/TilingInterface.cpp
index 4183d4fd59cd6..87e0af3b5cbc0 100644
--- a/mlir/lib/Interfaces/TilingInterface.cpp
+++ b/mlir/lib/Interfaces/TilingInterface.cpp
@@ -12,7 +12,9 @@
#include "mlir/Interfaces/TilingInterface.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVectorExtras.h"
+#include "llvm/ADT/StringSwitch.h"
using namespace mlir;
@@ -36,4 +38,58 @@ mlir::convertInnerTileAlignments(ArrayRef<int64_t> alignments) {
});
}
+StringRef mlir::stringifyInnerTileAlignment(InnerTileAlignment alignment) {
+ switch (alignment) {
+ case InnerTileAlignment::Unknown:
+ return "Unknown";
+ case InnerTileAlignment::Multiple:
+ return "Multiple";
+ case InnerTileAlignment::Equal:
+ return "Equal";
+ }
+ llvm_unreachable("unknown InnerTileAlignment");
+}
+
+std::optional<InnerTileAlignment>
+mlir::symbolizeInnerTileAlignment(StringRef keyword) {
+ return llvm::StringSwitch<std::optional<InnerTileAlignment>>(keyword)
+ .Case("Unknown", InnerTileAlignment::Unknown)
+ .Case("Multiple", InnerTileAlignment::Multiple)
+ .Case("Equal", InnerTileAlignment::Equal)
+ .Default(std::nullopt);
+}
+
+ParseResult mlir::parseInnerTileAlignmentArray(OpAsmParser &parser,
+ DenseI64ArrayAttr &alignments) {
+ SmallVector<int64_t> values;
+ auto parseEntry = [&]() -> ParseResult {
+ StringRef keyword;
+ llvm::SMLoc loc = parser.getCurrentLocation();
+ if (parser.parseKeyword(&keyword))
+ return failure();
+ std::optional<InnerTileAlignment> alignment =
+ symbolizeInnerTileAlignment(keyword);
+ if (!alignment)
+ return parser.emitError(loc)
+ << "expected one of 'Unknown', 'Multiple' or 'Equal', but got '"
+ << keyword << "'";
+ values.push_back(static_cast<int64_t>(*alignment));
+ return success();
+ };
+ if (parser.parseCommaSeparatedList(AsmParser::Delimiter::Square, parseEntry))
+ return failure();
+ alignments = DenseI64ArrayAttr::get(parser.getContext(), values);
+ return success();
+}
+
+void mlir::printInnerTileAlignmentArray(OpAsmPrinter &printer, Operation *,
+ DenseI64ArrayAttr alignments) {
+ printer << "[";
+ llvm::interleaveComma(alignments.asArrayRef(), printer, [&](int64_t value) {
+ printer << stringifyInnerTileAlignment(
+ static_cast<InnerTileAlignment>(value));
+ });
+ printer << "]";
+}
+
#include "mlir/Interfaces/TilingInterface.cpp.inc"
diff --git a/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir b/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir
index 6584596cdfdb2..5a6e0cc62b95c 100644
--- a/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir
@@ -100,3 +100,39 @@ transform.sequence failures(propagate) {
// expected-error at below {{expected '('}}
%res = transform.structured.generalize %arg0 : !transform.any_op -> !transform.any_op
}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{expected one of 'Unknown', 'Multiple' or 'Equal', but got 'Foo'}}
+ %1, %loop = transform.structured.tile_using_for %arg0 tile_sizes [8] inner_tile_alignments = [Foo] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{expected one of 'Unknown', 'Multiple' or 'Equal', but got 'Bogus'}}
+ %1, %loop = transform.structured.fuse %arg0 tile_sizes [8] inner_tile_alignments = [Bogus] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.match ops{["tensor.empty"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ // expected-error at below {{expected one of 'Unknown', 'Multiple' or 'Equal', but got 'Nope'}}
+ %fused, %new = transform.structured.fuse_into_containing_op %0 into %1 inner_tile_alignments = [Nope] : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ %0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ // expected-error at below {{expected one of 'Unknown', 'Multiple' or 'Equal', but got 'Xyz'}}
+ %a, %b = transform.test.fuse_consumer %1 into (%0) inner_tile_alignments = [Xyz] : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
diff --git a/mlir/test/Dialect/Linalg/transform-ops.mlir b/mlir/test/Dialect/Linalg/transform-ops.mlir
index 06a89fccd5c38..c92f29b441f72 100644
--- a/mlir/test/Dialect/Linalg/transform-ops.mlir
+++ b/mlir/test/Dialect/Linalg/transform-ops.mlir
@@ -16,6 +16,28 @@ transform.sequence failures(propagate) {
%2, %3:2 = transform.structured.tile_using_for %0 tile_sizes [0, 5, 3] {test_attr3 = 1 : i64, test_attr4}: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
}
+// Check that the `inner_tile_alignments` hint round-trips as a keyword list
+// (Unknown / Multiple / Equal); see `InnerTileAlignment`.
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+ // CHECK: transform.structured.tile_using_for %arg0 tile_sizes [8, 4] inner_tile_alignments = [Equal, Multiple]
+ %0, %1:2 = transform.structured.tile_using_for %arg0 tile_sizes [8, 4] inner_tile_alignments = [Equal, Multiple] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+ // CHECK: transform.structured.fuse %arg0 tile_sizes [8] inner_tile_alignments = [Multiple, Unknown]
+ %0, %1 = transform.structured.fuse %arg0 tile_sizes [8] inner_tile_alignments = [Multiple, Unknown] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.match ops{["tensor.empty"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ // CHECK: transform.structured.fuse_into_containing_op %{{.*}} into %{{.*}} inner_tile_alignments = [Equal, Equal]
+ %fused, %new = transform.structured.fuse_into_containing_op %0 into %1 inner_tile_alignments = [Equal, Equal] : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
transform.sequence failures(propagate) {
^bb1(%arg0: !transform.any_op):
%t = transform.structured.split %arg0 after 42 { dimension = 0 } : !transform.any_op
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
index b7cb810cda0e6..9467c925e543c 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
@@ -177,12 +177,20 @@ static LogicalResult
applyFuseConsumer(RewriterBase &rewriter, Operation *transformOp,
Operation *consumer,
MutableArrayRef<LoopLikeOpInterface> loops,
+ ArrayRef<InnerTileAlignment> innerTileAlignments,
TransformResults &transformResults) {
SmallVector<Operation *> fusedConsumerOps;
rewriter.setInsertionPoint(consumer);
+ // Drive the SCF helper with a constant control function returning the fixed
+ // `inner_tile_alignments` array for every op (consulted only by pack/unpack).
+ scf::InnerTileAlignmentFnTy alignmentFn =
+ [alignments = llvm::to_vector(innerTileAlignments)](
+ TilingInterface, ArrayRef<OpFoldResult>, ArrayRef<Operation *>) {
+ return alignments;
+ };
FailureOr<scf::SCFFuseConsumerOfSliceResult> fuseConsumerResults =
- scf::tileAndFuseConsumer(rewriter, consumer, loops);
+ scf::tileAndFuseConsumer(rewriter, consumer, loops, alignmentFn);
if (failed(fuseConsumerResults))
return consumer->emitOpError("failed to fuse consumer of slice");
@@ -198,6 +206,10 @@ applyFuseConsumer(RewriterBase &rewriter, Operation *transformOp,
return success();
}
+LogicalResult transform::TestFuseConsumerOp::verify() {
+ return verifyInnerTileAlignments(getOperation(), getInnerTileAlignments());
+}
+
DiagnosedSilenceableFailure
transform::TestFuseConsumerOp::apply(TransformRewriter &rewriter,
TransformResults &transformResults,
@@ -215,8 +227,11 @@ transform::TestFuseConsumerOp::apply(TransformRewriter &rewriter,
}
loops.push_back(loopLikeOp);
}
- LogicalResult result = applyFuseConsumer(rewriter, getOperation(), consumer,
- loops, transformResults);
+ SmallVector<InnerTileAlignment> innerTileAlignments =
+ convertInnerTileAlignments(getInnerTileAlignments());
+ LogicalResult result =
+ applyFuseConsumer(rewriter, getOperation(), consumer, loops,
+ innerTileAlignments, transformResults);
return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()
: DiagnosedSilenceableFailure::success();
}
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td
index 8c4f64de47795..efa16212f0fef 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td
@@ -89,18 +89,26 @@ def TestFuseConsumerOp : Op<Transform_Dialect, "test.fuse_consumer",
Returns a handle to the consumer operation after fusion and the loops that might be
modified.
+
+ The optional `inner_tile_alignments` array is supplied to the SCF driver as a
+ constant `scf::InnerTileAlignmentFnTy` control function (returning it for
+ every op, see `InnerTileAlignment`).
}];
- let arguments = (ins
+ let arguments = (ins
TransformHandleTypeInterface:$consumer,
- Variadic<TransformHandleTypeInterface>:$loops);
+ Variadic<TransformHandleTypeInterface>:$loops,
+ DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$inner_tile_alignments);
let results = (outs TransformHandleTypeInterface:$fused_consumer,
Variadic<TransformHandleTypeInterface>:$result_loops);
let assemblyFormat = [{
$consumer `into` `(` $loops `)`
+ (`inner_tile_alignments` `=`
+ custom<InnerTileAlignmentArray>($inner_tile_alignments)^)?
attr-dict `:` functional-type(operands, results)
}];
+ let hasVerifier = 1;
}
>From 87107fff669ebdb51b30abed273decd7dd980837 Mon Sep 17 00:00:00 2001
From: Ege Beysel <beyselege at gmail.com>
Date: Wed, 10 Jun 2026 17:07:06 +0000
Subject: [PATCH 4/4] test(mlir/linalg): cover scalable pack/unpack
tiling+fusion via InnerTileAlignment
Add tests for standalone tiling and producer/consumer fusion of scalable
linalg.pack/linalg.unpack driven by the inner_tile_alignments hint.
Signed-off-by: Ege Beysel <beyselege at gmail.com>
---
.../Linalg/scalable-pack-consumer-fusion.mlir | 373 ++++++++++++++++++
.../Linalg/scalable-pack-producer-fusion.mlir | 111 ++++++
.../Dialect/Linalg/scalable-pack-tiling.mlir | 87 ++++
.../scalable-unpack-consumer-fusion.mlir | 209 ++++++++++
.../scalable-unpack-producer-fusion.mlir | 364 +++++++++++++++++
.../Linalg/scalable-unpack-tiling.mlir | 122 ++++++
6 files changed, 1266 insertions(+)
create mode 100644 mlir/test/Dialect/Linalg/scalable-pack-consumer-fusion.mlir
create mode 100644 mlir/test/Dialect/Linalg/scalable-pack-producer-fusion.mlir
create mode 100644 mlir/test/Dialect/Linalg/scalable-pack-tiling.mlir
create mode 100644 mlir/test/Dialect/Linalg/scalable-unpack-consumer-fusion.mlir
create mode 100644 mlir/test/Dialect/Linalg/scalable-unpack-producer-fusion.mlir
create mode 100644 mlir/test/Dialect/Linalg/scalable-unpack-tiling.mlir
diff --git a/mlir/test/Dialect/Linalg/scalable-pack-consumer-fusion.mlir b/mlir/test/Dialect/Linalg/scalable-pack-consumer-fusion.mlir
new file mode 100644
index 0000000000000..de619d1fb9aa4
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/scalable-pack-consumer-fusion.mlir
@@ -0,0 +1,373 @@
+// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file --verify-diagnostics | FileCheck %s
+
+// Consumer fusion - linalg.pack with scalable inner tiles. Producer step (8*vscale)
+// equals the pack inner tile size(8*vscale) on the tiled source dimension, so the
+// outer dim of the fused pack tile is statically 1. This information is passed as
+// an inner tile alignment hint `Equal`.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func.func @fuse_scalable_pack_consumer_equal
+// CHECK-SAME: %[[ARG0:.+]]: tensor<256x128xf32>, %[[ARG1:.+]]: tensor<256x128xf32>, %[[ARG2:.+]]: tensor<256x128xf32>, %[[DEST:.+]]: tensor<?x?x?x?xf32>
+func.func @fuse_scalable_pack_consumer_equal(
+ %arg0: tensor<256x128xf32>, %arg1: tensor<256x128xf32>,
+ %arg2: tensor<256x128xf32>, %dest: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c256 = arith.constant 256 : index
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %c8_vscale = arith.muli %c8, %vscale : index
+
+ // Loop tile size is equal to the inner tile size of the consumer `linalg.pack` (8 * vscale).
+ %0 = scf.for %iv = %c0 to %c256 step %c8_vscale iter_args(%out = %arg2) -> (tensor<256x128xf32>) {
+ %sz = affine.min affine_map<(d0)[s0] -> (-d0 + 256, s0)>(%iv)[%c8_vscale]
+ %ext_out = tensor.extract_slice %out[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%ext_a, %ext_b : tensor<?x128xf32>, tensor<?x128xf32>)
+ outs(%ext_out : tensor<?x128xf32>) {
+ ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+ %mul = arith.mulf %in0, %in1 : f32
+ linalg.yield %mul : f32
+ } -> tensor<?x128xf32>
+ %inserted = tensor.insert_slice %computed into %out[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<?x128xf32> into tensor<256x128xf32>
+ scf.yield %inserted : tensor<256x128xf32>
+ }
+
+ %pack = linalg.pack %0 outer_dims_perm = [0, 1]
+ inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, %c4_vscale]
+ into %dest : tensor<256x128xf32> -> tensor<?x?x?x?xf32>
+ return %pack : tensor<?x?x?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %pack = transform.structured.match ops{["linalg.pack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %loop = transform.structured.match ops{["scf.for"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // The `Equal` hint is passed to hint the equality between the loop tile size 8 * vscale
+ // and the inner tile size 8 * vscale.
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) inner_tile_alignments = [Equal, Unknown]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // CHECK: %[[C8:.*]] = arith.constant 8 : index
+ // CHECK: %[[VSCALE:.*]] = vector.vscale
+ // CHECK: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+ // CHECK: %[[RES:.*]]:2 = scf.for {{.*}} step %[[C8_VSCALE]]
+ // CHECK-SAME: iter_args(%{{.*}} = %[[ARG2]], %{{.*}} = %[[DEST]])
+ // CHECK: %[[GENERIC:.*]] = linalg.generic
+ // CHECK: %[[PACK:.*]] = linalg.pack %[[GENERIC]]
+ // CHECK-SAME: inner_tiles = [%[[C8_VSCALE]], %{{.*}}]
+ // CHECK-SAME: -> tensor<1x?x?x?xf32>
+ // CHECK: scf.yield {{.*}}, %{{.*}} :
+ // CHECK: return %[[RES]]#1
+}
+
+// -----
+
+// Consumer fusion with a static producer step (64) and a scalable pack inner
+// tile (8*vscale), hinted `Multiple`. Fusion honors the hint and takes the aligned
+// (non-equal) path and the outer dim of the fused pack tile is dynamic (`64 ceildiv 8*vscale`).
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP_CEILDIV:.+]] = affine_map<()[s0] -> (64 ceildiv s0)>
+// CHECK-LABEL: func.func @fuse_scalable_pack_consumer_aligned
+// CHECK-SAME: %[[ARG0:.+]]: tensor<256x128xf32>, %[[ARG1:.+]]: tensor<256x128xf32>, %[[ARG2:.+]]: tensor<256x128xf32>, %[[DEST:.+]]: tensor<?x?x?x?xf32>
+func.func @fuse_scalable_pack_consumer_aligned(
+ %arg0: tensor<256x128xf32>, %arg1: tensor<256x128xf32>,
+ %arg2: tensor<256x128xf32>, %dest: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c64 = arith.constant 64 : index
+ %c256 = arith.constant 256 : index
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %c8_vscale = arith.muli %c8, %vscale : index
+
+ %0 = scf.for %iv = %c0 to %c256 step %c64 iter_args(%out = %arg2) -> (tensor<256x128xf32>) {
+ %ext_out = tensor.extract_slice %out[%iv, 0] [64, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<64x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [64, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<64x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [64, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<64x128xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%ext_a, %ext_b : tensor<64x128xf32>, tensor<64x128xf32>)
+ outs(%ext_out : tensor<64x128xf32>) {
+ ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+ %mul = arith.mulf %in0, %in1 : f32
+ linalg.yield %mul : f32
+ } -> tensor<64x128xf32>
+ %inserted = tensor.insert_slice %computed into %out[%iv, 0] [64, 128] [1, 1]
+ : tensor<64x128xf32> into tensor<256x128xf32>
+ scf.yield %inserted : tensor<256x128xf32>
+ }
+
+ %pack = linalg.pack %0 outer_dims_perm = [0, 1]
+ inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, %c4_vscale]
+ into %dest : tensor<256x128xf32> -> tensor<?x?x?x?xf32>
+ return %pack : tensor<?x?x?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %pack = transform.structured.match ops{["linalg.pack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %loop = transform.structured.match ops{["scf.for"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // The `Multiple` hint is passed to hint the alignment between the loop tile size 64
+ // and the inner tile size 8 * vscale.
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) inner_tile_alignments = [Multiple, Unknown]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+ // CHECK-DAG: %[[C64:.*]] = arith.constant 64 : index
+ // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale
+ // CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+ // CHECK-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+ // CHECK: %[[RES:.*]]:2 = scf.for {{.*}} step %[[C64]]
+ // CHECK-SAME: iter_args(%{{.*}} = %[[ARG2]], %{{.*}} = %[[DEST]])
+ // CHECK: %[[GENERIC:.*]] = linalg.generic
+ // CHECK: %[[OUTER:.*]] = affine.apply #[[$MAP_CEILDIV]]()[%[[C8_VSCALE]]]
+ // CHECK: %[[PACK_DEST:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, 0, 0, 0] [%[[OUTER]], %{{.*}}, %{{.*}}, %{{.*}}] [1, 1, 1, 1]
+ // CHECK: %[[PACK:.*]] = linalg.pack %[[GENERIC]]
+ // CHECK-SAME: inner_tiles = [%[[C8_VSCALE]], %[[C4_VSCALE]]]
+ // CHECK-SAME: into %[[PACK_DEST]]
+ // CHECK-SAME: -> tensor<?x?x?x?xf32>
+ // CHECK: scf.yield {{.*}}, %{{.*}} :
+ // CHECK: return %[[RES]]#1
+}
+
+// -----
+
+// Consumer fusion - both producer step and pack inner tile are scalable, step (8*vscale)
+// is an integer multiple of the inner tile (4*vscale) but not equal to it. The corresponding
+// `Multiple` hint is passed to the tiling interface. Fusion succeeds via the aligned (non-equal)
+// path, so the outer dim of the fused pack tile stays dynamic (`8*vscale ceildiv 4*vscale`).
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP_MIN:.+]] = affine_map<(d0)[s0] -> (-d0 + 256, s0)>
+// CHECK: #[[$MAP_CEILDIV:.+]] = affine_map<(d0)[s0] -> (d0 ceildiv s0)>
+// CHECK-LABEL: func.func @fuse_scalable_pack_consumer_aligned_scalable
+// CHECK-SAME: %[[ARG0:.+]]: tensor<256x128xf32>, %[[ARG1:.+]]: tensor<256x128xf32>, %[[ARG2:.+]]: tensor<256x128xf32>, %[[DEST:.+]]: tensor<?x?x?x?xf32>
+func.func @fuse_scalable_pack_consumer_aligned_scalable(
+ %arg0: tensor<256x128xf32>, %arg1: tensor<256x128xf32>,
+ %arg2: tensor<256x128xf32>, %dest: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c256 = arith.constant 256 : index
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %c8_vscale = arith.muli %c8, %vscale : index
+
+ // Loop tile size (8 * vscale) is a multiple of the inner tile size of the consumer `linalg.pack` (4 * vscale).
+ %0 = scf.for %iv = %c0 to %c256 step %c8_vscale iter_args(%out = %arg2) -> (tensor<256x128xf32>) {
+ %sz = affine.min affine_map<(d0)[s0] -> (-d0 + 256, s0)>(%iv)[%c8_vscale]
+ %ext_out = tensor.extract_slice %out[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%ext_a, %ext_b : tensor<?x128xf32>, tensor<?x128xf32>)
+ outs(%ext_out : tensor<?x128xf32>) {
+ ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+ %mul = arith.mulf %in0, %in1 : f32
+ linalg.yield %mul : f32
+ } -> tensor<?x128xf32>
+ %inserted = tensor.insert_slice %computed into %out[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<?x128xf32> into tensor<256x128xf32>
+ scf.yield %inserted : tensor<256x128xf32>
+ }
+
+ %pack = linalg.pack %0 outer_dims_perm = [0, 1]
+ inner_dims_pos = [0, 1] inner_tiles = [%c4_vscale, %c4_vscale]
+ into %dest : tensor<256x128xf32> -> tensor<?x?x?x?xf32>
+ return %pack : tensor<?x?x?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %pack = transform.structured.match ops{["linalg.pack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %loop = transform.structured.match ops{["scf.for"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // The `Multiple` hint is passed to hint the alignment between the loop tile size
+ // 8 * vscale and the inner tile size 4 * vscale.
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) inner_tile_alignments = [Multiple, Unknown]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+ // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale
+ // CHECK-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+ // CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+ // CHECK: %[[RES:.*]]:2 = scf.for {{.*}} step %[[C8_VSCALE]]
+ // CHECK-SAME: iter_args(%{{.*}} = %[[ARG2]], %{{.*}} = %[[DEST]])
+ // CHECK: %[[SZ:.*]] = affine.min #[[$MAP_MIN]](%{{.*}})[%[[C8_VSCALE]]]
+ // CHECK: %[[GENERIC:.*]] = linalg.generic
+ // CHECK: %[[OUTER:.*]] = affine.apply #[[$MAP_CEILDIV]](%[[SZ]])[%[[C4_VSCALE]]]
+ // CHECK: %[[PACK_DEST:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, 0, 0, 0] [%[[OUTER]], %{{.*}}, %{{.*}}, %{{.*}}] [1, 1, 1, 1]
+ // CHECK: %[[PACK:.*]] = linalg.pack %[[GENERIC]]
+ // CHECK-SAME: inner_tiles = [%[[C4_VSCALE]], %[[C4_VSCALE]]]
+ // CHECK-SAME: into %[[PACK_DEST]]
+ // CHECK-SAME: -> tensor<?x?x?x?xf32>
+ // CHECK: scf.yield {{.*}}, %{{.*}} :
+ // CHECK: return %[[RES]]#1
+}
+
+// -----
+
+// Consumer fusion (negative): linalg.pack with scalable inner tiles and no alignment hint.
+// The relationship between the loop tile size (8 * vscale) and the inner tile size (8 * vscale)
+// cannot be decided statically. Without a user hint that asserts this, fusion fails.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @negative_fuse_scalable_pack_consumer_no_hint(
+ %arg0: tensor<?x128xf32>, %arg1: tensor<?x128xf32>,
+ %arg2: tensor<?x128xf32>, %dest: tensor<?x?x?x?xf32>, %ub: index)
+ -> tensor<?x?x?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %c8_vscale = arith.muli %c8, %vscale : index
+
+ %0 = scf.for %iv = %c0 to %ub step %c8_vscale iter_args(%out = %arg2) -> (tensor<?x128xf32>) {
+ %sz = affine.min affine_map<(d0)[s0, s1] -> (s1 - d0, s0)>(%iv)[%c8_vscale, %ub]
+ %ext_out = tensor.extract_slice %out[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<?x128xf32> to tensor<?x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<?x128xf32> to tensor<?x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<?x128xf32> to tensor<?x128xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%ext_a, %ext_b : tensor<?x128xf32>, tensor<?x128xf32>)
+ outs(%ext_out : tensor<?x128xf32>) {
+ ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+ %mul = arith.mulf %in0, %in1 : f32
+ linalg.yield %mul : f32
+ } -> tensor<?x128xf32>
+ %inserted = tensor.insert_slice %computed into %out[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<?x128xf32> into tensor<?x128xf32>
+ scf.yield %inserted : tensor<?x128xf32>
+ }
+
+ // expected-error @below {{'linalg.pack' op failed to fuse consumer of slice}}
+ %pack = linalg.pack %0 outer_dims_perm = [0, 1]
+ inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, %c4_vscale]
+ into %dest : tensor<?x128xf32> -> tensor<?x?x?x?xf32>
+ return %pack : tensor<?x?x?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %pack = transform.structured.match ops{["linalg.pack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %loop = transform.structured.match ops{["scf.for"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // No inner_tile_alignments hint passed to signal the equality between the loop tile size
+ // 8 * vscale and the inner tile size 8 * vscale.
+ %a, %b = transform.test.fuse_consumer %pack into (%loop)
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Consumer fusion with a transposing outer_dims_perm. The hint is read in source
+// (operand) order, so `Equal` on source dim 0 (loop step 8*vscale == inner tile
+// 8*vscale) collapses that outer tile to 1; outer_dims_perm = [1, 0] then places
+// it at result position 1 (`tensor<?x1x?x?xf32>`), not position 0.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func.func @fuse_pack_consumer_transposed_outer
+func.func @fuse_pack_consumer_transposed_outer(
+ %arg0: tensor<256x128xf32>, %arg1: tensor<256x128xf32>,
+ %arg2: tensor<256x128xf32>, %dest: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c256 = arith.constant 256 : index
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %c8_vscale = arith.muli %c8, %vscale : index
+
+ %0 = scf.for %iv = %c0 to %c256 step %c8_vscale iter_args(%out = %arg2) -> (tensor<256x128xf32>) {
+ // 256 is not a static multiple of 8*vscale, so a real tiling clamps the
+ // per-iteration tile to affine.min(256 - iv, 8*vscale).
+ %sz = affine.min affine_map<(d0)[s0] -> (-d0 + 256, s0)>(%iv)[%c8_vscale]
+ %ext_out = tensor.extract_slice %out[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%ext_a, %ext_b : tensor<?x128xf32>, tensor<?x128xf32>)
+ outs(%ext_out : tensor<?x128xf32>) {
+ ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+ %mul = arith.mulf %in0, %in1 : f32
+ linalg.yield %mul : f32
+ } -> tensor<?x128xf32>
+ %inserted = tensor.insert_slice %computed into %out[%iv, 0] [%sz, 128] [1, 1]
+ : tensor<?x128xf32> into tensor<256x128xf32>
+ scf.yield %inserted : tensor<256x128xf32>
+ }
+
+ %pack = linalg.pack %0 outer_dims_perm = [1, 0]
+ inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, %c4_vscale]
+ into %dest : tensor<256x128xf32> -> tensor<?x?x?x?xf32>
+ return %pack : tensor<?x?x?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %pack = transform.structured.match ops{["linalg.pack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %loop = transform.structured.match ops{["scf.for"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // The `Equal` hint is passed to hint the equality between the loop tile size 8 * vscale
+ // and the inner tile size 8 * vscale.
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) inner_tile_alignments = [Equal, Unknown]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+ // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale
+ // CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+ // CHECK-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+ // CHECK: scf.for {{.*}} step %[[C8_VSCALE]]
+ // CHECK: linalg.pack
+ // CHECK-SAME: outer_dims_perm = [1, 0]
+ // CHECK-SAME: inner_tiles = [%[[C8_VSCALE]], %[[C4_VSCALE]]]
+ // CHECK-SAME: -> tensor<?x1x?x?xf32>
+ // CHECK: scf.yield
+}
diff --git a/mlir/test/Dialect/Linalg/scalable-pack-producer-fusion.mlir b/mlir/test/Dialect/Linalg/scalable-pack-producer-fusion.mlir
new file mode 100644
index 0000000000000..12544b04d8478
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/scalable-pack-producer-fusion.mlir
@@ -0,0 +1,111 @@
+// RUN: mlir-opt %s --transform-interpreter=entry-point=ones_equal -canonicalize -cse | FileCheck %s --check-prefixes=ONES_FUSED
+// RUN: mlir-opt %s --transform-interpreter=entry-point=ones_multiple -canonicalize -cse | FileCheck %s --check-prefixes=ONES_UNFUSED
+// RUN: mlir-opt %s --transform-interpreter=entry-point=ones_nohint -canonicalize -cse | FileCheck %s --check-prefixes=ONES_UNFUSED
+// RUN: mlir-opt %s --transform-interpreter=entry-point=inner_equal -canonicalize -cse | FileCheck %s --check-prefixes=INNER
+// RUN: mlir-opt %s --transform-interpreter=entry-point=inner_multiple -canonicalize -cse | FileCheck %s --check-prefixes=INNER
+// RUN: mlir-opt %s --transform-interpreter=entry-point=inner_nohint -canonicalize -cse | FileCheck %s --check-prefixes=INNER
+
+// Producer fusion of a scalable `linalg.pack` into a tiled elementwise consumer.
+// A pack producer can only be fused when the consumer requests full inner tiles.
+//
+// The `inner_tile_alignments` hint is a per-dimension keyword list:
+// - `Equal`: the loop tile size equals the pack/unpack inner tile size.
+// - `Multiple`: the loop tile size is an integer multiple of the inner tile.
+// - `Unknown`: the default; nothing is asserted for that dimension.
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func.func @pack_elemwise(%src: tensor<?x?xf32>, %pack_dest: tensor<?x?x?x?xf32>,
+ %out: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+ %c8 = arith.constant 8 : index
+ %c4 = arith.constant 4 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %pack = linalg.pack %src inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %pack_dest
+ : tensor<?x?xf32> -> tensor<?x?x?x?xf32>
+ %generic = linalg.generic {indexing_maps = [#map, #map],
+ iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+ ins(%pack : tensor<?x?x?x?xf32>) outs(%out : tensor<?x?x?x?xf32>) {
+ ^bb0(%in: f32, %o: f32):
+ %e = math.exp %in : f32
+ linalg.yield %e : f32
+ } -> tensor<?x?x?x?xf32>
+ return %generic : tensor<?x?x?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ // Tile the consumer's outer dims by 1. `Equal` -> the pack fuses.
+ transform.named_sequence @ones_equal(%arg1: !transform.any_op {transform.readonly}) {
+ %g = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %t, %loops:2 = transform.structured.fuse %g tile_sizes [1, 1, 0, 0] inner_tile_alignments = [Equal, Equal]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // ONES_FUSED-LABEL: func.func @pack_elemwise
+ // ONES_FUSED-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // ONES_FUSED: %[[RES:.*]] = scf.for {{.*}} step %[[C1]]
+ // ONES_FUSED: scf.for {{.*}} step %[[C1]]
+ // The pack is fused into the loop: it reads a source slice and produces a
+ // collapsed 1x1 outer tile that the generic consumes directly.
+ // ONES_FUSED: %[[SRC:.*]] = tensor.extract_slice %{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+ // ONES_FUSED: %[[PACK:.*]] = linalg.pack %[[SRC]]
+ // ONES_FUSED-SAME: -> tensor<1x1x?x?xf32>
+ // ONES_FUSED: linalg.generic {{.*}} ins(%[[PACK]]
+ // ONES_FUSED: return %[[RES]]
+
+ // Same input, `Multiple` hint: not sufficient, the pack is not fused.
+ transform.named_sequence @ones_multiple(%arg1: !transform.any_op {transform.readonly}) {
+ %g = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %t, %loops:2 = transform.structured.fuse %g tile_sizes [1, 1, 0, 0] inner_tile_alignments = [Multiple, Multiple]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // Same input, no hint: the pack is not fused.
+ transform.named_sequence @ones_nohint(%arg1: !transform.any_op {transform.readonly}) {
+ %g = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %t, %loops:2 = transform.structured.fuse %g tile_sizes [1, 1, 0, 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // ONES_UNFUSED-LABEL: func.func @pack_elemwise
+ // Without `Equal` the full-inner-tile requirement cannot be proven (scalable),
+ // so the pack stays outside the loop as a full pack and the tiled generic reads
+ // slices of its result.
+ // ONES_UNFUSED: %[[PACK:.*]] = linalg.pack
+ // ONES_UNFUSED-SAME: -> tensor<?x?x?x?xf32>
+ // ONES_UNFUSED: %[[RES:.*]] = scf.for
+ // ONES_UNFUSED: scf.for
+ // ONES_UNFUSED: tensor.extract_slice %[[PACK]]
+ // ONES_UNFUSED: linalg.generic
+ // ONES_UNFUSED: return %[[RES]]
+
+ // Tile the packed inner-tile dims. The pack is never fused (its inner tiles are not tileable),
+ // regardless of the hint - only the generic is tiled. All three sequences produce the same IR.
+
+ transform.named_sequence @inner_equal(%arg1: !transform.any_op {transform.readonly}) {
+ %g = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %t, %loops:2 = transform.structured.fuse %g tile_sizes [0, 0, 8, 4] inner_tile_alignments = [Equal, Equal]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.named_sequence @inner_multiple(%arg1: !transform.any_op {transform.readonly}) {
+ %g = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %t, %loops:2 = transform.structured.fuse %g tile_sizes [0, 0, 8, 4] inner_tile_alignments = [Multiple, Multiple]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.named_sequence @inner_nohint(%arg1: !transform.any_op {transform.readonly}) {
+ %g = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %t, %loops:2 = transform.structured.fuse %g tile_sizes [0, 0, 8, 4]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // INNER-LABEL: func.func @pack_elemwise
+ // INNER: %[[PACK:.*]] = linalg.pack
+ // INNER-SAME: -> tensor<?x?x?x?xf32>
+ // INNER: scf.for
+ // INNER: scf.for
+ // INNER: tensor.extract_slice %[[PACK]]
+ // INNER: linalg.generic
+}
diff --git a/mlir/test/Dialect/Linalg/scalable-pack-tiling.mlir b/mlir/test/Dialect/Linalg/scalable-pack-tiling.mlir
new file mode 100644
index 0000000000000..bc1e98fbc91c7
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/scalable-pack-tiling.mlir
@@ -0,0 +1,87 @@
+// RUN: mlir-opt %s --transform-interpreter=entry-point=ones_equal -canonicalize -cse | FileCheck %s --check-prefixes=ONES
+// RUN: mlir-opt %s --transform-interpreter=entry-point=ones_multiple -canonicalize -cse | FileCheck %s --check-prefixes=ONES
+// RUN: mlir-opt %s --transform-interpreter=entry-point=ones_nohint -canonicalize -cse | FileCheck %s --check-prefixes=ONES
+// RUN: mlir-opt %s --transform-interpreter=entry-point=twofour_equal -canonicalize -cse | FileCheck %s --check-prefixes=ALIGNED
+// RUN: mlir-opt %s --transform-interpreter=entry-point=twofour_multiple -canonicalize -cse | FileCheck %s --check-prefixes=ALIGNED
+// RUN: mlir-opt %s --transform-interpreter=entry-point=twofour_nohint -canonicalize -cse | FileCheck %s --check-prefixes=ALIGNED
+
+// Tiling of a scalable `linalg.pack`. A pack's tiling interface iterates over the
+// packed (destination) OUTER dims only, the inner tiles are not part of the
+// iteration domain - so tiling is completely independent of any
+// `inner_tile_alignments` hint.
+
+func.func @pack(%src: tensor<128x256xf32>, %dest: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+ %c8 = arith.constant 8 : index
+ %c4 = arith.constant 4 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %pack = linalg.pack %src inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %dest
+ : tensor<128x256xf32> -> tensor<?x?x?x?xf32>
+ return %pack : tensor<?x?x?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ // Tile the two outer dims by 1. Hint-independent, the three sequences below produce identical IR.
+ transform.named_sequence @ones_equal(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [1, 1] inner_tile_alignments = [Equal, Equal]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.named_sequence @ones_multiple(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [1, 1] inner_tile_alignments = [Multiple, Multiple]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.named_sequence @ones_nohint(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [1, 1]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // ONES-LABEL: func.func @pack
+ // ONES-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // ONES: %[[RES:.*]] = scf.for {{.*}} step %[[C1]]
+ // ONES: scf.for {{.*}} step %[[C1]]
+ // ONES: %[[PACK:.*]] = linalg.pack
+ // ONES-SAME: tensor<?x?xf32> -> tensor<1x1x?x?xf32>
+ // ONES: tensor.insert_slice %[[PACK]]
+ // ONES: scf.yield
+ // ONES: scf.yield
+ // ONES: return %[[RES]]
+
+ // Tile the two outer dims by 2 and 4. Hint-independent.
+ transform.named_sequence @twofour_equal(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] inner_tile_alignments = [Equal, Equal]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.named_sequence @twofour_multiple(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] inner_tile_alignments = [Multiple, Multiple]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.named_sequence @twofour_nohint(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // ALIGNED-LABEL: func.func @pack
+ // ALIGNED-DAG: %[[C2:.*]] = arith.constant 2 : index
+ // ALIGNED-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // ALIGNED: %[[RES:.*]] = scf.for {{.*}} step %[[C2]]
+ // ALIGNED: scf.for {{.*}} step %[[C4]]
+ // ALIGNED: %[[SRC:.*]] = tensor.extract_slice %{{.*}} : tensor<128x256xf32> to tensor<?x?xf32>
+ // ALIGNED: %[[PACK:.*]] = linalg.pack %[[SRC]]
+ // ALIGNED-SAME: tensor<?x?xf32> -> tensor<?x?x?x?xf32>
+ // ALIGNED: tensor.insert_slice %[[PACK]]
+ // ALIGNED: scf.yield
+ // ALIGNED: scf.yield
+ // ALIGNED: return %[[RES]]
+}
diff --git a/mlir/test/Dialect/Linalg/scalable-unpack-consumer-fusion.mlir b/mlir/test/Dialect/Linalg/scalable-unpack-consumer-fusion.mlir
new file mode 100644
index 0000000000000..f0605d36780ba
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/scalable-unpack-consumer-fusion.mlir
@@ -0,0 +1,209 @@
+// RUN: mlir-opt %s -transform-interpreter -split-input-file -canonicalize -cse | FileCheck %s --check-prefixes=CHECK
+// RUN: not mlir-opt %s --transform-interpreter=entry-point=no_hint -split-input-file 2>&1 | FileCheck %s --check-prefixes=NO_HINT
+// RUN: not mlir-opt %s --transform-interpreter=entry-point=multiple -split-input-file 2>&1 | FileCheck %s --check-prefixes=MULTIPLE
+
+// Consumer fusion of a scalable `linalg.unpack` into an `scf.for` loop.
+// The `inner_tile_alignments` hint is a per-dimension keyword list:
+// - `Equal`: the loop tile size equals the pack/unpack inner tile size.
+// - `Multiple`: the loop tile size is an integer multiple of the inner tile.
+// - `Unknown`: the default; nothing is asserted for that dimension.
+
+// The loop tile size on the tiled (inner-tile) dimension is 8 * vscale, the same
+// scalable value as the unpack inner tile. The SAME input IR is fused three ways:
+// with the `Equal` hint, with the `Multiple` hint, and without a hint.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @fuse_scalable_unpack_consumer(
+ %arg0: tensor<32x?xf32>, %arg1: tensor<32x?xf32>,
+ %arg2: tensor<32x?xf32>) -> tensor<?xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %dim1 = tensor.dim %arg2, %c1 : tensor<32x?xf32>
+
+ // Loop tile size is equal to the inner tile size of the consumer
+ // `linalg.unpack` (8 * vscale).
+ %0 = scf.for %iv = %c0 to %dim1 step %c8_vscale iter_args(%out = %arg2) -> (tensor<32x?xf32>) {
+ %sz = affine.min affine_map<(d0)[s0, s1] -> (s1 - d0, s0)>(%iv)[%c8_vscale, %dim1]
+ %ext_a = tensor.extract_slice %arg0[0, %iv] [32, %sz] [1, 1]
+ : tensor<32x?xf32> to tensor<32x?xf32>
+ %ext_b = tensor.extract_slice %arg1[0, %iv] [32, %sz] [1, 1]
+ : tensor<32x?xf32> to tensor<32x?xf32>
+ %ext_out = tensor.extract_slice %out[0, %iv] [32, %sz] [1, 1]
+ : tensor<32x?xf32> to tensor<32x?xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%ext_a, %ext_b : tensor<32x?xf32>, tensor<32x?xf32>)
+ outs(%ext_out : tensor<32x?xf32>) {
+ ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+ %mul = arith.mulf %in0, %in1 : f32
+ linalg.yield %mul : f32
+ } -> tensor<32x?xf32>
+ %inserted = tensor.insert_slice %computed into %out[0, %iv] [32, %sz] [1, 1]
+ : tensor<32x?xf32> into tensor<32x?xf32>
+ scf.yield %inserted : tensor<32x?xf32>
+ }
+
+ %output = tensor.empty(%dim1) : tensor<?xf32>
+ %unpack = linalg.unpack %0 outer_dims_perm = [0]
+ inner_dims_pos = [0] inner_tiles = [%c8_vscale]
+ into %output : tensor<32x?xf32> -> tensor<?xf32>
+ return %unpack : tensor<?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ // Intended fusion: the loop tile equals the inner tile, asserted via `Equal`.
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %unpack = transform.structured.match ops{["linalg.unpack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %loop = transform.structured.match ops{["scf.for"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %a, %b = transform.test.fuse_consumer %unpack into (%loop) inner_tile_alignments = [Equal]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // CHECK-LABEL: func.func @fuse_scalable_unpack_consumer
+ // CHECK-SAME: %[[ARG0:.+]]: tensor<32x?xf32>, %[[ARG1:.+]]: tensor<32x?xf32>, %[[ARG2:.+]]: tensor<32x?xf32>
+ // CHECK: %[[VSCALE:.*]] = vector.vscale
+ // CHECK: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %{{.*}} : index
+ // CHECK: %[[RES:.*]]:2 = scf.for {{.*}} step %[[C8_VSCALE]]
+ // CHECK-SAME: iter_args(%{{.*}} = %[[ARG2]], %{{.*}} = %{{.*}})
+ // CHECK: %[[GENERIC:.*]] = linalg.generic
+ // CHECK: %[[UNPACK:.*]] = linalg.unpack %[[GENERIC]]
+ // CHECK-SAME: inner_tiles = [%[[C8_VSCALE]]]
+ // CHECK: scf.yield {{.*}}, %{{.*}} :
+ // CHECK: return %[[RES]]#1
+
+ // No hint: fusion fails (scalable loop tile vs scalable inner tile, undecidable).
+ transform.named_sequence @no_hint(%arg1: !transform.any_op {transform.readonly}) {
+ %unpack = transform.structured.match ops{["linalg.unpack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %loop = transform.structured.match ops{["scf.for"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %a, %b = transform.test.fuse_consumer %unpack into (%loop)
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // NO_HINT: 'linalg.unpack' op failed to fuse consumer of slice
+
+ // `Multiple` hint: insufficient for an unpack consumer's inner dim, fusion fails.
+ transform.named_sequence @multiple(%arg1: !transform.any_op {transform.readonly}) {
+ %unpack = transform.structured.match ops{["linalg.unpack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %loop = transform.structured.match ops{["scf.for"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %a, %b = transform.test.fuse_consumer %unpack into (%loop) inner_tile_alignments = [Multiple]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // MULTIPLE: 'linalg.unpack' op failed to fuse consumer of slice
+}
+
+// -----
+
+// A 3-op dispatch tiled from the mmt4d root, fusing its consumers, with an
+// inner-tile alignment hint driving the scalable unpack fusion:
+//
+// linalg.mmt4d (root, produces a packed [M, N, M0, N0] layout)
+// -> linalg.generic (transposing bias add: [M, N, M0, N0] -> [N, M, N0, M0])
+// -> linalg.unpack ([N, M, N0, M0] -> [N*N0, M*M0])
+//
+// The mmt4d is tiled along its scalable inner dim N0 (iteration dim 4) by
+// 8 * vscale; the generic and unpack are fused as consumers. The SAME input IR
+// is fused two ways: with the equal hint and without.
+
+#id4 = affine_map<(m, n, m0, n0) -> (m, n, m0, n0)>
+#tr4 = affine_map<(m, n, m0, n0) -> (n, m, n0, m0)>
+
+// CHECK: #[[$TR:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3, d2)>
+func.func @mmt4d_transpose_unpack(
+ %lhs: tensor<2x2x4x2xf32>, %rhs: tensor<2x2x?x2xf32>,
+ %acc: tensor<2x2x4x?xf32>, %bias: tensor<2x2x?x4xf32>,
+ %trinit: tensor<2x2x?x4xf32>, %out: tensor<?x8xf32>) -> tensor<?x8xf32> {
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+
+ // 1. mmt4d root: lhs[M,K,M0,K0] x rhs[N,K,N0,K0] -> out[M,N,M0,N0].
+ %mm = linalg.mmt4d ins(%lhs, %rhs : tensor<2x2x4x2xf32>, tensor<2x2x?x2xf32>)
+ outs(%acc : tensor<2x2x4x?xf32>) -> tensor<2x2x4x?xf32>
+
+ // 2. transposing bias add: [M,N,M0,N0] -> [N,M,N0,M0].
+ %tr = linalg.generic {
+ indexing_maps = [#id4, #tr4, #tr4],
+ iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+ ins(%mm, %bias : tensor<2x2x4x?xf32>, tensor<2x2x?x4xf32>)
+ outs(%trinit : tensor<2x2x?x4xf32>) {
+ ^bb0(%a: f32, %b: f32, %o: f32):
+ %s = arith.addf %a, %b : f32
+ linalg.yield %s : f32
+ } -> tensor<2x2x?x4xf32>
+
+ // 3. unpack: [N,M,N0,M0] -> [N*N0, M*M0] = [?, 8].
+ %unpack = linalg.unpack %tr outer_dims_perm = [0, 1] inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, 4] into %out
+ : tensor<2x2x?x4xf32> -> tensor<?x8xf32>
+ return %unpack : tensor<?x8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ // Intended fusion: tile the mmt4d root, fuse the transposing generic and the
+ // unpack, asserting `Equal` on the unpack's tiled (transposed) inner dim.
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // Tile the mmt4d root along its scalable inner dim N0 (iteration dim 4).
+ %tiled, %loop = transform.structured.tile_using_for %mmt4d tile_sizes [0, 0, 0, 0, [8], 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %gen = transform.structured.match ops{["linalg.generic"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %fg, %loop2 = transform.test.fuse_consumer %gen into (%loop)
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %unp = transform.structured.match ops{["linalg.unpack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %fu, %loop3 = transform.test.fuse_consumer %unp into (%loop2) inner_tile_alignments = [Equal, Unknown]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // CHECK-LABEL: func.func @mmt4d_transpose_unpack
+ // CHECK-SAME: %[[LHS:.+]]: tensor<2x2x4x2xf32>, %[[RHS:.+]]: tensor<2x2x?x2xf32>, %[[ACC:.+]]: tensor<2x2x4x?xf32>, %[[BIAS:.+]]: tensor<2x2x?x4xf32>, %[[TRINIT:.+]]: tensor<2x2x?x4xf32>, %[[OUT:.+]]: tensor<?x8xf32>
+ // CHECK: %[[VSCALE:.*]] = vector.vscale
+ // CHECK: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %{{.*}} : index
+ // CHECK: %[[RES:.*]]:3 = scf.for {{.*}} step %[[C8_VSCALE]]
+ // CHECK-SAME: iter_args(%{{.*}} = %[[ACC]], %{{.*}} = %[[TRINIT]], %{{.*}} = %[[OUT]])
+ // CHECK: %[[MM:.*]] = linalg.mmt4d
+ // CHECK: %[[GENERIC:.*]] = linalg.generic
+ // CHECK-SAME: indexing_maps = [#{{.+}}, #[[$TR]], #[[$TR]]]
+ // CHECK-SAME: ins(%[[MM]],
+ // CHECK: %[[UNPACK:.*]] = linalg.unpack %[[GENERIC]]
+ // CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1]
+ // CHECK-SAME: inner_tiles = [%[[C8_VSCALE]], 4]
+ // CHECK: scf.yield {{.*}}, {{.*}}, %{{.*}} :
+ // CHECK: return %[[RES]]#2
+
+ // No hint on the unpack: unpack fusion fails.
+ transform.named_sequence @no_hint(%arg1: !transform.any_op {transform.readonly}) {
+ %mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_for %mmt4d tile_sizes [0, 0, 0, 0, [8], 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %gen = transform.structured.match ops{["linalg.generic"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %fg, %loop2 = transform.test.fuse_consumer %gen into (%loop)
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %unp = transform.structured.match ops{["linalg.unpack"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %fu, %loop3 = transform.test.fuse_consumer %unp into (%loop2)
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // NO_HINT: 'linalg.unpack' op failed to fuse consumer of slice
+
+ // Dummy entry point so the `multiple` RUN line resolves here.
+ transform.named_sequence @multiple(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/scalable-unpack-producer-fusion.mlir b/mlir/test/Dialect/Linalg/scalable-unpack-producer-fusion.mlir
new file mode 100644
index 0000000000000..96dc4937176fc
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/scalable-unpack-producer-fusion.mlir
@@ -0,0 +1,364 @@
+// RUN: mlir-opt %s --transform-interpreter=entry-point=aligned -split-input-file -canonicalize -cse | FileCheck %s --check-prefixes=ALIGNED
+// RUN: mlir-opt %s --transform-interpreter=entry-point=equal -split-input-file -canonicalize -cse | FileCheck %s --check-prefixes=EQUAL
+// RUN: mlir-opt %s --transform-interpreter=entry-point=unaligned -split-input-file -canonicalize -cse | FileCheck %s --check-prefixes=UNALIGNED
+// RUN: mlir-opt %s -transform-interpreter -split-input-file -canonicalize -cse | FileCheck %s --check-prefixes=CHECK
+
+// Producer fusion of a scalable `linalg.unpack` into a tiled consumer.
+// The `inner_tile_alignments` hint is a per-dimension keyword list:
+// - `Equal`: the loop tile size equals the pack/unpack inner tile size.
+// - `Multiple`: the loop tile size is an integer multiple of the inner tile.
+// - `Unknown`: the default; nothing is asserted for that dimension.
+
+// The elemwise consumer is tiled three ways from the same input IR: aligned, equal, and unaligned.
+
+// ALIGNED-DAG: #[[$MAP_CEILDIV:.+]] = affine_map<(d0)[s0] -> (d0 ceildiv s0)>
+func.func @unpack_elemwise_scalable(%arg0: tensor<4x8x?x?xf32>, %arg1: tensor<?x?xf32>, %arg2 : index, %arg3 : index) -> tensor<?x?xf32> {
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ // 16 * vscale is only consumed as a loop tile size by `@aligned`; `@equal` and
+ // `@unaligned` leave it dead and it is folded away.
+ %c16_vscale = arith.muli %c16, %vscale : index
+ %0 = tensor.empty(%arg2, %arg3) : tensor<?x?xf32>
+ %1 = linalg.unpack %arg0 inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %0
+ : tensor<4x8x?x?xf32> -> tensor<?x?xf32>
+ %2 = linalg.exp ins(%1: tensor<?x?xf32>)
+ outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %2 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ // Aligned: loop tile sizes (16 * vscale, 8 * vscale) are integer multiples of
+ // the inner tiles, asserted via `Multiple`. The outer dims become
+ // `ceilDiv(loop tile, inner tile)`.
+ transform.named_sequence @aligned(%arg1: !transform.any_op {transform.readonly}) {
+ %exp = transform.structured.match ops{["linalg.exp"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %mulis = transform.structured.match ops{["arith.muli"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %h8, %h4, %h16 = transform.split_handle %mulis : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ %tiled, %loops:2 = transform.structured.fuse %exp tile_sizes [%h16, %h8] interchange [0, 1]
+ inner_tile_alignments = [Multiple, Multiple]
+ : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // ALIGNED-LABEL: func.func @unpack_elemwise_scalable
+ // ALIGNED-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // ALIGNED-DAG: %[[C8:.*]] = arith.constant 8 : index
+ // ALIGNED-DAG: %[[C16:.*]] = arith.constant 16 : index
+ // ALIGNED-DAG: %[[VSCALE:.*]] = vector.vscale
+ // ALIGNED-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+ // ALIGNED-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+ // ALIGNED-DAG: %[[C16_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C16]] : index
+ // ALIGNED: %[[RES:.*]] = scf.for {{.*}} step %[[C16_VSCALE]]
+ // ALIGNED: scf.for {{.*}} step %[[C8_VSCALE]]
+ // ALIGNED: %[[OUTER_DIM0:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C8_VSCALE]]]
+ // ALIGNED: %[[OUTER_DIM1:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C4_VSCALE]]]
+ // ALIGNED: %[[SRC:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, %{{.*}}, 0, 0] [%[[OUTER_DIM0]], %[[OUTER_DIM1]], %[[C8_VSCALE]], %[[C4_VSCALE]]]
+ // ALIGNED: %[[UNPACK:.*]] = linalg.unpack %[[SRC]]
+ // ALIGNED-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ // ALIGNED-NOT: tensor.extract_slice %[[UNPACK]]
+ // ALIGNED: linalg.exp ins(%[[UNPACK]]
+ // ALIGNED: scf.yield
+ // ALIGNED: scf.yield
+ // ALIGNED: return %[[RES]]
+
+ // Equal: loop tile sizes (8 * vscale, 4 * vscale) equal the inner tiles,
+ // asserted via `Equal`; the fused unpack's outer dims collapse to 1.
+ transform.named_sequence @equal(%arg1: !transform.any_op {transform.readonly}) {
+ %exp = transform.structured.match ops{["linalg.exp"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %mulis = transform.structured.match ops{["arith.muli"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %h8, %h4, %h16 = transform.split_handle %mulis : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ %tiled, %loops:2 = transform.structured.fuse %exp tile_sizes [%h8, %h4] interchange [0, 1]
+ inner_tile_alignments = [Equal, Equal]
+ : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // EQUAL-LABEL: func.func @unpack_elemwise_scalable
+ // EQUAL-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // EQUAL-DAG: %[[C8:.*]] = arith.constant 8 : index
+ // EQUAL-DAG: %[[VSCALE:.*]] = vector.vscale
+ // EQUAL-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+ // EQUAL-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+ // EQUAL: %[[RES:.*]] = scf.for {{.*}} step %[[C8_VSCALE]]
+ // EQUAL: scf.for {{.*}} step %[[C4_VSCALE]]
+ // EQUAL: %[[UNPACK:.*]] = linalg.unpack
+ // EQUAL-SAME: tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+ // EQUAL-NOT: tensor.extract_slice %[[UNPACK]]
+ // EQUAL: linalg.exp ins(%[[UNPACK]]
+ // EQUAL: scf.yield
+ // EQUAL: scf.yield
+ // EQUAL: return %[[RES]]
+
+ // Unaligned: static tile sizes (7, 5) are not aligned to the scalable inner
+ // tiles and no hint is passed, so the fused unpack over-computes and a trailing
+ // extract_slice recovers the needed slice.
+ transform.named_sequence @unaligned(%arg1: !transform.any_op {transform.readonly}) {
+ %exp = transform.structured.match ops{["linalg.exp"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %tiled, %loops:2 = transform.structured.fuse %exp tile_sizes [7, 5] interchange [0, 1]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // UNALIGNED-LABEL: func.func @unpack_elemwise_scalable
+ // UNALIGNED-DAG: %[[C7:.*]] = arith.constant 7 : index
+ // UNALIGNED-DAG: %[[C5:.*]] = arith.constant 5 : index
+ // UNALIGNED: %[[RES:.*]] = scf.for {{.*}} step %[[C7]]
+ // UNALIGNED: scf.for {{.*}} step %[[C5]]
+ // UNALIGNED: %[[UNPACK:.*]] = linalg.unpack
+ // UNALIGNED-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ // UNALIGNED: %[[EXTRACT:.*]] = tensor.extract_slice %[[UNPACK]]
+ // UNALIGNED-NOT: linalg.exp ins(%[[UNPACK]]
+ // UNALIGNED: linalg.exp ins(%[[EXTRACT]]
+ // UNALIGNED: scf.yield
+ // UNALIGNED: scf.yield
+ // UNALIGNED: return %[[RES]]
+
+ // Dummy entry point so the default RUN line resolves here.
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+}
+
+// -----
+
+// Fusing a scalable `linalg.unpack` producer into a containing `scf.forall`. The
+// same input IR is fused three ways: aligned, equal, and unaligned.
+
+func.func @fuse_unpack_into_containing(
+ %src: tensor<?x?x?x?xf32>, %unpack_empty: tensor<?x?xf32>,
+ %out: tensor<?x?xf32>, %ub0: index, %ub1: index) -> tensor<?x?xf32> {
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ // Inner tile sizes of the unpack.
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ // Containing-loop tile sizes: numerically the same as the inner tiles, but
+ // deliberately distinct SSA values so the `@aligned` (`Multiple`) path cannot
+ // statically fold `ceilDiv(tile, inner)` to 1 and keeps a dynamic outer dim.
+ %c8_vscale_step = arith.muli %c8, %vscale : index
+ %c4_vscale_step = arith.muli %c4, %vscale : index
+ %unpack = linalg.unpack %src inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %unpack_empty
+ : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ // Loop tile sizes (8 * vscale, 4 * vscale) equal the inner tile sizes of the
+ // consumer `linalg.unpack` in value.
+ %res = scf.forall (%i, %j) = (0, 0) to (%ub0, %ub1) step (%c8_vscale_step, %c4_vscale_step)
+ shared_outs(%o = %out) -> (tensor<?x?xf32>) {
+ %slice = tensor.extract_slice %unpack[%i, %j] [%c8_vscale_step, %c4_vscale_step] [1, 1]
+ : tensor<?x?xf32> to tensor<?x?xf32>
+ %oslice = tensor.extract_slice %o[%i, %j] [%c8_vscale_step, %c4_vscale_step] [1, 1]
+ : tensor<?x?xf32> to tensor<?x?xf32>
+ %0 = linalg.exp ins(%slice : tensor<?x?xf32>) outs(%oslice : tensor<?x?xf32>) -> tensor<?x?xf32>
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %0 into %o[%i, %j] [%c8_vscale_step, %c4_vscale_step] [1, 1]
+ : tensor<?x?xf32> into tensor<?x?xf32>
+ }
+ }
+ return %res : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ // Equal tiles: the fused unpack collapses to a 1x1 outer tile.
+ transform.named_sequence @equal(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %fused, %newc = transform.structured.fuse_into_containing_op %0 into %1
+ inner_tile_alignments = [Equal, Equal]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // EQUAL-LABEL: func.func @fuse_unpack_into_containing
+ // EQUAL: scf.forall
+ // EQUAL: %[[UNPACK:.+]] = linalg.unpack
+ // EQUAL-SAME: : tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+ // EQUAL-NOT: tensor.extract_slice %[[UNPACK]]
+ // EQUAL: linalg.exp ins(%[[UNPACK]]
+
+ // No hint: general (unaligned) tiling with a trailing result slice.
+ transform.named_sequence @unaligned(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %fused, %newc = transform.structured.fuse_into_containing_op %0 into %1
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // UNALIGNED-LABEL: func.func @fuse_unpack_into_containing
+ // UNALIGNED: scf.forall
+ // UNALIGNED: %[[UNPACK:.+]] = linalg.unpack
+ // UNALIGNED-SAME: : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ // UNALIGNED: %[[EXTRACTED:.+]] = tensor.extract_slice %[[UNPACK]]
+ // UNALIGNED-NOT: linalg.exp ins(%[[UNPACK]]
+ // UNALIGNED: linalg.exp ins(%[[EXTRACTED]]
+
+ // Aligned: the `Multiple` hints assert the containing loop's tile sizes are
+ // multiples of the inner tiles. Although the loop tile size and inner tile sizes
+ // are equal but distinct SSA values, the equivalence cannot be inferred in the
+ // absence of the `Equal` hint.
+ transform.named_sequence @aligned(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %fused, %newc = transform.structured.fuse_into_containing_op %0 into %1
+ inner_tile_alignments = [Multiple, Multiple]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // ALIGNED-LABEL: func.func @fuse_unpack_into_containing
+ // ALIGNED: scf.forall
+ // The Multiple hint keeps a dynamic ceilDiv outer dim (step and inner tile are
+ // distinct SSA values, so it does not fold to 1) but needs no result slice.
+ // ALIGNED: %[[OUTER:.+]] = affine.apply {{.*}}[%{{.*}}, %{{.*}}]
+ // ALIGNED-NOT: linalg.unpack {{.*}} : tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+ // ALIGNED: %[[UNPACK:.+]] = linalg.unpack
+ // ALIGNED-SAME: : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ // ALIGNED-NOT: tensor.extract_slice %[[UNPACK]]
+ // ALIGNED: linalg.exp ins(%[[UNPACK]]
+
+ // Dummy entry point so the default RUN line resolves here.
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+}
+
+// -----
+
+// Producer fusion with a transposing consumer, `inner_tile_alignments` is
+// a caller-asserted hint indexed in the unpack's dest-dim order, so the
+// caller must arrange it for the transpose. Source dim 0 is `Multiple`
+// so -> outer size is `ceilDiv(loop tile,8 * vscale)` (dynamic);
+// source dim 1 is `Equal` -> outer size collapses to 1.
+
+// CHECK-DAG: #[[$MAP_CEILDIV:.+]] = affine_map<(d0)[s0] -> (d0 ceildiv s0)>
+func.func @unpack_transposed_consumer_scalable(%arg0: tensor<2x4x?x?xf32>, %out: tensor<?x?xf32>, %d0: index, %d1: index) -> tensor<?x?xf32> {
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %t0 = arith.muli %c4, %vscale : index
+ %t1 = arith.muli %c16, %vscale : index
+ %0 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
+ %unpack = linalg.unpack %arg0 inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %0
+ : tensor<2x4x?x?xf32> -> tensor<?x?xf32>
+ // Consumer reads %unpack transposed.
+ %1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%unpack : tensor<?x?xf32>) outs(%out : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %o: f32):
+ linalg.yield %in : f32
+ } -> tensor<?x?xf32>
+ return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func.func @unpack_transposed_consumer_scalable
+// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
+// CHECK-DAG: %[[VSCALE:.*]] = vector.vscale
+// CHECK-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+// CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+// CHECK-DAG: %[[C16_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C16]] : index
+// CHECK: %[[RES:.*]] = scf.for {{.*}} step %[[C4_VSCALE]]
+// CHECK: scf.for {{.*}} step %[[C16_VSCALE]]
+// CHECK: %[[OUTER_DIM0:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C8_VSCALE]]]
+// CHECK: %[[SRC:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, %{{.*}}, 0, 0] [%[[OUTER_DIM0]], 1, %{{.*}}, %{{.*}}] [1, 1, 1, 1]
+// CHECK-SAME: tensor<2x4x?x?xf32> to tensor<?x1x?x?xf32>
+// CHECK: %[[UNPACK:.*]] = linalg.unpack %[[SRC]]
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: linalg.generic
+// CHECK: scf.yield
+// CHECK: scf.yield
+// CHECK: return %[[RES]]
+
+module attributes {transform.with_named_sequence} {
+ // Hint indexed in the unpack's dest-dim order: dim0 Multiple, dim1 Equal.
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %gen = transform.structured.match ops{["linalg.generic"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %mulis = transform.structured.match ops{["arith.muli"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %i0, %i1, %h0, %h1 = transform.split_handle %mulis
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op,
+ !transform.any_op, !transform.any_op)
+ %tiled, %loops:2 = transform.structured.fuse %gen tile_sizes [%h0, %h1]
+ inner_tile_alignments = [Multiple, Equal]
+ : (!transform.any_op, !transform.any_op, !transform.any_op)
+ -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+
+ // Dummy entry points so the aligned/equal/unaligned RUN lines resolve here.
+ transform.named_sequence @aligned(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+ transform.named_sequence @equal(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+ transform.named_sequence @unaligned(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+}
+
+// -----
+
+// The hint also reaches the block-argument fusion path, when the producer is the
+// `scf.forall` init (used through the block argument rather than via a direct
+// extract use), `inner_tile_alignments` still yields the aligned tiling.
+
+func.func @fuse_unpack_through_block_arg(
+ %src: tensor<?x?x?x?xf32>, %unpack_empty: tensor<?x?xf32>,
+ %ub0: index, %ub1: index) -> tensor<?x?xf32> {
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %unpack = linalg.unpack %src inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %unpack_empty
+ : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ %res = scf.forall (%i, %j) = (0, 0) to (%ub0, %ub1) step (%c8_vscale, %c4_vscale)
+ shared_outs(%o = %unpack) -> (tensor<?x?xf32>) {
+ %slice = tensor.extract_slice %o[%i, %j] [%c8_vscale, %c4_vscale] [1, 1]
+ : tensor<?x?xf32> to tensor<?x?xf32>
+ %0 = linalg.exp ins(%slice : tensor<?x?xf32>) outs(%slice : tensor<?x?xf32>) -> tensor<?x?xf32>
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %0 into %o[%i, %j] [%c8_vscale, %c4_vscale] [1, 1]
+ : tensor<?x?xf32> into tensor<?x?xf32>
+ }
+ }
+ return %res : tensor<?x?xf32>
+}
+// CHECK-LABEL: func.func @fuse_unpack_through_block_arg
+// CHECK: scf.forall
+// CHECK: %[[UNPACK:.+]] = linalg.unpack
+// CHECK-SAME: : tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: linalg.exp ins(%[[UNPACK]]
+
+module attributes {transform.with_named_sequence} {
+ // The `Equal` hints assert the containing loop's tile sizes equal the unpack
+ // inner tile sizes (8 * vscale, 4 * vscale).
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %fused, %newc = transform.structured.fuse_into_containing_op %0 into %1
+ inner_tile_alignments = [Equal, Equal]
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+
+ // Dummy entry points so the aligned/equal/unaligned RUN lines resolve here.
+ transform.named_sequence @aligned(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+ transform.named_sequence @equal(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+ transform.named_sequence @unaligned(%arg1: !transform.any_op {transform.readonly}) {
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/scalable-unpack-tiling.mlir b/mlir/test/Dialect/Linalg/scalable-unpack-tiling.mlir
new file mode 100644
index 0000000000000..57bbc3af79499
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/scalable-unpack-tiling.mlir
@@ -0,0 +1,122 @@
+// RUN: mlir-opt %s --transform-interpreter=entry-point=equal -canonicalize -cse | FileCheck %s --check-prefixes=ALL,EQUAL
+// RUN: mlir-opt %s --transform-interpreter=entry-point=aligned -canonicalize -cse | FileCheck %s --check-prefixes=ALL,ALIGNED
+// RUN: mlir-opt %s --transform-interpreter=entry-point=unaligned -canonicalize -cse | FileCheck %s --check-prefixes=ALL,UNALIGNED
+
+// Pure tiling of a scalable `linalg.unpack` (no producer/consumer fusion). The
+// same input IR is tiled three ways, each selected by a different transform
+// entry point (see the named sequences and RUN lines):
+// - `@equal` tiles with [2, 4] * vscale, equal to the inner tiles, and
+// asserts that via `[Equal, Equal]`. Perfect tiling: the tiled
+// unpack's outer dims collapse to 1 and no remainder is needed.
+// - `@aligned` tiles with [4, 8] * vscale, integer multiples of the inner
+// tiles, asserted via `[Multiple, Multiple]`. The outer dims
+// become `ceilDiv(loop tile, inner tile)` and no remainder is
+// needed.
+// - `@unaligned` tiles with static [7, 5], not aligned to the scalable inner
+// tiles and with no hint, so the tiled unpack keeps a trailing
+// remainder `tensor.extract_slice`.
+//
+// The `inner_tile_alignments` hint is a per-dimension keyword list:
+// - `Equal`: the loop tile size equals the pack/unpack inner tile size.
+// - `Multiple`: the loop tile size is an integer multiple of the inner tile.
+// - `Unknown`: the default; nothing is asserted for that dimension.
+
+// ALIGNED-DAG: #[[$MAP_CEILDIV:.+]] = affine_map<(d0)[s0] -> (d0 ceildiv s0)>
+// ALL-LABEL: func.func @CKkc_to_KC_scalable
+
+func.func @CKkc_to_KC_scalable(%source: tensor<32x4x?x?xf32>, %dest: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %c2 = arith.constant 2 : index
+ %c4 = arith.constant 4 : index
+ %vscale = vector.vscale
+ %c2_vscale = arith.muli %c2, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
+ inner_tiles = [%c2_vscale, %c4_vscale] into %dest
+ : tensor<32x4x?x?xf32> -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ // Perfect tiling: the scalable tile sizes ([2] = 2 * vscale, [4] = 4 * vscale)
+ // equal the scalable inner tiles, asserted via `Equal`.
+ transform.named_sequence @equal(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [[2], [4]]
+ inner_tile_alignments = [Equal, Equal]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // --- @equal: perfect tiling, outer dims collapse to 1. ---
+ // EQUAL-DAG: %[[C2:.*]] = arith.constant 2 : index
+ // EQUAL-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // EQUAL-DAG: %[[VSCALE:.*]] = vector.vscale
+ // EQUAL-DAG: %[[C2_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C2]] : index
+ // EQUAL-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+ // EQUAL: %[[RES:.*]] = scf.for {{.*}} step %[[C2_VSCALE]]
+ // EQUAL: scf.for {{.*}} step %[[C4_VSCALE]]
+ // EQUAL: %[[UNPACK:.*]] = linalg.unpack
+ // EQUAL-SAME: tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+ // For the perfect tiling case, the outer dims equal 1 and the unpack can be
+ // consumed directly - no extra extract_slice on the result is needed.
+ // EQUAL-NOT: tensor.extract_slice %[[UNPACK]]
+ // EQUAL: tensor.insert_slice %[[UNPACK]]
+ // EQUAL: scf.yield
+ // EQUAL: scf.yield
+ // EQUAL: return %[[RES]]
+
+ // Aligned tiling: the scalable tile sizes ([4] = 4 * vscale, [8] = 8 * vscale)
+ // are integer multiples of the scalable inner tiles (2 * vscale, 4 * vscale),
+ // asserted via `Multiple`.
+ transform.named_sequence @aligned(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [[4], [8]]
+ inner_tile_alignments = [Multiple, Multiple]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // --- @aligned: multiple tiling, outer dims are ceilDiv(loop tile, inner tile). ---
+ // ALIGNED-DAG: %[[C2:.*]] = arith.constant 2 : index
+ // ALIGNED-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // ALIGNED-DAG: %[[VSCALE:.*]] = vector.vscale
+ // ALIGNED-DAG: %[[C2_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C2]] : index
+ // ALIGNED-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+ // ALIGNED-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %{{.*}} : index
+ // ALIGNED: %[[RES:.*]] = scf.for {{.*}} step %[[C4_VSCALE]]
+ // ALIGNED: scf.for {{.*}} step %[[C8_VSCALE]]
+ // The outer dims are ceilDiv(loop tile, inner tile); the transpose (outer_dims_perm
+ // = [1, 0]) swaps them back into source-dim order in the extracted slice.
+ // ALIGNED: %[[CEIL_C2:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C2_VSCALE]]]
+ // ALIGNED: %[[CEIL_C4:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C4_VSCALE]]]
+ // ALIGNED: %[[SRC:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, %{{.*}}, 0, 0] [%[[CEIL_C4]], %[[CEIL_C2]], %[[C2_VSCALE]], %[[C4_VSCALE]]]
+ // ALIGNED: %[[UNPACK:.*]] = linalg.unpack %[[SRC]]
+ // ALIGNED-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ // For the aligned tiling case, the unpack can also be consumed directly.
+ // ALIGNED-NOT: tensor.extract_slice %[[UNPACK]]
+ // ALIGNED: tensor.insert_slice %[[UNPACK]]
+ // ALIGNED: scf.yield
+ // ALIGNED: scf.yield
+ // ALIGNED: return %[[RES]]
+
+ // Unaligned tiling: static tile sizes (7, 5) are not aligned to the scalable
+ // inner tiles and no hint is passed, so the tiled unpack keeps a remainder.
+ transform.named_sequence @unaligned(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [7, 5]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ // --- @unaligned: static tiling, keeps a trailing remainder slice. ---
+ // UNALIGNED-DAG: %[[C7:.*]] = arith.constant 7 : index
+ // UNALIGNED-DAG: %[[C5:.*]] = arith.constant 5 : index
+ // UNALIGNED: %[[RES:.*]] = scf.for {{.*}} step %[[C7]]
+ // UNALIGNED: scf.for {{.*}} step %[[C5]]
+ // UNALIGNED: %[[UNPACK:.*]] = linalg.unpack
+ // UNALIGNED-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ // For the unaligned tiling case the unpack over-computes, so a trailing
+ // extract_slice recovers the correct slice from the result.
+ // UNALIGNED: %[[EXTRACT:.*]] = tensor.extract_slice %[[UNPACK]]
+ // UNALIGNED: tensor.insert_slice %[[EXTRACT]]
+ // UNALIGNED: scf.yield
+ // UNALIGNED: scf.yield
+ // UNALIGNED: return %[[RES]]
+}
More information about the Mlir-commits
mailing list