Ege Beysel
3824108316
# Inner tile alignment hints for scalable `linalg.pack`/`linalg.unpack` tiling and fusion ## Overview Tiling and fusing `linalg.pack`/`linalg.unpack` produces a clean result only when the tiling implementation can tell how a loop tile size relates to the op's inner tile size. When both are statically known this is decided by comparing the constants. But with **scalable** (and, more generally, dynamic) sizes, e.g. a loop tile of `8 * vscale` against an inner tile of `8 * vscale`, that relationship is symbolic and cannot cleanly be recovered from the IR, so the implementation conservatively falls back to a dynamic, over-allocated tile. See #150185 for more details. This PR adds an optional **inner tile alignment hint**: a per-dimension caller assertion about that relationship, threaded from the transform-op surface through the SCF driver into the pack/unpack `TilingInterface` implementations. The caller is already aware of both the loop tile sizes, and the inner tile sizes, and therefore can set these hints easily. With the hint, scalable pack/unpack tiling and fusion produce the same tight IR they produce for the statically-aligned case. ## The inner tile alignment hint `InnerTileAlignment` is a per-iteration-domain-dimension enum: - `Unknown` (`0`) — no information; the implementation keeps its prior behavior. - `Multiple` (`1`) — the loop tile size is an integer multiple of the inner tile. - `Equal` (`2`) — the loop tile size equals the inner tile size. It is a *caller assertion*, not a checked fact. It is consulted only when the relationship cannot be decided from the IR; when the tile and inner-tile sizes are both statically known, that static comparison wins and a contradicting hint is ignored rather than allowed to corrupt the result. Otherwise the hint is not verified, so an incorrect assertion silently produces invalid tiling, its correctness is owned by the caller. Only pack/unpack consult it; every other op ignores it. ## TilingInterface API Four methods gain a second overload taking an extra `ArrayRef<InnerTileAlignment>` hint: - `getTiledImplementation` - `generateResultTileValue` - `getTiledImplementationFromOperandTiles` - `getIterationDomainTileFromOperandTiles` The hint-bearing overload defaults to forwarding to the hint-less one, so the change is additive: existing implementers and callers that never mention the hint are unaffected, and the rest of the interface is untouched. ### Where it helps The loop tiles the packed dimension by `8 * vscale`, so the size handed to the pack/unpack is a **full-or-remainder** `affine.min` against that step: ```mlir // #map = affine_map<(d0)[s0] -> (-d0 + 4096, s0)> // min(4096 - iv, 8*vscale) %ts = affine.min #map(%iv)[%c8_vscale] // ... %ts feeds the extracted slice fed to the pack/unpack, whose inner tile // on this dimension is also 8*vscale. ``` To fold the tiled op's outer (tile-count) dimension to a static `1`, the current implementation asks `ValueBoundsConstraintSet::computeConstantBound` for a *constant* upper bound on `%ts` and compares it against the inner tile size, essentially `UB(%ts) % innerTile == 0`. With a scalable `8 * vscale` inner tile, neither the upper bound of `%ts` nor the inner tile is a constant, so the comparison is undecidable and the tile conservatively falls back to a dynamic outer dim (`tensor<?x...>`), over-allocating the tile. `inner_tile_alignments = [Equal]` asserts `%ts == 8 * vscale`, so the outer dim folds to a static `1` (`tensor<1x...>`); `[Multiple]` instead takes the `ceilDiv` branch. The hint thus supplies exactly the relationship the upper-bound comparison cannot recover for symbolic sizes. ### How the hint is arranged - One entry per **iteration-domain dimension**, in **pre-interchange** order (`interchange` only reorders the generated loops, not the hint). - Entries are **not** remapped through indexing maps or `outer_dims_perm`: for a transposing pack they stay in source (pre-permutation) order. - For the consumer-fusion (`*FromOperandTiles`) methods, the indexing follows the op's iteration domain: for a pack it coincides with the unpacked source dims; for an unpack the entry for the i-th inner tile sits at its destination dimension `inner_dims_pos[i]`. The caller must pre-arrange the array to match this order. See the tests for more examples. ### Out-of-tree implementers Because the two overloads share a source name: - An op that lists one of the four methods in `DeclareOpInterfaceMethods` now has **both** arities declared on it, so it must also define the hint-bearing overload (typically a one-line forward to the hint-less one). - An `ExternalModel` that defines the hint-less overload **hides** the inherited defaulted hint-bearing overload (C++ name hiding); add `using Base::<method>;` to keep it visible. Implementers that do not care about the hint need no functional change beyond the above visibility boilerplate. ## SCF tiling driver API The producer- and consumer-fusion driver entry points (`tileAndFuseProducerOfSlice`, `tileAndFuseConsumerOfSlices`, `tileAndFuseConsumer`) and `SCFTilingOptions` accept an optional **control function** rather than a single fixed array. A single tile-and-fuse invocation can tile/fuse several ops with **different iteration domains**, so no single array can be correct for all of them. The driver therefore invokes the control function once per op being tiled/fused and expects the returned array to be in *that op's own* domain, it does not remap it. ### Why a control function and not one array A single tile-and-fuse can pull in more than one pack/unpack, and each consults the hint in *its own* iteration domain: ```mlir %unpacked = linalg.unpack %p0 inner_dims_pos = [0] inner_tiles = [%c8_vscale] ... // domain A %mul = linalg.generic ins(%unpacked, ...) ... // tiled root %packed = linalg.pack %mul inner_dims_pos = [1] inner_tiles = [%c4_vscale] ... // domain B // One tile-and-fuse: tile %mul, fuse %unpacked (producer) and %packed (consumer). ``` Both `%unpacked` and `%packed` consume the hint, but their domains differ, a different dimension is tiled and the inner-tile arrangement is different, so the array indexed in `%unpacked`'s domain is simply wrong for `%packed`. No single fixed array can be correct for both. The driver therefore invokes the control function once per op being tiled/fused and the caller returns each op's array in its own domain. ### Why the function gets slices (not tile sizes) for fusion For pure tiling the driver knows the explicit tile sizes and passes them. In fusion there are no caller-specified tile sizes for the fused op, its tile is induced by the `tensor.extract_slice` (producer) or `insert_slice` (consumer) being fused. The control function is therefore handed those slices, from which the caller can derive the alignment in the fused op's domain. ## Transform ops `transform.structured.tile_using_for`, `transform.structured.fuse`, and `transform.structured.fuse_into_containing_op` gain an optional `inner_tile_alignments` array attribute (empty by default, preserving prior behavior). It is wired to the driver as a **constant** control function (the same array for every op). This is sufficient because each of these ops drives tiling/fusion at single-op granularity; the per-op control function remains available at the library/driver level for callers that fuse multiple ops with differing domains. The test-only `transform.test.fuse_consumer` is wired the same way. ## Testing - **`linalg.unpack`**: standalone tiling, producer fusion (`structured.fuse`, `fuse_into_containing_op`), and consumer fusion (`test.fuse_consumer`), with both `Multiple` and `Equal` hints. - **`linalg.pack`**: standalone tiling, producer fusion, and consumer fusion (`test.fuse_consumer`), with both `Multiple` and `Equal` hints. - Negative cases: missing hint falls back to the conservative tiling; a hint that contradicts statically-known sizes is ignored; transposing `outer_dims_perm`/consumers confirm source-order (un-permuted) indexing. - Transform-op attribute verification (`inner_tile_alignments` validity). Assisted-by: Claude Code --------- Signed-off-by: Ege Beysel <beyselege@gmail.com>
96 lines
3.4 KiB
C++
96 lines
3.4 KiB
C++
//===- TilingInterface.cpp - Tiling interface -------------------*- C++ -*-===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file contains the definitions of the interface in `TilingInterface.td`.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Interfaces/TilingInterface.h"
|
|
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/ADT/SmallVectorExtras.h"
|
|
#include "llvm/ADT/StringSwitch.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);
|
|
});
|
|
}
|
|
|
|
StringRef mlir::stringifyInnerTileAlignment(InnerTileAlignment alignment) {
|
|
switch (alignment) {
|
|
case InnerTileAlignment::Unknown:
|
|
return "Unknown";
|
|
case InnerTileAlignment::Multiple:
|
|
return "Multiple";
|
|
case InnerTileAlignment::Equal:
|
|
return "Equal";
|
|
}
|
|
llvm_unreachable("unknown InnerTileAlignment");
|
|
}
|
|
|
|
std::optional<InnerTileAlignment>
|
|
mlir::symbolizeInnerTileAlignment(StringRef keyword) {
|
|
return llvm::StringSwitch<std::optional<InnerTileAlignment>>(keyword)
|
|
.Case("Unknown", InnerTileAlignment::Unknown)
|
|
.Case("Multiple", InnerTileAlignment::Multiple)
|
|
.Case("Equal", InnerTileAlignment::Equal)
|
|
.Default(std::nullopt);
|
|
}
|
|
|
|
ParseResult mlir::parseInnerTileAlignmentArray(OpAsmParser &parser,
|
|
DenseI64ArrayAttr &alignments) {
|
|
SmallVector<int64_t> values;
|
|
auto parseEntry = [&]() -> ParseResult {
|
|
StringRef keyword;
|
|
llvm::SMLoc loc = parser.getCurrentLocation();
|
|
if (parser.parseKeyword(&keyword))
|
|
return failure();
|
|
std::optional<InnerTileAlignment> alignment =
|
|
symbolizeInnerTileAlignment(keyword);
|
|
if (!alignment)
|
|
return parser.emitError(loc)
|
|
<< "expected one of 'Unknown', 'Multiple' or 'Equal', but got '"
|
|
<< keyword << "'";
|
|
values.push_back(static_cast<int64_t>(*alignment));
|
|
return success();
|
|
};
|
|
if (parser.parseCommaSeparatedList(AsmParser::Delimiter::Square, parseEntry))
|
|
return failure();
|
|
alignments = DenseI64ArrayAttr::get(parser.getContext(), values);
|
|
return success();
|
|
}
|
|
|
|
void mlir::printInnerTileAlignmentArray(OpAsmPrinter &printer, Operation *,
|
|
DenseI64ArrayAttr alignments) {
|
|
printer << "[";
|
|
llvm::interleaveComma(alignments.asArrayRef(), printer, [&](int64_t value) {
|
|
printer << stringifyInnerTileAlignment(
|
|
static_cast<InnerTileAlignment>(value));
|
|
});
|
|
printer << "]";
|
|
}
|
|
|
|
#include "mlir/Interfaces/TilingInterface.cpp.inc"
|