[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; }