[Mlir-commits] [mlir] [mlir][linalg/scf/transform] scalable tiling and fusion for pack/unpack ops (PR #204007)
Ege Beysel
llvmlistbot at llvm.org
Wed Jun 24 02:48:23 PDT 2026
https://github.com/egebeysel updated https://github.com/llvm/llvm-project/pull/204007
>From 1bcb7d29f3d88c6406e0a20648d5f6063aba6c89 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. pack/unpack honor the hint only
when the relationship is not statically decidable; a hint that contradicts
statically known sizes is ignored.
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 | 51 ++++
.../mlir/Interfaces/TilingInterface.td | 111 ++++++++
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 24 ++
.../Linalg/Transforms/TilingInterfaceImpl.cpp | 237 +++++++++++++++---
.../Tensor/IR/TensorTilingInterfaceImpl.cpp | 4 +
mlir/lib/Interfaces/TilingInterface.cpp | 22 ++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 6 +
mlir/unittests/Dialect/Linalg/CMakeLists.txt | 1 +
8 files changed, 424 insertions(+), 32 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/TilingInterface.h b/mlir/include/mlir/Interfaces/TilingInterface.h
index 8693cbea7f0b0..66bd970767f3d 100644
--- a/mlir/include/mlir/Interfaces/TilingInterface.h
+++ b/mlir/include/mlir/Interfaces/TilingInterface.h
@@ -66,6 +66,57 @@ 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 inner tile.
+/// - `Equal`: the loop tile size equals the inner tile size.
+///
+/// This is a caller assertion, not a checked fact: it is only consulted when
+/// the relationship cannot be decided from the IR (e.g., scalable or dynamic
+/// sizes). When the tile and inner-tile sizes are both statically known,
+/// implementations trust that static comparison instead, so a hint that
+/// contradicts statically known sizes is ignored rather than allowed to produce
+/// incorrect tiling. The hint is otherwise never verified, so an incorrect
+/// assertion produces silently invalid tiling.
+///
+/// Entries are indexed by the dimensions the consulting method reasons about,
+/// i.e. the op's iteration domain (in pre-interchange order -- `interchange`
+/// reorders the generated loops only). This also holds for the
+/// `*FromOperandTiles` consumer-fusion methods: for a pack the iteration domain
+/// coincides with the unpacked operand's source dimensions, while for an unpack
+/// the entry for the i-th inner tile sits at its dest dimension
+/// `inner_dims_pos[i]` (a dimension of the unpacked tensor, not of the packed
+/// operand). Entries are not remapped through indexing maps or
+/// `outer_dims_perm` (for a transposing pack they stay in source order,
+/// pre-permutation), so the caller must pre-arrange them to match that order.
+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..d125313f7b84f 100644
--- a/mlir/include/mlir/Interfaces/TilingInterface.td
+++ b/mlir/include/mlir/Interfaces/TilingInterface.td
@@ -59,6 +59,23 @@ 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. As a
+ consequence:
+ - an `ExternalModel` that defines the hint-less overload should add a
+ `using Base::<method>;` declaration to keep the defaulted overload
+ visible;
+ - an op that lists `<method>` in `DeclareOpInterfaceMethods` has both
+ arities declared on it, so it must also define the hint-bearing
+ overload.
}];
let cppNamespace = "::mlir";
let methods = [
@@ -115,6 +132,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 +238,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 +290,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 +388,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/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 1a56c5a483e73..6bf705ef06c39 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2940,6 +2940,12 @@ SmallVector<utils::IteratorType> SoftmaxOp::getLoopIteratorTypes() {
return iteratorTypes;
}
+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,
@@ -3299,6 +3305,12 @@ LogicalResult WinogradFilterTransformOp::getResultTilePosition(
/// Users can specify the tile sizes of F and C.
/// `offsets` are the values for the offsets of F, KH, KW, C for one tile.
/// `sizes` are the values for the sizes of F, KH, KW, C for one tile.
+FailureOr<TilingResult> WinogradFilterTransformOp::getTiledImplementation(
+ OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
+ return getTiledImplementation(builder, offsets, sizes);
+}
+
FailureOr<TilingResult> WinogradFilterTransformOp::getTiledImplementation(
OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) {
@@ -3452,6 +3464,12 @@ LogicalResult WinogradInputTransformOp::getResultTilePosition(
/// C) Users can specify the tile sizes of tileH, tileW, N, and C. `offsets` are
/// the values for the offsets of tileH, tileW, N, C for one tile. `sizes` are
/// the values for the sizes of tileH, tileW, N, C for one tile.
+FailureOr<TilingResult> WinogradInputTransformOp::getTiledImplementation(
+ OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
+ return getTiledImplementation(builder, offsets, sizes);
+}
+
FailureOr<TilingResult>
WinogradInputTransformOp::getTiledImplementation(OpBuilder &builder,
ArrayRef<OpFoldResult> offsets,
@@ -3647,6 +3665,12 @@ LogicalResult WinogradOutputTransformOp::getResultTilePosition(
/// specify the tile sizes of tileH, tileW, N, and F. `offsets` are the values
/// for the offsets of tileH, tileW, N, F for one tile. `sizes` are the values
/// for the sizes of tileH, tileW, N, F for one tile.
+FailureOr<TilingResult> WinogradOutputTransformOp::getTiledImplementation(
+ OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
+ return getTiledImplementation(builder, offsets, sizes);
+}
+
FailureOr<TilingResult> WinogradOutputTransformOp::getTiledImplementation(
OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) {
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index 4eaa7bf0233c6..545177baf52eb 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -96,6 +96,15 @@ template <typename LinalgOpTy>
struct LinalgOpTilingInterface
: public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
LinalgOpTy> {
+ using Base =
+ TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
+ LinalgOpTy>;
+ // Inherit the defaulted hint-bearing overloads; this op ignores the hint.
+ 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);
@@ -943,6 +952,17 @@ struct PackOpTiling
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,
+ // TODO(egebeysel): support alignment hints for the pack producer tiling
+ // case. Currently, it is only consulted on the consumer-fusion paths
+ // (getIterationDomainTileFromOperandTiles / *FromOperandTiles).
+ ArrayRef<InnerTileAlignment>) const {
auto packOp = cast<PackOp>(op);
// TODO: Support Memref PackOp. Temporarily return failure.
if (!packOp.hasPureTensorSemantics())
@@ -1055,10 +1075,18 @@ 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))
@@ -1070,8 +1098,9 @@ struct PackOpTiling
if (!isEqualConstantIntOrValue(std::get<0>(iter), std::get<1>(iter)))
return failure();
- FailureOr<TilingResult> tilingResult = getTiledImplementation(
- op, b, offsets.drop_back(numTiles), sizes.drop_back(numTiles));
+ FailureOr<TilingResult> tilingResult =
+ getTiledImplementation(op, b, offsets.drop_back(numTiles),
+ sizes.drop_back(numTiles), innerTileAlignments);
if (failed(tilingResult))
return failure();
return tilingResult.value();
@@ -1123,15 +1152,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 +1203,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 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 +1218,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 +1253,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). It is only consulted when the relationship is
+ // not statically decidable; when both sizes are known constants we
+ // trust the static divisibility check below, so a contradicting hint
+ // cannot fuse into an invalid pack.
+ bool assumeInnerTileSizesMatchTiles = false;
+ bool staticallyDecidable =
+ !failed(cstTileSize) && cstInnerSize.has_value();
+ if (innerTileAlignment == InnerTileAlignment::Unknown ||
+ staticallyDecidable) {
+ if (failed(cstTileSize) || !cstInnerSize.has_value() ||
+ *cstTileSize % *cstInnerSize != 0)
+ return failure();
+ } else {
+ // The hint is consulted (the relationship is not statically
+ // decidable) and trusted as a caller assertion: an `Equal` assertion
+ // collapses the outer dim to a static 1 (loop tile == inner tile, so
+ // exactly one inner tile); a `Multiple` assertion takes the ceilDiv
+ // branch below.
+ assumeInnerTileSizesMatchTiles =
+ innerTileAlignment == InnerTileAlignment::Equal;
+ }
using AV = affine::AffineValueExpr;
affine::AffineBuilder ab(b, loc);
@@ -1212,7 +1287,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 +1303,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 +1345,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 +1384,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 +1417,33 @@ 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). It
+ // is only consulted when the relationship is not statically decidable (the
+ // scalable case the upper-bound computation above cannot handle); when both
+ // sizes are known constants we trust the static comparison below, so a hint
+ // that contradicts statically known sizes is ignored. When `Unknown`, fall
+ // back to the static upper-bound path below.
+ bool assumeInnerTileSizesMatchTiles = false;
+ bool staticallyDecidable = !failed(cstSize) && cstInnerSize.has_value();
+ if (innerTileAlignment != InnerTileAlignment::Unknown &&
+ !staticallyDecidable) {
+ info.isAlignedToInnerTileSize = true;
+ // The hint is consulted (the relationship is not statically decidable) and
+ // trusted as a caller assertion: an `Equal` assertion collapses the tiled
+ // source outer dim to a static 1; a `Multiple` assertion takes the ceilDiv
+ // branch below.
+ assumeInnerTileSizesMatchTiles =
+ innerTileAlignment == InnerTileAlignment::Equal;
+ }
+ 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);
@@ -1421,6 +1529,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 +1556,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 +1629,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();
@@ -1573,6 +1700,17 @@ struct UnPackOpTiling
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.
LogicalResult getIterationDomainTileFromOperandTiles(
@@ -1580,7 +1718,9 @@ struct UnPackOpTiling
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
SmallVectorImpl<OpFoldResult> &resultOffsets,
- SmallVectorImpl<OpFoldResult> &resultSizes) const {
+ SmallVectorImpl<OpFoldResult> &resultSizes,
+ ArrayRef<InnerTileAlignment> /*innerTileAlignments*/) const {
+ // This pure coordinate remapping does not consult the alignment hint.
if (operandNumbers.size() != 1) {
LLVM_DEBUG({ llvm::dbgs() << "unable to handle multiple operands"; });
return failure();
@@ -1642,11 +1782,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 +1809,36 @@ 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)))
+ 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) {
+ if (isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i]))
+ continue;
+ // If both sizes are statically known they are (given the check above)
+ // statically unequal: the inner dim is genuinely tiled and a
+ // contradicting Equal hint must not override that. Only honor the hint
+ // when equality is not statically decidable (e.g., scalable/dynamic
+ // sizes).
+ if (getConstantIntValue(mixedTiles[i]) &&
+ getConstantIntValue(innerSizes[i]))
return failure();
+ // `innerTileAlignments` is indexed by the unpack iteration domain (the
+ // dest dims); the i-th inner tile lives on dest dim `innerDimsPos[i]`,
+ // which is not necessarily a trailing dim. Only `Equal` is meaningful
+ // here: `Multiple` would mean the inner dim is tiled by a non-unit
+ // multiple of the inner tile, which is not fusible as an unpack consumer,
+ // so it falls through to the failure below.
+ int64_t destDim = innerDimsPos[i];
+ if (destDim < static_cast<int64_t>(innerTileAlignments.size()) &&
+ innerTileAlignments[destDim] == InnerTileAlignment::Equal)
+ continue;
+ return failure();
}
Location loc = unPackOp.getLoc();
@@ -1675,7 +1848,7 @@ struct UnPackOpTiling
SmallVector<OpFoldResult> outputOffsets, outputSizes;
if (failed(getIterationDomainTileFromOperandTiles(
op, b, operandNumbers, allOffsets, allSizes, outputOffsets,
- outputSizes)))
+ outputSizes, innerTileAlignments)))
return failure();
auto oneAttr = b.getI64IntegerAttr(1);
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
index 124a63281a37c..0b8fb252645ca 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
@@ -22,6 +22,10 @@ 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 ignores 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/mlir/unittests/Dialect/Linalg/CMakeLists.txt b/mlir/unittests/Dialect/Linalg/CMakeLists.txt
index a7da4e07c2551..df0cec31a4106 100644
--- a/mlir/unittests/Dialect/Linalg/CMakeLists.txt
+++ b/mlir/unittests/Dialect/Linalg/CMakeLists.txt
@@ -8,4 +8,5 @@ mlir_target_link_libraries(MLIRLinalgTests
MLIRFuncDialect
MLIRLinalgDialect
MLIRTensorDialect
+ MLIRTilingInterface
)
>From ed05d4c8350235d62fc18876768b96c1b4be4a11 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 | 73 +++++++++++++++++--
.../Dialect/Tensor/Transforms/Transforms.h | 12 +--
.../SCF/Transforms/TileUsingInterface.cpp | 63 +++++++++++-----
.../SwapExtractSliceWithProducerPatterns.cpp | 11 ++-
4 files changed, 126 insertions(+), 33 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
index 0005fad3d5c01..917911182a022 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
@@ -31,6 +31,19 @@ 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 to use (see
+/// `InnerTileAlignment`). The driver invokes it once per tiling/fusion of an
+/// op, so the returned array is always in *that op's own* domain -- the driver
+/// does not remap it. `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 +78,34 @@ struct SCFTilingOptions {
return *this;
}
+ /// Optional control function returning, per tiled/fused op, the caller
+ /// assertion of how each loop tile size relates to a tiled
+ /// `linalg.pack`/`linalg.unpack` op's inner tile size, in that op's own
+ /// iteration domain (see `InnerTileAlignmentFnTy` and `InnerTileAlignment`).
+ /// The driver invokes it once per tiling/fusion of an op. A null function
+ /// (the default) means "no information", in which case the relationship is
+ /// derived from the IR. Consulted only by pack/unpack implementations;
+ /// ignored by every other op.
+ InnerTileAlignmentFnTy innerTileAlignmentFn = nullptr;
+ SCFTilingOptions &setInnerTileAlignmentFn(InnerTileAlignmentFnTy fn) {
+ innerTileAlignmentFn = std::move(fn);
+ return *this;
+ }
+ /// Convenience setter installing 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); callers fusing several ops with
+ /// differing iteration domains should use `setInnerTileAlignmentFn` so each
+ /// op gets an array in its own domain.
+ 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 +328,27 @@ 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 (see
+/// `InnerTileAlignmentFnTy`). A null `fn` (the default) means "no hint".
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 +474,29 @@ 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 (see
+/// `InnerTileAlignmentFnTy`). A null `fn` (the default) means "no hint".
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 (see
+/// `InnerTileAlignmentFnTy`). A null `fn` (the default) means "no hint".
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..ff139c8269d01 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());
@@ -1189,11 +1191,17 @@ mlir::scf::tileUsingSCF(RewriterBase &rewriter, TilingInterface op,
return success();
}
- // 5c. Tile the cloned operation.
- tilingResult =
- getTiledImplementation(rewriter, clonedOp, options.reductionStrategy,
- regionIterArgs, tileOffsetsVec, tileSizesVec,
- ivs, numThreads, givenTileSizes, reductionDims);
+ // 5c. Tile the cloned operation. Resolve the inner-tile alignment hint for
+ // this op, in its own iteration domain, via the control function (if set).
+ 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 +1354,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 +1365,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 +1411,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 +1836,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 +2223,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 +2307,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 +2355,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 +2446,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 +2482,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 +2547,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 +2593,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 833618db571a95125e4be7b1e9f6c8744bcd0911 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 | 32 +++-
.../TransformOps/LinalgTransformOps.cpp | 46 ++++--
...nd-fuse-consumer-inner-tile-alignment.mlir | 145 ++++++++++++++++++
.../TestTilingInterfaceTransformOps.cpp | 21 ++-
.../TestTilingInterfaceTransformOps.td | 10 +-
5 files changed, 238 insertions(+), 16 deletions(-)
create mode 100644 mlir/test/Interfaces/TilingInterface/tile-and-fuse-consumer-inner-tile-alignment.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 7f155c7f75d75..745234432cc66 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -432,6 +432,14 @@ 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.
+ Entries are indexed in iteration-domain order; an empty array (the default)
+ preserves the prior behavior and the hint is ignored by ops that do not
+ consume it. The hint is a caller assertion: for genuinely scalable/dynamic
+ sizes it cannot be verified, so an incorrect entry silently produces invalid
+ tiling -- the caller owns its correctness.
}];
let arguments =
@@ -441,6 +449,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,
@@ -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 each fused producer for pack/unpack tiling.
+ Entries are indexed in the producer's iteration-domain order; an empty array
+ (the default) preserves the prior behavior. The hint is ignored by producers
+ that do not consume it. The hint is a caller assertion: for genuinely
+ scalable/dynamic sizes it cannot be verified, so an incorrect entry silently
+ produces invalid tiling -- the caller owns its correctness.
}];
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 hasVerifier = 1;
let builders = [
OpBuilder<(ins "Value":$producerOp, "Value":$containingOp)>
@@ -2266,13 +2285,22 @@ 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.
+ Entries are indexed in iteration-domain order; an empty array (the default)
+ preserves the prior behavior and the hint is ignored by ops that do not
+ consume it. The hint is a caller assertion: for genuinely scalable/dynamic
+ sizes it cannot be verified, so an incorrect entry silently produces invalid
+ tiling -- the caller owns its correctness.
}];
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 = [
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index f44693096b26b..2e329ddccbce8 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);
}
@@ -727,6 +728,10 @@ static LogicalResult applyTilingToAll(
return success();
}
+// `verifyInnerTileAlignments` and `convertInnerTileAlignments` are shared with
+// the test transform ops; they live next to `InnerTileAlignment` in
+// TilingInterface.h.
+
DiagnosedSilenceableFailure
transform::FuseOp::apply(transform::TransformRewriter &rewriter,
mlir::transform::TransformResults &transformResults,
@@ -757,6 +762,13 @@ 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). `setInnerTileAlignments` installs a constant control
+ // function returning this one array for every fused op; a caller needing a
+ // per-op hint sets `innerTileAlignmentFn` on the options instead (see
+ // tileConsumerAndFuseProducersUsingSCF).
+ tileAndFuseOptions.tilingOptions.setInnerTileAlignments(
+ convertInnerTileAlignments(getInnerTileAlignments()));
if (getApplyCleanup()) {
MLIRContext *context = rewriter.getContext();
@@ -821,7 +833,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 +1003,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 +1080,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 +1128,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 +1205,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 +1281,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 +1299,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 +1355,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 +1382,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 +3572,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 +3702,11 @@ transform::TileUsingForOp::apply(transform::TransformRewriter &rewriter,
}
tilingOptions.setInterchange(getInterchange());
+ // Optional caller-asserted per-dimension inner-tile alignment (see
+ // InnerTileAlignment), indexed in iteration-domain order, so pack/unpack
+ // tiling need not re-derive it from scalable/dynamic IR.
+ tilingOptions.setInnerTileAlignments(
+ convertInnerTileAlignments(getInnerTileAlignments()));
FailureOr<scf::SCFTilingResult> maybeTilingResult =
tileUsingSCF(rewriter, tilingInterface, tilingOptions);
if (failed(maybeTilingResult))
diff --git a/mlir/test/Interfaces/TilingInterface/tile-and-fuse-consumer-inner-tile-alignment.mlir b/mlir/test/Interfaces/TilingInterface/tile-and-fuse-consumer-inner-tile-alignment.mlir
new file mode 100644
index 0000000000000..d6530a4e286be
--- /dev/null
+++ b/mlir/test/Interfaces/TilingInterface/tile-and-fuse-consumer-inner-tile-alignment.mlir
@@ -0,0 +1,145 @@
+// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file --verify-diagnostics | FileCheck %s
+
+// 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 loop tile
+// (8*vscale) and the unpack inner tile (8*vscale) are both scalable, so their
+// relationship is not statically decidable and the unpack fusion needs the
+// `inner_tile_alignments` hint (supplied to the SCF driver as a control
+// function). The transposing bias add moves the scalable N0 inner tile onto the
+// unpack's dest dim 0, so `Equal` must sit at index 0 (`array<i64: 2, 0>`); the
+// generic ignores the hint. The negative case below shows the hint is
+// load-bearing.
+
+#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)>
+// 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
+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} {
+ 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)
+ // Fuse the transposing bias add (ignores the hint).
+ %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)
+ // Fuse the unpack. The transpose puts the tiled scalable N0 inner tile on
+ // dest dim 0, so Equal sits at index 0.
+ %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 = array<i64: 2, 0>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Negative (hint is load-bearing): same chain fused without any alignment hint.
+// The loop tile (8*vscale) and the unpack inner tile (8*vscale) are both
+// scalable, so the relationship is not statically decidable and fusion of the
+// unpack fails.
+
+#id4 = affine_map<(m, n, m0, n0) -> (m, n, m0, n0)>
+#tr4 = affine_map<(m, n, m0, n0) -> (n, m, n0, m0)>
+
+func.func @mmt4d_transpose_unpack_no_hint(
+ %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
+
+ %mm = linalg.mmt4d ins(%lhs, %rhs : tensor<2x2x4x2xf32>, tensor<2x2x?x2xf32>)
+ outs(%acc : tensor<2x2x4x?xf32>) -> tensor<2x2x4x?xf32>
+
+ %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>
+
+ // expected-error @below {{'linalg.unpack' op failed to fuse consumer of slice}}
+ %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} {
+ 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
+ %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
+ }
+}
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..659d24609c515 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td
@@ -89,11 +89,16 @@ 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);
@@ -101,6 +106,7 @@ def TestFuseConsumerOp : Op<Transform_Dialect, "test.fuse_consumer",
$consumer `into` `(` $loops `)`
attr-dict `:` functional-type(operands, results)
}];
+ let hasVerifier = 1;
}
>From 143b8da000c66b66f4dcb05e611a4d438a6a1876 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, including
negative cases where a hint that contradicts statically known sizes is ignored.
Signed-off-by: Ege Beysel <beyselege at gmail.com>
---
mlir/test/Dialect/Linalg/scalable-pack.mlir | 457 ++++++++++++
mlir/test/Dialect/Linalg/scalable-unpack.mlir | 692 ++++++++++++++++++
.../Dialect/Linalg/transform-ops-invalid.mlir | 36 +
3 files changed, 1185 insertions(+)
create mode 100644 mlir/test/Dialect/Linalg/scalable-pack.mlir
create mode 100644 mlir/test/Dialect/Linalg/scalable-unpack.mlir
diff --git a/mlir/test/Dialect/Linalg/scalable-pack.mlir b/mlir/test/Dialect/Linalg/scalable-pack.mlir
new file mode 100644
index 0000000000000..e5eb6047e6223
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/scalable-pack.mlir
@@ -0,0 +1,457 @@
+// 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.
+
+#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>
+// 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
+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
+
+ %0 = scf.for %iv = %c0 to %c256 step %c8_vscale iter_args(%out = %arg2) -> (tensor<256x128xf32>) {
+ %ext_out = tensor.extract_slice %out[%iv, 0] [%c8_vscale, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [%c8_vscale, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [%c8_vscale, 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] [%c8_vscale, 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
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) {inner_tile_alignments = array<i64: 2, 0>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Consumer fusion with a static producer step (64) and a scalable pack inner
+// tile (8*vscale), hinted `Multiple`. This is a caller assertion, the relationship
+// is not statically decidable, so fusion 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-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>
+// CHECK: %[[C64:.*]] = arith.constant 64 : index
+// CHECK: %[[RES:.*]]:2 = scf.for {{.*}} step %[[C64]]
+// CHECK-SAME: iter_args(%{{.*}} = %[[ARG2]], %{{.*}} = %[[DEST]])
+// CHECK: %[[GENERIC:.*]] = linalg.generic
+// CHECK: %[[PACK:.*]] = linalg.pack %[[GENERIC]]
+// CHECK-SAME: -> tensor<?x?x?x?xf32>
+// CHECK: scf.yield {{.*}}, %{{.*}} :
+// CHECK: return %[[RES]]#1
+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
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) {inner_tile_alignments = array<i64: 1, 0>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// 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. 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-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>
+// 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: -> tensor<?x?x?x?xf32>
+// CHECK: scf.yield {{.*}}, %{{.*}} :
+// CHECK: return %[[RES]]#1
+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
+
+ %0 = scf.for %iv = %c0 to %c256 step %c8_vscale iter_args(%out = %arg2) -> (tensor<256x128xf32>) {
+ %ext_out = tensor.extract_slice %out[%iv, 0] [%c8_vscale, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [%c8_vscale, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [%c8_vscale, 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] [%c8_vscale, 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
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) {inner_tile_alignments = array<i64: 1, 0>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Consumer fusion (negative): linalg.pack with scalable inner tiles and no
+// alignment hint. The loop step (12*vscale) and the inner tile (8*vscale) are
+// both dynamic, so without a hint their relationship is not statically decidable
+// and fusion falls through to failure.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @fuse_scalable_pack_consumer_mismatch(
+ %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
+ %c12 = arith.constant 12 : index
+ %c256 = arith.constant 256 : index
+ %vscale = vector.vscale
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c12_vscale = arith.muli %c12, %vscale : index
+
+ %0 = scf.for %iv = %c0 to %c256 step %c12_vscale iter_args(%out = %arg2) -> (tensor<256x128xf32>) {
+ %ext_out = tensor.extract_slice %out[%iv, 0] [%c12_vscale, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [%c12_vscale, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [%c12_vscale, 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] [%c12_vscale, 128] [1, 1]
+ : tensor<?x128xf32> into tensor<256x128xf32>
+ scf.yield %inserted : tensor<256x128xf32>
+ }
+
+ // 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<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
+ %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 (negative, static contradiction): the loop step (48) and the
+// pack inner tile (32) on the tiled dimension are both static and 48 is not a
+// multiple of 32. With both sizes static the divisibility check governs and the
+// contradicting `Multiple` hint is ignored, so fusion fails.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @fuse_pack_consumer_static_contradicting_hint(
+ %arg0: tensor<96x128xf32>, %arg1: tensor<96x128xf32>,
+ %arg2: tensor<96x128xf32>, %dest: tensor<?x?x32x4xf32>) -> tensor<?x?x32x4xf32> {
+ %c0 = arith.constant 0 : index
+ %c48 = arith.constant 48 : index
+ %c96 = arith.constant 96 : index
+
+ %0 = scf.for %iv = %c0 to %c96 step %c48 iter_args(%out = %arg2) -> (tensor<96x128xf32>) {
+ %ext_out = tensor.extract_slice %out[%iv, 0] [48, 128] [1, 1]
+ : tensor<96x128xf32> to tensor<48x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [48, 128] [1, 1]
+ : tensor<96x128xf32> to tensor<48x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [48, 128] [1, 1]
+ : tensor<96x128xf32> to tensor<48x128xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%ext_a, %ext_b : tensor<48x128xf32>, tensor<48x128xf32>)
+ outs(%ext_out : tensor<48x128xf32>) {
+ ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+ %mul = arith.mulf %in0, %in1 : f32
+ linalg.yield %mul : f32
+ } -> tensor<48x128xf32>
+ %inserted = tensor.insert_slice %computed into %out[%iv, 0] [48, 128] [1, 1]
+ : tensor<48x128xf32> into tensor<96x128xf32>
+ scf.yield %inserted : tensor<96x128xf32>
+ }
+
+ // 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 = [32, 4]
+ into %dest : tensor<96x128xf32> -> tensor<?x?x32x4xf32>
+ return %pack : tensor<?x?x32x4xf32>
+}
+
+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
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) {inner_tile_alignments = array<i64: 1, 0>}
+ : (!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
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[VSCALE:.*]] = vector.vscale
+// CHECK: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+// CHECK: scf.for {{.*}} step %[[C8_VSCALE]]
+// CHECK: linalg.pack
+// CHECK-SAME: outer_dims_perm = [1, 0]
+// CHECK-SAME: -> tensor<?x1x?x?xf32>
+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>) {
+ %ext_out = tensor.extract_slice %out[%iv, 0] [%c8_vscale, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_a = tensor.extract_slice %arg0[%iv, 0] [%c8_vscale, 128] [1, 1]
+ : tensor<256x128xf32> to tensor<?x128xf32>
+ %ext_b = tensor.extract_slice %arg1[%iv, 0] [%c8_vscale, 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] [%c8_vscale, 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
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) {inner_tile_alignments = array<i64: 2, 0>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Consumer fusion where the loop tile is a BOUNDED scalable size: 256 is not a
+// static multiple of 8*vscale, so the per-iteration tile is
+// affine.min(256 - iv, 8*vscale), whose value-bounds upper bound (256) equals
+// the source dim size. The dim must still be treated as tiled and the `Equal`
+// hint applied (outer dim collapses to a static 1) -- the bounded upper bound
+// must not make the scalable dim look untiled and bypass the hint.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func.func @fuse_pack_consumer_bounded_scalable_equal
+// CHECK: %[[PACK:.*]] = linalg.pack
+// CHECK-SAME: -> tensor<1x?x?x?xf32>
+func.func @fuse_pack_consumer_bounded_scalable_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
+
+ %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
+ // d0 tiled by 8*vscale == inner tile -> Equal; d1 untiled -> Unknown.
+ %a, %b = transform.test.fuse_consumer %pack into (%loop) {inner_tile_alignments = array<i64: 2, 0>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/scalable-unpack.mlir b/mlir/test/Dialect/Linalg/scalable-unpack.mlir
new file mode 100644
index 0000000000000..1b5fa8491ff5d
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/scalable-unpack.mlir
@@ -0,0 +1,692 @@
+// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file --verify-diagnostics | FileCheck %s
+
+// Perfect scalable tiling - scalable tile sizes equal scalable inner
+// tiles. Outer sizes of the tiled unpack should be 1's.
+
+// CHECK-LABEL: func.func @perfect_CKkc_to_KC_scalable
+// CHECK: %[[RES:.*]] = scf.for
+// CHECK: scf.for
+// CHECK: %[[UNPACK:.*]] = linalg.unpack
+// CHECK-SAME: tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: tensor.insert_slice %[[UNPACK]]
+// CHECK: return %[[RES]]
+func.func @perfect_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} {
+ 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, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [[2], [4]]
+ {inner_tile_alignments = array<i64: 2, 2>}
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Aligned scalable tiling - scalable tile sizes ([16] = 16*vscale, [8] =
+// 8*vscale) that are integer multiples of the scalable inner tiles (2x 8*vscale
+// and 2x 4*vscale), so tiling is aligned and the tiled unpack needs no trailing
+// remainder slice.
+
+// CHECK-LABEL: func.func @NCnc_to_NC_scalable_aligned
+// CHECK: %[[RES:.*]] = scf.for
+// CHECK: scf.for
+// CHECK: %[[UNPACK:.*]] = linalg.unpack
+// CHECK-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: tensor.insert_slice %[[UNPACK]]
+// CHECK: return %[[RES]]
+func.func @NCnc_to_NC_scalable_aligned(%source: tensor<4x8x?x?xf32>, %dest: tensor<?x?xf32>) -> tensor<?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
+ %0 = linalg.unpack %source inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %dest
+ : tensor<4x8x?x?xf32> -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [[16], [8]]
+ {inner_tile_alignments = array<i64: 1, 1>}
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Unaligned scalable tiling - static tile sizes not aligned to scalable
+// inner tiles.
+
+// CHECK-LABEL: func.func @NCnc_to_NC_scalable_unaligned
+// CHECK: %[[RES:.*]] = scf.for
+// CHECK: scf.for
+// CHECK: %[[UNPACK:.*]] = linalg.unpack
+// CHECK-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[UNPACK]]
+// CHECK: tensor.insert_slice %[[EXTRACT]]
+// CHECK: return %[[RES]]
+func.func @NCnc_to_NC_scalable_unaligned(%source: tensor<4x8x?x?xf32>, %dest: tensor<?x?xf32>) -> tensor<?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
+ %0 = linalg.unpack %source inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %dest
+ : tensor<4x8x?x?xf32> -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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, %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
+ }
+}
+
+// -----
+
+// Producer fusion - linalg.unpack with scalable inner tiles fused as a producer
+// into an elementwise consumer that is tiled with scalable tile
+// sizes that are an integer multiple of the inner tiles (16*vscale of 8*vscale,
+// 8*vscale of 4*vscale). `Multiple` hint is passed accordingly.
+
+// CHECK-LABEL: func.func @unpack_elemwise_scalable_multiple
+// CHECK: %[[RES:.*]] = scf.for
+// CHECK: scf.for
+// CHECK: %[[UNPACK:.*]] = linalg.unpack
+// Multiple (not Equal): outer dims stay dynamic (ceilDiv), not collapsed to 1.
+// CHECK-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: linalg.exp ins(%[[UNPACK]]
+// CHECK: return %[[RES]]
+func.func @unpack_elemwise_scalable_multiple(%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
+ %t0 = arith.muli %c16, %vscale : index
+ %t1 = arith.muli %c8, %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} {
+ transform.named_sequence @__transform_main(%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
+ %i0, %i1, %t0h, %t1h = 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 %exp tile_sizes [%t0h, %t1h] interchange [0, 1]
+ {inner_tile_alignments = array<i64: 1, 1>}
+ : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Producer fusion - same as above but the consumer is tiled with scalable tile
+// sizes equal to the unpack inner tiles (8*vscale, 4*vscale). The `Equal` hint
+// collapses the fused unpack's outer dims to 1 even though the tile-size and
+// inner-tile SSA values are not provably equal.
+
+// CHECK-LABEL: func.func @unpack_elemwise_scalable_equal
+// CHECK: %[[RES:.*]] = scf.for
+// CHECK: scf.for
+// CHECK: %[[UNPACK:.*]] = linalg.unpack
+// CHECK-SAME: tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: linalg.exp ins(%[[UNPACK]]
+// CHECK: return %[[RES]]
+func.func @unpack_elemwise_scalable_equal(%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
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %t0 = arith.muli %c8, %vscale : index
+ %t1 = arith.muli %c4, %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} {
+ transform.named_sequence @__transform_main(%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
+ %i0, %i1, %t0h, %t1h = 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 %exp tile_sizes [%t0h, %t1h] interchange [0, 1]
+ {inner_tile_alignments = array<i64: 2, 2>}
+ : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ 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.
+
+// CHECK-LABEL: func.func @unpack_transposed_consumer_scalable
+// CHECK: %[[RES:.*]] = scf.for
+// CHECK: scf.for
+// CHECK: %[[SRC:.*]] = tensor.extract_slice
+// 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: return %[[RES]]
+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>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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)
+ // Hint indexed in the unpack's dest-dim order: dim0 Multiple, dim1 Equal.
+ %tiled, %loops:2 = transform.structured.fuse %gen tile_sizes [%h0, %h1]
+ {inner_tile_alignments = array<i64: 1, 2>}
+ : (!transform.any_op, !transform.any_op, !transform.any_op)
+ -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Consumer fusion - linalg.unpack with scalable inner tiles fused as a
+// consumer into an scf.for loop. The loop step on the inner tile dimension
+// equals the unpack inner tile size (8*vscale), so fusion succeeds.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// 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
+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>
+
+ %0 = scf.for %iv = %c0 to %dim1 step %c8_vscale iter_args(%out = %arg2) -> (tensor<32x?xf32>) {
+ %extracted = tensor.extract_slice %out[0, %iv] [32, %c8_vscale] [1, 1]
+ : tensor<32x?xf32> to tensor<32x?xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0, %arg1 : tensor<32x?xf32>, tensor<32x?xf32>)
+ outs(%extracted : 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, %c8_vscale] [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} {
+ 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)
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Consumer fusion (negative) - linalg.unpack with scalable inner tiles and no
+// alignment hint. Without a hint equality is not statically provable
+// and fusion falls through to failure.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @fuse_scalable_unpack_consumer_mismatch(
+ %arg0: tensor<32x?xf32>, %arg1: tensor<32x?xf32>,
+ %arg2: tensor<32x?xf32>) -> tensor<?xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : 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
+ %dim1 = tensor.dim %arg2, %c1 : tensor<32x?xf32>
+
+ %0 = scf.for %iv = %c0 to %dim1 step %c4_vscale iter_args(%out = %arg2) -> (tensor<32x?xf32>) {
+ %extracted = tensor.extract_slice %out[0, %iv] [32, %c4_vscale] [1, 1]
+ : tensor<32x?xf32> to tensor<32x?xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0, %arg1 : tensor<32x?xf32>, tensor<32x?xf32>)
+ outs(%extracted : 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, %c4_vscale] [1, 1]
+ : tensor<32x?xf32> into tensor<32x?xf32>
+ scf.yield %inserted : tensor<32x?xf32>
+ }
+
+ %output = tensor.empty(%dim1) : tensor<?xf32>
+ // expected-error @below {{'linalg.unpack' op failed to fuse consumer of slice}}
+ %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} {
+ 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)
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Standalone tiling (static contradiction) - the loop tile sizes (7, 5) are
+// statically unequal to the inner tiles (8, 4), so an `Equal` hint must be
+// ignored and the unpack remainder slice must survive.
+
+// CHECK-LABEL: func.func @unpack_static_contradicting_equal_hint
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: %[[UNPACK:.*]] = linalg.unpack
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[UNPACK]]
+// CHECK: tensor.insert_slice %[[EXTRACT]]
+func.func @unpack_static_contradicting_equal_hint(%source: tensor<?x?x8x4xf32>, %dest: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [8, 4]
+ into %dest : tensor<?x?x8x4xf32> -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [7, 5]
+ {inner_tile_alignments = array<i64: 2, 2>}
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Consumer fusion (negative, static contradiction) - the loop step (4) and the
+// unpack inner tile (8) are statically unequal, so an `Equal` hint must be
+// ignored and the consumer cannot be fused.
+
+#map_static = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func.func @fuse_unpack_consumer_static_contradicting_equal_hint(
+ %arg0: tensor<4x16x8xf32>, %arg1: tensor<4x16x8xf32>,
+ %arg2: tensor<4x16x8xf32>) -> tensor<32x16xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+
+ %0 = scf.for %iv = %c0 to %c8 step %c4 iter_args(%out = %arg2) -> (tensor<4x16x8xf32>) {
+ %ext_out = tensor.extract_slice %out[0, 0, %iv] [4, 16, 4] [1, 1, 1]
+ : tensor<4x16x8xf32> to tensor<4x16x4xf32>
+ %ext_a = tensor.extract_slice %arg0[0, 0, %iv] [4, 16, 4] [1, 1, 1]
+ : tensor<4x16x8xf32> to tensor<4x16x4xf32>
+ %ext_b = tensor.extract_slice %arg1[0, 0, %iv] [4, 16, 4] [1, 1, 1]
+ : tensor<4x16x8xf32> to tensor<4x16x4xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map_static, #map_static, #map_static],
+ iterator_types = ["parallel", "parallel", "parallel"]}
+ ins(%ext_a, %ext_b : tensor<4x16x4xf32>, tensor<4x16x4xf32>)
+ outs(%ext_out : tensor<4x16x4xf32>) {
+ ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+ %mul = arith.mulf %in0, %in1 : f32
+ linalg.yield %mul : f32
+ } -> tensor<4x16x4xf32>
+ %inserted = tensor.insert_slice %computed into %out[0, 0, %iv] [4, 16, 4] [1, 1, 1]
+ : tensor<4x16x4xf32> into tensor<4x16x8xf32>
+ scf.yield %inserted : tensor<4x16x8xf32>
+ }
+
+ %output = tensor.empty() : tensor<32x16xf32>
+ // expected-error @below {{'linalg.unpack' op failed to fuse consumer of slice}}
+ %unpack = linalg.unpack %0 inner_dims_pos = [0] inner_tiles = [8]
+ into %output : tensor<4x16x8xf32> -> tensor<32x16xf32>
+ return %unpack : tensor<32x16xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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 = array<i64: 2, 0>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Consumer fusion (negative) - the loop step (16*vscale) is a non-unit multiple
+// of the unpack inner tile (8*vscale), so the inner dim is genuinely tiled. Only
+// `Equal` is meaningful for an unpack consumer's inner dim; a `Multiple` hint is
+// not sufficient, so fusion must still fail.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @fuse_scalable_unpack_consumer_multiple_hint_fails(
+ %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
+ %c16 = arith.constant 16 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c16_vscale = arith.muli %c16, %vscale : index
+ %dim1 = tensor.dim %arg2, %c1 : tensor<32x?xf32>
+
+ %0 = scf.for %iv = %c0 to %dim1 step %c16_vscale iter_args(%out = %arg2) -> (tensor<32x?xf32>) {
+ %extracted = tensor.extract_slice %out[0, %iv] [32, %c16_vscale] [1, 1]
+ : tensor<32x?xf32> to tensor<32x?xf32>
+ %computed = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0, %arg1 : tensor<32x?xf32>, tensor<32x?xf32>)
+ outs(%extracted : 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, %c16_vscale] [1, 1]
+ : tensor<32x?xf32> into tensor<32x?xf32>
+ scf.yield %inserted : tensor<32x?xf32>
+ }
+
+ %output = tensor.empty(%dim1) : tensor<?xf32>
+ // expected-error @below {{'linalg.unpack' op failed to fuse consumer of slice}}
+ %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} {
+ 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 = array<i64: 1>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Standalone tiling with a static loop tile and a scalable inner tile under an
+// `Equal` hint. The tile (8/4) and the inner tile (8*vscale/4*vscale) cannot be
+// related statically, so the hint is consulted and trusted as a caller
+// assertion: the unpack is aligned (no trailing remainder slice).
+
+// CHECK-LABEL: func.func @unpack_static_tile_scalable_inner_hint
+// CHECK: %[[RES:.*]] = scf.for
+// CHECK: scf.for
+// CHECK: %[[UNPACK:.*]] = linalg.unpack
+// CHECK-SAME: tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: tensor.insert_slice %[[UNPACK]]
+// CHECK: return %[[RES]]
+func.func @unpack_static_tile_scalable_inner_hint(
+ %source: tensor<4x8x?x?xf32>, %dest: tensor<?x?xf32>) -> 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
+ %0 = linalg.unpack %source inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %dest
+ : tensor<4x8x?x?xf32> -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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
+ // Static tile [8, 4], scalable inner [8*vscale, 4*vscale] -> not statically
+ // decidable; Equal hint -> aligned (ceilDiv), but bounded tile -> no collapse.
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [8, 4]
+ {inner_tile_alignments = array<i64: 2, 2>}
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Fusing a scalable `linalg.unpack` producer into a containing `scf.forall`
+// with equal loop and inner tile sizes.
+
+// CHECK-LABEL: func.func @fuse_unpack_into_containing_aligned
+// 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]]
+func.func @fuse_unpack_into_containing_aligned(
+ %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
+ %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 = %out) -> (tensor<?x?xf32>) {
+ %slice = tensor.extract_slice %unpack[%i, %j] [%c8_vscale, %c4_vscale] [1, 1]
+ : tensor<?x?xf32> to tensor<?x?xf32>
+ %oslice = 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(%oslice : 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>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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 = array<i64: 2, 2>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Same fusion WITHOUT the hint: the producer falls back to the general (unaligned)
+// tiling, which over-allocates the unpacked tile (`tensor<?x?x?x?xf32>` source) and
+// recovers the needed slice with a trailing `tensor.extract_slice` on the result.
+
+// CHECK-LABEL: func.func @fuse_unpack_into_containing_unaligned
+// CHECK: scf.forall
+// CHECK: %[[UNPACK:.+]] = linalg.unpack
+// CHECK-SAME: : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+// CHECK: tensor.extract_slice %[[UNPACK]]
+func.func @fuse_unpack_into_containing_unaligned(
+ %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
+ %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 = %out) -> (tensor<?x?xf32>) {
+ %slice = tensor.extract_slice %unpack[%i, %j] [%c8_vscale, %c4_vscale] [1, 1]
+ : tensor<?x?xf32> to tensor<?x?xf32>
+ %oslice = 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(%oslice : 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>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ 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.
+
+// CHECK-LABEL: func.func @fuse_unpack_through_block_arg_aligned
+// 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]]
+func.func @fuse_unpack_through_block_arg_aligned(
+ %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>
+}
+
+module attributes {transform.with_named_sequence} {
+ 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 = array<i64: 2, 2>}
+ : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir b/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir
index 6584596cdfdb2..b089a26ef5a40 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 inner_tile_alignments entries to be one of 0 (Unknown), 1 (Multiple) or 2 (Equal), but got 5}}
+ %1, %loop = transform.structured.tile_using_for %arg0 tile_sizes [8] {inner_tile_alignments = array<i64: 5>} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{expected inner_tile_alignments entries to be one of 0 (Unknown), 1 (Multiple) or 2 (Equal), but got -1}}
+ %1, %loop = transform.structured.fuse %arg0 tile_sizes [8] {inner_tile_alignments = array<i64: -1>} : (!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 inner_tile_alignments entries to be one of 0 (Unknown), 1 (Multiple) or 2 (Equal), but got 7}}
+ %fused, %new = transform.structured.fuse_into_containing_op %0 into %1 {inner_tile_alignments = array<i64: 7>} : (!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 inner_tile_alignments entries to be one of 0 (Unknown), 1 (Multiple) or 2 (Equal), but got 3}}
+ %a, %b = transform.test.fuse_consumer %1 into (%0) {inner_tile_alignments = array<i64: 3>} : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
More information about the Mlir-commits
mailing list