[LLVMCPU] add inner-tile-alignment hint attribute (#24806)
Introduce the representation for per-tiling-level inner-tile alignment
hints
as introduced in llvm/llvm-project#204007:
#iree_cpu.inner_tile_alignments<level = [Unknown|Multiple|Equal, ...]>
These will only be set on pack and unpack operations with scalable inner
tile sizes, and will not be for other operations that do not need them.
Assisted-by: Claude Code
---
<sub>Stack created with <a
href="https://github.com/github/gh-stack">GitHub Stacks CLI</a> • <a
href="https://gh.io/stacks-feedback">Give Feedback 💬</a></sub>
Signed-off-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel
index 8ef2115..ac9b672 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel
@@ -77,6 +77,7 @@
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TensorDialect",
+ "@llvm-project//mlir:TilingInterface",
"@llvm-project//mlir:VectorDialect",
],
)
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt
index 6ff9835..ea6059e 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt
@@ -40,6 +40,7 @@
MLIRParser
MLIRSupport
MLIRTensorDialect
+ MLIRTilingInterface
MLIRVectorDialect
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::Dialect::Codegen::Utils
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
index 8997f88..ed1d263 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
@@ -28,6 +28,7 @@
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/ValueRange.h"
+#include "mlir/Interfaces/TilingInterface.h"
#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUEnums.cpp.inc"
#define GET_ATTRDEF_CLASSES
@@ -323,6 +324,110 @@
}
//===----------------------------------------------------------------------===//
+// InnerTileAlignmentsAttr
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+InnerTileAlignmentsAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+ DictionaryAttr alignments) {
+ if (!alignments || alignments.empty()) {
+ return emitError() << "expected at least one tiling level";
+ }
+ for (NamedAttribute entry : alignments) {
+ auto arr = dyn_cast<DenseI64ArrayAttr>(entry.getValue());
+ if (!arr) {
+ return emitError()
+ << "expected a per-dimension InnerTileAlignment array for '"
+ << entry.getName().getValue() << "'";
+ }
+ for (int64_t value : arr.asArrayRef()) {
+ if (!mlir::isValidInnerTileAlignment(value)) {
+ return emitError() << "invalid InnerTileAlignment value: " << value;
+ }
+ }
+ }
+ return success();
+}
+
+Attribute InnerTileAlignmentsAttr::parse(AsmParser &parser, Type) {
+ MLIRContext *ctx = parser.getContext();
+ if (parser.parseLess()) {
+ return {};
+ }
+ SmallVector<NamedAttribute> items;
+ bool first = true;
+ while (failed(parser.parseOptionalGreater())) {
+ if (!first && parser.parseComma()) {
+ return {};
+ }
+ first = false;
+ // <level_name> = [Unknown, Equal, ...]
+ std::string keyStr;
+ if (parser.parseKeywordOrString(&keyStr) || parser.parseEqual()) {
+ return {};
+ }
+ SmallVector<int64_t> alignments;
+ if (parser.parseCommaSeparatedList(
+ AsmParser::Delimiter::Square, [&]() -> ParseResult {
+ StringRef keyword;
+ if (parser.parseKeyword(&keyword)) {
+ return failure();
+ }
+ std::optional<mlir::InnerTileAlignment> kind =
+ mlir::symbolizeInnerTileAlignment(keyword);
+ if (!kind) {
+ return parser.emitError(parser.getCurrentLocation(),
+ "expected an InnerTileAlignment "
+ "(Unknown|Multiple|Equal), got: ")
+ << keyword;
+ }
+ alignments.push_back(static_cast<int64_t>(*kind));
+ return success();
+ })) {
+ return {};
+ }
+ items.emplace_back(StringAttr::get(ctx, keyStr),
+ DenseI64ArrayAttr::get(ctx, alignments));
+ }
+ return parser.getChecked<InnerTileAlignmentsAttr>(
+ ctx, DictionaryAttr::get(ctx, items));
+}
+
+void InnerTileAlignmentsAttr::print(AsmPrinter &printer) const {
+ printer << "<";
+ llvm::interleaveComma(getAlignments(), printer, [&](NamedAttribute entry) {
+ // `.str()` avoids wrapping the tiling-level key with `"`.
+ printer << entry.getName().str() << " = [";
+ ArrayRef<int64_t> alignments =
+ cast<DenseI64ArrayAttr>(entry.getValue()).asArrayRef();
+ llvm::interleaveComma(alignments, printer, [&](int64_t value) {
+ printer << mlir::stringifyInnerTileAlignment(
+ static_cast<mlir::InnerTileAlignment>(value));
+ });
+ printer << "]";
+ });
+ printer << ">";
+}
+
+void InnerTileAlignmentsAttr::setOnOp(
+ Operation *op,
+ ArrayRef<std::pair<TilingLevel, SmallVector<int64_t>>> perLevel) {
+ MLIRContext *ctx = op->getContext();
+ SmallVector<NamedAttribute> entries;
+ entries.reserve(perLevel.size());
+ for (auto &[level, alignments] : perLevel) {
+ entries.emplace_back(StringAttr::get(ctx, getTilingLevelName(level)),
+ DenseI64ArrayAttr::get(ctx, alignments));
+ }
+ op->setAttr(getMnemonic(), InnerTileAlignmentsAttr::get(
+ ctx, DictionaryAttr::get(ctx, entries)));
+}
+
+InnerTileAlignmentsAttr InnerTileAlignmentsAttr::getFromOp(Operation *op) {
+ return op->getAttrOfType<InnerTileAlignmentsAttr>(getMnemonic());
+}
+
+//===----------------------------------------------------------------------===//
// CPU MMA intrinsic layout (MxNxK shape and element types)
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td
index 9629fd8..5fcfe91 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td
@@ -249,6 +249,35 @@
let genVerifyDecl = 1;
}
+def IREECPU_InnerTileAlignmentsAttr :
+ AttrDef<IREECPU_Dialect, "InnerTileAlignments"> {
+ let mnemonic = "inner_tile_alignments";
+ let summary = [{Per-tiling-level inner-tile alignment hints for pack/unpack ops.}];
+ let description = [{
+ A discardable hint attached to a scalable `linalg.pack`/`linalg.unpack`
+ during tile-size selection, recording, per tiling level, the
+ `mlir::InnerTileAlignment` (`Unknown`, `Multiple` or `Equal`) of each
+ iteration-domain dimension.
+ }];
+ let parameters = (ins
+ AttrParameter<"DictionaryAttr",
+ "Maps a tiling-level name to its per-dim InnerTileAlignment array.">:$alignments
+ );
+ let extraClassDeclaration = [{
+ /// Builds a hint from per-tiling-level `InnerTileAlignment` values and attaches
+ /// it to `op` under the attribute's mnemonic, overwriting any existing one.
+ static void setOnOp(
+ ::mlir::Operation *op,
+ ::llvm::ArrayRef<
+ ::std::pair<TilingLevel, ::llvm::SmallVector<int64_t>>> perLevel);
+
+ /// Returns the hint attached to `op`, or a null attribute if it has none.
+ static InnerTileAlignmentsAttr getFromOp(::mlir::Operation *op);
+ }];
+ let hasCustomAssemblyFormat = 1;
+ let genVerifyDecl = 1;
+}
+
//===----------------------------------------------------------------------===//
// Encoding Resolvers.
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/test/invalid.mlir b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/test/invalid.mlir
index 05e1980..0a6b31f 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/test/invalid.mlir
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/test/invalid.mlir
@@ -122,3 +122,23 @@
} : tensor<1x1x1x1xf32>, tensor<1x1x4x1xf32> into tensor<1x1x1x?xf32>
return %0 : tensor<1x1x1x?xf32>
}
+
+// -----
+
+// The inner-tile alignment array only accepts InnerTileAlignment keywords
+// (Unknown, Multiple, Equal).
+// expected-error@+1 {{expected an InnerTileAlignment}}
+#invalid_inner_tile_alignment_name = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Bogus]>
+
+// -----
+
+// Each alignment must be encapsulated in an array, e.g. `[Multiple]`, not
+// given as a bare keyword.
+// expected-error@+1 {{}}
+#invalid_inner_tile_alignment_value = #iree_cpu.inner_tile_alignments<vector_common_parallel = Multiple>
+
+// -----
+
+// An empty hint (no tiling levels) is meaningless and is rejected.
+// expected-error@+1 {{expected at least one tiling level}}
+#invalid_inner_tile_alignment_empty = #iree_cpu.inner_tile_alignments<>
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/test/roundtrip.mlir b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/test/roundtrip.mlir
index 29d2d44..5a2845a 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/test/roundtrip.mlir
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/test/roundtrip.mlir
@@ -82,3 +82,35 @@
}
// CHECK-LABEL: @test_ukernel_provider()
// CHECK-SAME: iree_codegen.ukernel_provider = #iree_cpu.ukernel_provider
+
+// -----
+
+// Round-trip the inner-tile alignment hint attribute. It records, per tiling
+// level, the InnerTileAlignment (Unknown|Multiple|Equal) of each
+// iteration-domain dimension.
+func.func @test_inner_tile_alignments() attributes {
+ inner_tile_alignments =
+ #iree_cpu.inner_tile_alignments<vector_common_parallel = [Unknown, Equal]>} {
+ return
+}
+// CHECK-LABEL: @test_inner_tile_alignments()
+// CHECK-SAME: inner_tile_alignments = #iree_cpu.inner_tile_alignments<
+// CHECK-SAME: vector_common_parallel = [Unknown, Equal]>
+
+// -----
+
+// Multiple levels and all three alignment kinds, each level carrying one entry
+// per iteration-domain dimension (here 2). Levels are printed sorted by name,
+// regardless of input order.
+func.func @test_inner_tile_alignments_multi_level() attributes {
+ inner_tile_alignments = #iree_cpu.inner_tile_alignments<
+ vector_inner_parallel = [Equal, Unknown],
+ vector_common_parallel = [Multiple, Unknown],
+ distribution = [Unknown, Equal]>} {
+ return
+}
+// CHECK-LABEL: @test_inner_tile_alignments_multi_level()
+// CHECK-SAME: inner_tile_alignments = #iree_cpu.inner_tile_alignments<
+// CHECK-SAME: distribution = [Unknown, Equal],
+// CHECK-SAME: vector_common_parallel = [Multiple, Unknown],
+// CHECK-SAME: vector_inner_parallel = [Equal, Unknown]>
