[LLVMCPU][NFC] Separate the intrinsic MNK shape from its register layout (#24861)
`getIntrinsicMNKShape` states how many M/N/K elements one invocation of
an intrinsic computes, independent of how those elements are laid out in
registers. Selection scores candidates on this shape and the
sibling-pairing validator checks M and N against it, so every intrinsic
declares one, including scalable ones, which declare their base sizes. A
scalable and a non-scalable intrinsic should never compete for the same
matmul, so comparisons stay like-for-like.
The hand-rolled swizzles for intrinsics whose tiles are not row-major
move into `getNonRowMajorIntrinsicSwizzle`.
Assisted-by: Claude Code
Signed-off-by: Ege Beysel <beyselege@gmail.com>
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 ed1d263..61ae8a6 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
@@ -432,19 +432,12 @@
//===----------------------------------------------------------------------===//
// Helper function for getIntrinsicSwizzle.
-// Allows succinctly specifying the tiles for the common case of row-major tiles
-// i.e. when the TileSwizzle merely encodes the 2D tile shape and no further
-// expand/transpose.
-// This case is very common for CPU MMA intrinsics due to:
-// 1. CPU matmul instructions having simple tile layouts.
-// 2. The inner_tiled op RHS being transposed, meaning that the row-major RHS
-// tile really is a column-major RHS matmul tile.
-//
-// If you're trying to add support for a new intrinsic that doesn't have a
-// row-major tile layout, don't look at this function, go directly to
-// getIntrinsicSwizzle.
+// Returns how many M/N/K elements one invocation of `intrinsic` computes,
+// independent of how those elements are laid out in registers. Every intrinsic
+// declares its shape here, including scalable ones, which declare their base
+// sizes.
std::optional<std::tuple<int64_t, int64_t, int64_t>>
-getRowMajorTilesMNKShape(MMAIntrinsic intrinsic) {
+getIntrinsicMNKShape(MMAIntrinsic intrinsic) {
using Tuple = std::tuple<int64_t, int64_t, int64_t>;
switch (intrinsic) {
case MMAIntrinsic::None:
@@ -487,6 +480,11 @@
// ACC tile has a non-row-major layout, hand-rolled in `getIntrinsicSwizzle`.
case MMAIntrinsic::MMA_X86_AVX512VNNI_16x16x2_I32_I8_CASTI16:
return Tuple{16, 16, 2};
+ // Base sizes, i.e. the shape at the minimum vector length.
+ case MMAIntrinsic::MMA_ARM_SVE_FMLA_1x4VLx1_F32_F32:
+ return Tuple{1, 4, 1};
+ case MMAIntrinsic::MMA_ARM_SVE_FMLA_4VLx1x1_F32_F32:
+ return Tuple{4, 1, 1};
default:
if (isGenericScalar(intrinsic)) {
return Tuple{1, 1, 1};
@@ -559,16 +557,18 @@
return swizzle;
}
-Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma,
- int operandIdx) {
+/// Helper for `getIntrinsicSwizzle`:
+/// Returns the hand-rolled TileSwizzle for the intrinsics whose tiles are not
+/// row-major, or nullopt for every other intrinsic.
+static std::optional<Codegen::TileSwizzle>
+getNonRowMajorIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma, int operandIdx) {
using TileSwizzle = Codegen::TileSwizzle;
using Dim = TileSwizzle::Dim;
- // SVE has scalable dims that the row-major path below can't express, so we
- // hand-roll the swizzle. The natural orientation is 1×4VL×1: 4VL lives on
- // RHS dim 0 (N) and ACC dim 0 (the convention is `(N, M)` here). The
- // transposed orientation 4VL×1×1 mirrors that: 4VL on LHS dim 0 (M) and
- // ACC dim 0 (now `(M, N)`).
+ // SVE has scalable dims that the row-major path can't express. The natural
+ // orientation is 1×4VL×1: 4VL lives on RHS dim 0 (N) and ACC dim 0 (the
+ // convention is `(N, M)` here). The transposed orientation 4VL×1×1 mirrors
+ // that: 4VL on LHS dim 0 (M) and ACC dim 0 (now `(M, N)`).
if (mma == MMAIntrinsic::MMA_ARM_SVE_FMLA_1x4VLx1_F32_F32 ||
mma == MMAIntrinsic::MMA_ARM_SVE_FMLA_4VLx1x1_F32_F32) {
bool transposed = (mma == MMAIntrinsic::MMA_ARM_SVE_FMLA_4VLx1x1_F32_F32);
@@ -601,11 +601,30 @@
return swizzle;
}
- auto maybeMnkTuple = getRowMajorTilesMNKShape(mma);
+ return std::nullopt;
+}
+
+Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma,
+ int operandIdx) {
+ using TileSwizzle = Codegen::TileSwizzle;
+ if (std::optional<TileSwizzle> swizzle =
+ getNonRowMajorIntrinsicSwizzle(mma, operandIdx)) {
+ return *swizzle;
+ }
+
+ // What follows specifies the tiles for the common case of row-major tiles,
+ // i.e. when the TileSwizzle merely encodes the 2D tile shape and no further
+ // expand/transpose. This case is very common for CPU MMA intrinsics due to:
+ // 1. CPU matmul instructions having simple tile layouts.
+ // 2. The inner_tiled op RHS being transposed, meaning that the row-major RHS
+ // tile really is a column-major RHS matmul tile.
+ //
+ // If you're trying to add support for a new intrinsic that doesn't have a
+ // row-major tile layout, its swizzle belongs in
+ // `getNonRowMajorIntrinsicSwizzle`.
+ auto maybeMnkTuple = getIntrinsicMNKShape(mma);
if (!maybeMnkTuple) {
- // Whenever one adds support for a new intrinsic that doesn't have a
- // row-major tile layout, new logic goes here.
- assert(false && "Non-row-major-tile intrinsics not yet implemented.");
+ assert(false && "intrinsic does not declare an MNK shape.");
return TileSwizzle();
}
auto [mSize, nSize, kSize] = *maybeMnkTuple;
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h
index c0a1425..e0a6da3 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h
@@ -93,10 +93,10 @@
std::tuple<Type, Type, Type> getABCElementTypes(MLIRContext *ctx,
MMAIntrinsic intrinsic);
-// Returns the (M, N, K) tile shape for a row-major-tile intrinsic, or
-// nullopt otherwise (e.g. SVE).
+// Returns how many M/N/K elements one invocation of `intrinsic` computes.
+// Scalable intrinsics declare their base sizes.
std::optional<std::tuple<int64_t, int64_t, int64_t>>
-getRowMajorTilesMNKShape(MMAIntrinsic intrinsic);
+getIntrinsicMNKShape(MMAIntrinsic intrinsic);
// Idempotently attaches the bitcode for ukernel `name` as
// `hal.executable.objects` on `op`, looking it up first in any
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
index 0a5c8abe..e922b33 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -454,8 +454,8 @@
// `ui8` LHS visibly swaps with its sibling's `ui8` RHS. Read directly
// from `getABCElementTypes`; `getUndistributedTileTypes` would strip it.
auto getShapeAndTypes = [&](MMAIntrinsic intr) {
- auto mnk = IREE::CPU::getRowMajorTilesMNKShape(intr);
- assert(mnk && "validator only handles row-major-tile intrinsics");
+ auto mnk = IREE::CPU::getIntrinsicMNKShape(intr);
+ assert(mnk && "intrinsic does not declare an MNK shape");
auto [m, n, k] = *mnk;
auto [lhs, rhs, acc] = IREE::CPU::getABCElementTypes(ctx, intr);
return std::make_tuple(m, n, k, lhs, rhs, acc);
@@ -599,7 +599,7 @@
accTy != elementTypes[2]) {
return std::nullopt;
}
- auto mnk = IREE::CPU::getRowMajorTilesMNKShape(intr);
+ auto mnk = IREE::CPU::getIntrinsicMNKShape(intr);
if (!mnk) {
return std::nullopt;
}
