diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingPatterns.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingPatterns.cpp index 259af263664a..0fae21185edb 100644 --- a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingPatterns.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingPatterns.cpp @@ -662,6 +662,10 @@ class MaterializeLinalgOp } Operation *newOp = layoutAttr.lowerOp(rewriter, op, convertedResTypes, operands); + if (!newOp) { + return op->emitOpError( + "unsupported operation for encoding materialization"); + } rewriter.replaceOp(op, newOp->getResults()); return success(); } diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir index 576530a2ace1..159f1b6612e8 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir +++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir @@ -148,6 +148,98 @@ func.func @pack_gemm_fill_dynamic(%arg0 : tensor, %arg1 : tensor (d0, d2)> +#map_it1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map_it2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#encoding_lhs_it = #iree_encoding.encoding +#encoding_rhs_it = #iree_encoding.encoding +#encoding_result_it = #iree_encoding.encoding +func.func @pack_gemm_fill_dynamic_inner_tiled_avx512(%arg0 : tensor, %arg1 : tensor, %m: index, %n: index, %k: index) -> tensor attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f", enable_inner_tiled = true, iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %cst = arith.constant 0.0 : f32 + %d0 = tensor.dim %arg0, %c0 : tensor + %d1 = tensor.dim %arg1, %c1 : tensor + %0 = iree_encoding.set_encoding %arg0 encoding_dims{%m, %n, %k} : tensor -> tensor + %1 = iree_encoding.set_encoding %arg1 encoding_dims{%m, %n, %k} : tensor -> tensor + %2 = tensor.empty(%d0, %d1) : tensor + %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor) + -> tensor + %4 = linalg.matmul ins(%0, %1 : tensor, tensor) + outs(%3 : tensor) -> tensor + %5 = iree_encoding.unset_encoding %4 encoding_dims{%m, %n, %k} : tensor -> tensor{%d0, %d1} + return %5 : tensor +} +// CHECK-DAG: #[[$MAP_N:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)> +// CHECK-LABEL: func @pack_gemm_fill_dynamic_inner_tiled_avx512( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor +// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index +// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index +// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]] +// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG1]], %[[C1]] +// CHECK-DAG: %[[OUT_D1:.+]] = affine.apply #[[$MAP_N]]()[%[[D1]]] +// CHECK-DAG: %[[PACK_LHS:.+]] = linalg.pack {{.*}}%[[ARG0]] +// CHECK: %[[PACK_RHS:.+]] = linalg.pack +// CHECK-SAME: %[[ARG1]] +// CHECK-DAG: %[[EMPTY:.+]] = tensor.empty(%[[D0]], %[[OUT_D1]]) : tensor +// CHECK: %[[FILL:.+]] = linalg.fill +// CHECK-SAME: outs(%[[EMPTY]] : +// CHECK: %[[INNER:.+]] = iree_codegen.inner_tiled ins(%[[PACK_LHS]], %[[PACK_RHS]]) outs(%[[FILL]]) +// CHECK-SAME: kind = #iree_cpu.data_tiled_mma_layout, semantics = #iree_cpu.mma_semantics<> +// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[INNER]] +// CHECK: return %[[UNPACK]] + +// ----- + +#map_it_se = affine_map<(d0, d1, d2) -> (d0, d2)> +#map_it_se1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map_it_se2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#encoding_it_lhs = #iree_encoding.encoding +#encoding_it_rhs = #iree_encoding.encoding +#encoding_it_res = #iree_encoding.encoding +func.func @set_encoding_matmul_LHS_inner_tiled_avx512(%arg0: tensor<127x255xf32>, %k: index) -> tensor<127x255xf32, #encoding_it_lhs> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f", enable_inner_tiled = true, iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %0 = iree_encoding.set_encoding %arg0 encoding_dims{%k} : tensor<127x255xf32> -> tensor<127x255xf32, #encoding_it_lhs> + return %0 : tensor<127x255xf32, #encoding_it_lhs> +} +func.func @set_encoding_matmul_RHS_inner_tiled_avx512(%arg0: tensor<127x255xf32>, %k: index) -> tensor<127x255xf32, #encoding_it_rhs> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f", enable_inner_tiled = true, iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %0 = iree_encoding.set_encoding %arg0 encoding_dims{%k} : tensor<127x255xf32> -> tensor<127x255xf32, #encoding_it_rhs> + return %0 : tensor<127x255xf32, #encoding_it_rhs> +} +func.func @unset_encoding_matmul_RESULT_inner_tiled_avx512(%arg0: tensor<127x255xf32, #encoding_it_res>, %k: index) -> tensor<127x255xf32> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f", enable_inner_tiled = true, iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %0 = iree_encoding.unset_encoding %arg0 encoding_dims{%k} : tensor<127x255xf32, #encoding_it_res> -> tensor<127x255xf32> + return %0 : tensor<127x255xf32> +} +// CHECK-LABEL: func @set_encoding_matmul_LHS_inner_tiled_avx512( +// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<127x255xf32> +// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<127x255x1x1xf32> +// CHECK: %[[PACK:.+]] = linalg.pack %[[INPUT]] outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [1, 1] into %[[EMPTY]] : tensor<127x255xf32> -> tensor<127x255x1x1xf32> +// CHECK: return %[[PACK]] : tensor<127x255x1x1xf32> +// CHECK-LABEL: func @set_encoding_matmul_RHS_inner_tiled_avx512( +// CHECK-SAME: %[[INPUT_R:[a-zA-Z0-9]+]]: tensor<127x255xf32> +// CHECK-DAG: %[[CST_R:.+]] = arith.constant 0.0 +// CHECK: %[[EMPTY_R:.+]] = tensor.empty() : tensor<16x127x16x1xf32> +// CHECK: %[[PACK_R:.+]] = linalg.pack %[[INPUT_R]] padding_value(%[[CST_R]] : f32) outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 1] into %[[EMPTY_R]] : tensor<127x255xf32> -> tensor<16x127x16x1xf32> +// CHECK: return %[[PACK_R]] : tensor<16x127x16x1xf32> +// CHECK-LABEL: func @unset_encoding_matmul_RESULT_inner_tiled_avx512( +// CHECK-SAME: %[[PACKED:[a-zA-Z0-9]+]]: tensor<127x16x1x16xf32> +// CHECK: %[[EMPTY_U:.+]] = tensor.empty() : tensor<127x255xf32> +// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[PACKED]] outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [1, 16] into %[[EMPTY_U]] : tensor<127x16x1x16xf32> -> tensor<127x255xf32> +// CHECK: return %[[UNPACK]] + +// ----- + // It tests with bindings and checks that the reshape ops are folded into bindings. #executable_target_xyz = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> 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 55fdbe116101..8716af67f6e9 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp @@ -332,6 +332,8 @@ getRowMajorTilesMNKShape(MMAIntrinsic intrinsic) { switch (intrinsic) { case MMAIntrinsic::None: return Tuple{0, 0, 0}; + case MMAIntrinsic::MMA_X86_AVX2_FMA_1x8x1_F32_F32: + return Tuple{1, 8, 1}; case MMAIntrinsic::MMA_X86_AVX512_1x8x1_F64_F64: return Tuple{1, 8, 1}; case MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F32: @@ -370,7 +372,7 @@ static Codegen::TileSwizzle fixupSwizzle(Codegen::TileSwizzle swizzle) { } Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma, - int operandIdx) { + bool transposed, int operandIdx) { using TileSwizzle = Codegen::TileSwizzle; using Dim = TileSwizzle::Dim; @@ -394,6 +396,13 @@ Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma, return TileSwizzle(); } auto [mSize, nSize, kSize] = *maybeMnkTuple; + // In the transposed orientation, the intrinsic's hardware (M, N) roles are + // logically swapped: what the matmul code sees as M is driven by the + // intrinsic's N-dim, and vice versa. Swap the sizes up front so that the + // per-operand expansion code below can stay oblivious to orientation. + if (transposed) { + std::swap(mSize, nSize); + } TileSwizzle swizzle; swizzle.expandShape().resize(2); auto expandIfNonUnit = [](TileSwizzle &swizzle, int dim, int size) { @@ -402,6 +411,13 @@ Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma, } }; + // For every operand, expandShape[0] is the outer physical dim and + // expandShape[1] is the inner physical dim, with identity permutation. For + // the ACC in particular, `transposed_intrinsic` flips the logical (M, N) to + // physical (N, M), which we encode by swapping which logical dim fills each + // expandShape group rather than by a non-identity permutation. That way all + // three operand swizzles can be read on equal footing, just like LHS (M, K) + // and RHS (N, K). if (operandIdx == 0) { constexpr int M = 0, K = 1; expandIfNonUnit(swizzle, K, kSize); @@ -411,9 +427,10 @@ Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma, expandIfNonUnit(swizzle, K, kSize); expandIfNonUnit(swizzle, N, nSize); } else { - constexpr int M = 0, N = 1; - expandIfNonUnit(swizzle, N, nSize); - expandIfNonUnit(swizzle, M, mSize); + int64_t accOuter = transposed ? nSize : mSize; + int64_t accInner = transposed ? mSize : nSize; + expandIfNonUnit(swizzle, 1, accInner); + expandIfNonUnit(swizzle, 0, accOuter); } return fixupSwizzle(std::move(swizzle)); } @@ -421,7 +438,8 @@ Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma, Codegen::TileSwizzle getSwizzle(IREE::CPU::DataTiledMMAAttr mma, int operandIdx) { using TileSwizzle = Codegen::TileSwizzle; - TileSwizzle swizzle = getIntrinsicSwizzle(mma.getIntrinsic(), operandIdx); + TileSwizzle swizzle = getIntrinsicSwizzle( + mma.getIntrinsic(), mma.getTransposedIntrinsic(), operandIdx); TileSwizzle::Dim intrinsicsM = TileSwizzle::Dim::crossIntrinsic(mma.getIntrinsicsM()); TileSwizzle::Dim intrinsicsN = @@ -430,8 +448,9 @@ Codegen::TileSwizzle getSwizzle(IREE::CPU::DataTiledMMAAttr mma, TileSwizzle::Dim::crossIntrinsic(mma.getIntrinsicsK()); // Each swizzle is built as (outer physical dim, inner physical dim) in // expandShape[0], expandShape[1]. LHS is (M, K), RHS is (N, K), ACC is - // (M, N). The expansion below injects the intrinsics_* cross-intrinsic - // factors into whichever group represents each logical dim. + // (M, N) normally and (N, M) when `transposed_intrinsic` is set. The + // expansion below injects the intrinsics_* cross-intrinsic factors into + // whichever group represents each logical dim. if (operandIdx == 0) { constexpr int M = 0, K = 1; if (intrinsicsK.size() > 1) { @@ -449,12 +468,14 @@ Codegen::TileSwizzle getSwizzle(IREE::CPU::DataTiledMMAAttr mma, Codegen::expand(swizzle, N, intrinsicsN); } } else { - constexpr int M = 0, N = 1; - if (intrinsicsN.size() > 1) { - Codegen::expand(swizzle, N, intrinsicsN); + bool transposed = mma.getTransposedIntrinsic(); + TileSwizzle::Dim accOuterIntr = transposed ? intrinsicsN : intrinsicsM; + TileSwizzle::Dim accInnerIntr = transposed ? intrinsicsM : intrinsicsN; + if (accInnerIntr.size() > 1) { + Codegen::expand(swizzle, /*srcIdx=*/1, accInnerIntr); } - if (intrinsicsM.size() > 1) { - Codegen::expand(swizzle, M, intrinsicsM); + if (accOuterIntr.size() > 1) { + Codegen::expand(swizzle, /*srcIdx=*/0, accOuterIntr); } } return swizzle; @@ -472,6 +493,8 @@ static std::tuple getABCElementTypes(MLIRContext *context, switch (intrinsic) { case MMAIntrinsic::None: return {Type(), Type(), Type()}; + case MMAIntrinsic::MMA_X86_AVX2_FMA_1x8x1_F32_F32: + return {f32, f32, f32}; case MMAIntrinsic::MMA_X86_AVX512_1x8x1_F64_F64: return {f64, f64, f64}; case MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F32: @@ -536,7 +559,10 @@ void DataTiledMMAAttr::getUndistributedTileTypes( // inner physical dim) in expandShape[0], expandShape[1]. This mirrors GPU's // DataTiledMMA, where the tile types encode the layout directly and no // separate `permutations` attribute is needed on `inner_tiled`. LHS is - // (M, K), RHS is (N, K), ACC is (M, N). + // (M, K), RHS is (N, K), and ACC is (M, N) for non-transposed intrinsics + // and (N, M) for transposed ones; that transposition is baked into the ACC + // swizzle itself (see getIntrinsicSwizzle), so we can query all three + // operands uniformly here. auto tileType = [&](Codegen::TileSwizzle swizzle, Type elemType) { auto [outer, outerScalable] = getVectorAxisSizeAndScalability(swizzle.expandShape()[0]); 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 588b1fb30207..80548de5b85b 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td +++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td @@ -83,6 +83,12 @@ def IREECPU_DataTiledMMAAttr thread-distribution, no concept of subgroups and no interleaving of intrinsics' layout. + The `transposed_intrinsic` boolean toggles between the intrinsic's natural + orientation and an M↔N-swapped orientation (same hardware instruction with + LHS/RHS roles exchanged, accumulator laid out column-major). Orthogonal to + the intrinsic itself: any intrinsic may be used transposed, including + square ones where the effect is only the accumulator layout change. + Some GPU-specific methods in IREECodegen_InnerTileDescAttrInterface are left here but are unused. }]; @@ -99,7 +105,12 @@ def IREECPU_DataTiledMMAAttr "Intrinsic count along the N dimension.">:$intrinsics_n, DefaultValuedParameter< "int64_t", "1", - "Intrinsic count along the K dimension.">:$intrinsics_k); + "Intrinsic count along the K dimension.">:$intrinsics_k, + DefaultValuedParameter< + "bool", "false", + "If true, the intrinsic is used in an M↔N-swapped orientation " + "(LHS/RHS roles exchanged, accumulator " + "column-major).">:$transposed_intrinsic); } //===----------------------------------------------------------------------===// diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUEnums.td b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUEnums.td index 279c0a0a583f..b88c59ada90c 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUEnums.td +++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUEnums.td @@ -96,6 +96,18 @@ class IREECPU_I32EnumAttr; + +// X86 AVX/AVX2 + FMA3 (256-bit YMM registers). FMA3 is a separate ISA extension +// from AVX2 itself; these intrinsics lower to `vfmadd...`, so the name includes +// `AVX2_FMA_` and `getMmaIntrinsicRequiredFeatures` requires both `+avx2` and +// `+fma`. +def MMA_X86_AVX2_FMA_1x8x1_F32_F32 + : I32EnumAttrCase<"MMA_X86_AVX2_FMA_1x8x1_F32_F32", 0x1210>; + +// X86 AVX-512 (ZMM registers). Only the "natural" (canonical M<=N) orientation +// of each intrinsic is listed here; the M/N-swapped ("transposed") variant is +// expressed by the `transposed_intrinsic` boolean on DataTiledMMAAttr rather +// than as a separate enum value. def MMA_X86_AVX512_1x8x1_F64_F64 : I32EnumAttrCase<"MMA_X86_AVX512_1x8x1_F64_F64", 0x1300>; def MMA_X86_AVX512_1x16x1_F32_F32 @@ -125,6 +137,9 @@ def IREECPU_MMAIntrinsic "MMAIntrinsic", "Descriptor for different MMA intrinsics", [IREECPU_MMA_None, + // X86 AVX/AVX2 + FMA + MMA_X86_AVX2_FMA_1x8x1_F32_F32, + // X86 AVX-512 MMA_X86_AVX512_1x8x1_F64_F64, MMA_X86_AVX512_1x16x1_F32_F32, MMA_X86_AVX512_1x16x1_F32_F16_CASTF32, 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 38556164d18a..7d8e6498c08d 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h +++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h @@ -64,7 +64,11 @@ SmallVector getTilingLevelsAsInts(); StringRef getTilingLevelName(TilingLevel level); // Returns the TileSwizzle for the given intrinsic and operand index. -Codegen::TileSwizzle getIntrinsicSwizzle(MMAIntrinsic mma, int operandIdx); +// If `transposed` is true, the intrinsic is used in an M↔N-swapped +// orientation: the physical tile layouts reflect LHS/RHS roles being +// exchanged, and the accumulator is laid out column-major. +Codegen::TileSwizzle getIntrinsicSwizzle(MMAIntrinsic mma, bool transposed, + int operandIdx); // Returns the TileSwizzle for the given MMA attr and operand index. Codegen::TileSwizzle getSwizzle(DataTiledMMAAttr mma, int operandIdx); diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp index 6a07328b6abb..cb126fb4cd4a 100644 --- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp +++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp @@ -42,6 +42,7 @@ #include "iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.h" #include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h" +#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h" #include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" #include "iree/compiler/Codegen/ExternalInterfaces/Utils.h" @@ -50,9 +51,12 @@ #include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h" #include "iree/compiler/Dialect/Encoding/Utils/Utils.h" #include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h" +#include "llvm/ADT/StringExtras.h" #include "llvm/Support/DebugLog.h" #include "llvm/Support/InterleavedRange.h" +#include "mlir/IR/AffineMap.h" #include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" #define DEBUG_TYPE "iree-codegen-materialize-encoding" @@ -289,6 +293,288 @@ TileMxNxK chooseMatmulTile(ArrayRef enumeratedTiles, return bestRatedTile; } +static bool getEnableInnerTiledFromConfig(DictionaryAttr config) { + Attribute attr = config.get("enable_inner_tiled"); + if (auto battr = dyn_cast_if_present(attr)) { + return battr.getValue(); + } + return false; +} + +/// Returns the matmul {M, N, K} tile shape covered by a CPU DataTiledMMAAttr. +/// Derived directly from `getUndistributedTileTypes`: LHS is M×K (always) and +/// RHS is N×K (always), regardless of the `transposed_intrinsic` flag. The +/// ACC layout varies (M×N row-major vs. N×M column-major for transposed +/// intrinsics), so we derive M and N from the LHS and RHS tiles instead. +static IREE::Codegen::TileMxNxK getTileMxNxK(IREE::CPU::DataTiledMMAAttr mma) { + SmallVector tiles; + mma.getUndistributedTileTypes(tiles); + assert(tiles.size() == 3 && "Expected LHS, RHS, ACC tile types"); + ArrayRef lhsShape = tiles[0].getShape(); + ArrayRef rhsShape = tiles[1].getShape(); + return IREE::Codegen::TileMxNxK{lhsShape[0], rhsShape[0], lhsShape[1]}; +} + +/// Returns the set of `+feature` strings that must all be present in the +/// target config for `intr` to be usable. Most AVX-512 intrinsics map +/// 1:1 to a single `+feature`, but the AVX2 f32 MMA intrinsics require both +/// `+avx2` and `+fma` (FMA3 is a separate ISA extension from AVX2 itself, +/// and our AVX2 intrinsics lower to `vfmadd...`). +static SmallVector +getMmaIntrinsicRequiredFeatures(IREE::CPU::MMAIntrinsic intr) { + using IREE::CPU::MMAIntrinsic; + switch (intr) { + case MMAIntrinsic::MMA_X86_AVX2_FMA_1x8x1_F32_F32: + return {"+avx2", "+fma"}; + case MMAIntrinsic::MMA_X86_AVX512_1x8x1_F64_F64: + case MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F32: + case MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F16_CASTF32: + case MMAIntrinsic::MMA_X86_AVX512_1x16x2_I32_I16: + case MMAIntrinsic::MMA_X86_AVX512_1x16x2_I32_I8_CASTI16: + return {"+avx512f"}; + case MMAIntrinsic::MMA_X86_AVX512FP16_1x32x1_F16_F16: + return {"+avx512fp16"}; + case MMAIntrinsic::MMA_X86_AVX512BF16_1x16x2_F32_BF16: + return {"+avx512bf16"}; + case MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x2_I32_I16: + case MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x2_I32_I8_CASTI16: + return {"+avx512vnni"}; + default: + return {}; + } +} + +/// Returns x86 `MMAIntrinsic` cases whose required ISA extensions are all +/// present in `config` (`cpu_features` / target features). Only the "natural" +/// (M<=N) intrinsic orientations are listed; the M↔N-swapped orientation is +/// expressed by the `transposed_intrinsic` flag on DataTiledMMAAttr, enumerated +/// separately by the cost model. +static SmallVector +getMmaIntrinsicsForTargetConfig(DictionaryAttr config) { + using IREE::CPU::MMAIntrinsic; + SmallVector out; + if (!config) { + return out; + } + if (!isX86(config)) { + return out; + } + static const MMAIntrinsic kAllX86[] = { + MMAIntrinsic::MMA_X86_AVX2_FMA_1x8x1_F32_F32, + MMAIntrinsic::MMA_X86_AVX512_1x8x1_F64_F64, + MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F32, + MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F16_CASTF32, + MMAIntrinsic::MMA_X86_AVX512FP16_1x32x1_F16_F16, + MMAIntrinsic::MMA_X86_AVX512BF16_1x16x2_F32_BF16, + MMAIntrinsic::MMA_X86_AVX512_1x16x2_I32_I16, + MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x2_I32_I16, + MMAIntrinsic::MMA_X86_AVX512_1x16x2_I32_I8_CASTI16, + MMAIntrinsic::MMA_X86_AVX512VNNI_1x16x2_I32_I8_CASTI16, + }; + for (MMAIntrinsic intr : kAllX86) { + SmallVector required = getMmaIntrinsicRequiredFeatures(intr); + if (required.empty()) { + continue; + } + if (llvm::all_of(required, + [&](StringRef f) { return hasFeature(config, f); })) { + out.push_back(intr); + } + } + return out; +} + +/// Shape and element-bit-width summary of one intrinsic, with unroll factors +/// all 1, in a specific orientation. The LHS is (M, K), the RHS is (N, K) — +/// the `transposed` flag has already swapped M↔N inside the base attr. +/// Scalable dims (e.g. SVE) are treated as their 128-bit-minimum static +/// shape here; `getRegisterSpaceBytes` applies the matching simplification +/// on the capacity side. +struct IntrinsicInfo { + int64_t intrinsicM = 0, intrinsicN = 0, intrinsicK = 0; + int64_t lhsBits = 0, rhsBits = 0, accBits = 0; +}; + +/// Returns the IntrinsicInfo for `intr` in the given orientation, if its ABC +/// element types match `elementTypes`. Returns nullopt otherwise. +static std::optional +getIntrinsicInfo(MLIRContext *ctx, ArrayRef elementTypes, + IREE::CPU::MMAIntrinsic intr, bool transposed) { + auto base = IREE::CPU::DataTiledMMAAttr::get(ctx, intr, /*intrinsics_m=*/1, + /*intrinsics_n=*/1, + /*intrinsics_k=*/1, transposed); + SmallVector baseTiles; + base.getUndistributedTileTypes(baseTiles); + if (baseTiles.size() != 3) { + return std::nullopt; + } + if (baseTiles[0].getElementType() != elementTypes[0] || + baseTiles[1].getElementType() != elementTypes[1] || + baseTiles[2].getElementType() != elementTypes[2]) { + return std::nullopt; + } + IntrinsicInfo info; + info.intrinsicM = baseTiles[0].getShape()[0]; + info.intrinsicK = baseTiles[0].getShape()[1]; + info.intrinsicN = baseTiles[1].getShape()[0]; + info.lhsBits = baseTiles[0].getElementType().getIntOrFloatBitWidth(); + info.rhsBits = baseTiles[1].getElementType().getIntOrFloatBitWidth(); + info.accBits = baseTiles[2].getElementType().getIntOrFloatBitWidth(); + return info; +} + +/// Phase 1 of `chooseCpuInnerTiledMmaForEncoding`: jointly pick the +/// intrinsic and its orientation (`transposed_intrinsic`). We maximize +/// `usefulOps` per invocation, +/// usefulOps = min(intrinsicM, M) * min(intrinsicN, N) * intrinsicK +/// where the `min` clamps an intrinsicM/intrinsicN that exceeds a narrow +/// static matmul dim — i.e. it charges the intrinsic for the padding it +/// would force. A dynamic matmul dim is treated as "not narrow" (the full +/// intrinsic size counts toward usefulOps). Ties are broken by the raw +/// intrinsic size `intrinsicM*intrinsicN*intrinsicK` (prefer the bigger +/// intrinsic: it gives the Phase-2 register budget more tiles to play with) +/// and, finally, by iteration order (non-transposed first). `K` is never +/// narrow for optimization purposes, so it never enters the `min`. +/// +/// Returns nullopt if no compatible intrinsic exists. +static std::optional> +chooseIntrinsic(MLIRContext *ctx, ArrayRef elementTypes, + DictionaryAttr config, + const IREE::Encoding::BxMxNxKxKb &matmulSizes) { + auto usefulSize = [](int64_t matmulSize, int64_t intrinsicSize) -> int64_t { + return ShapedType::isDynamic(matmulSize) + ? intrinsicSize + : std::min(intrinsicSize, matmulSize); + }; + std::optional> best; + std::pair bestScore = {-1, -1}; + for (IREE::CPU::MMAIntrinsic intr : getMmaIntrinsicsForTargetConfig(config)) { + for (bool transposed : {false, true}) { + std::optional info = + getIntrinsicInfo(ctx, elementTypes, intr, transposed); + if (!info) { + continue; + } + int64_t usefulOps = usefulSize(matmulSizes.M, info->intrinsicM) * + usefulSize(matmulSizes.N, info->intrinsicN) * + info->intrinsicK; + int64_t rawOps = info->intrinsicM * info->intrinsicN * info->intrinsicK; + std::pair score = {usefulOps, rawOps}; + if (score > bestScore) { + bestScore = score; + best = {intr, transposed}; + } + } + } + return best; +} + +// Picks a CPU `DataTiledMMAAttr` for `iree_codegen.inner_tiled` given an +// encoding and target config. Picks the best (intrinsic, transposed) +// orientation via `chooseIntrinsic`, but leaves unroll factors at 1: an +// actual cost model for intrinsics_m/intrinsics_n is added in a follow-up +// commit. +static IREE::CPU::DataTiledMMAAttr +chooseCpuInnerTiledMmaForEncoding(MLIRContext *ctx, + IREE::Encoding::EncodingAttr encoding, + DictionaryAttr config) { + SmallVector elementTypes = encoding.getElementTypesArray(); + if (elementTypes.size() != 3) { + return {}; + } + FailureOr matmulSizes = + IREE::Encoding::getEncodingContractionLikeSizes(encoding); + if (failed(matmulSizes)) { + return {}; + } + std::optional> intrChoice = + chooseIntrinsic(ctx, elementTypes, config, *matmulSizes); + if (!intrChoice) { + return {}; + } + auto [intr, transposed] = *intrChoice; + return IREE::CPU::DataTiledMMAAttr::get(ctx, intr, /*intrinsics_m=*/1, + /*intrinsics_n=*/1, + /*intrinsics_k=*/1, transposed); +} + +/// Lowers a contraction under a `CPUEncodingResolverAttr` with +/// `enable_inner_tiled = true` to an `iree_codegen.inner_tiled` op whose kind +/// is a CPU `data_tiled_mma_layout`. Returns nullptr if no CPU MMA intrinsic +/// is available for the encoding/target. +static Operation *lowerContractionToInnerTiled( + OpBuilder &builder, linalg::LinalgOp linalgOp, ValueRange operands, + IREE::Encoding::LayoutMaterializerAttr layoutAttr) { + if (!linalgOp.hasPureTensorSemantics()) { + return nullptr; + } + + auto inputs = linalgOp.getDpsInputOperands(); + auto outputs = linalgOp.getDpsInits(); + + auto lhsType = cast(inputs[0]->get().getType()); + auto rhsType = cast(inputs[1]->get().getType()); + auto resultType = cast(outputs[0].getType()); + auto lhsEncoding = IREE::Encoding::getEncodingAttr(lhsType); + auto rhsEncoding = IREE::Encoding::getEncodingAttr(rhsType); + auto resultEncoding = IREE::Encoding::getEncodingAttr(resultType); + if (!lhsEncoding || !rhsEncoding || !resultEncoding) { + return nullptr; + } + + if (lhsEncoding.getOperandIndex().getValue() != IREE::Encoding::MATMUL_LHS || + rhsEncoding.getOperandIndex().getValue() != IREE::Encoding::MATMUL_RHS || + resultEncoding.getOperandIndex().getValue() != + IREE::Encoding::MATMUL_RESULT) { + return nullptr; + } + + auto cDims = linalg::inferContractionDims(linalgOp); + if (!cDims->batch.empty()) { + LDBG() << "inner_tiled lowering: batched contraction not implemented"; + return nullptr; + } + + MLIRContext *ctx = builder.getContext(); + Location loc = linalgOp.getLoc(); + AffineExpr d0 = builder.getAffineDimExpr(0); + AffineExpr d1 = builder.getAffineDimExpr(1); + AffineExpr d2 = builder.getAffineDimExpr(2); + SmallVector indexingMaps = { + AffineMap::get(3, 0, {d0, d2}, ctx), + AffineMap::get(3, 0, {d2, d1}, ctx), + AffineMap::get(3, 0, {d0, d1}, ctx), + }; + SmallVector iteratorTypes = { + utils::IteratorType::parallel, utils::IteratorType::parallel, + utils::IteratorType::reduction}; + + DictionaryAttr targetConfig; + if (auto cpuResolver = dyn_cast( + cast(layoutAttr))) { + targetConfig = cpuResolver.getConfiguration(); + } + + // Mirrors the GPU DataTiledMMA path: the MMA attribute is determined from + // the encoding + target config alone (the same inputs used by the encoding + // materialization side), not reverse-engineered from the already-packed + // operand tile shapes. + IREE::CPU::DataTiledMMAAttr chosenKind = + chooseCpuInnerTiledMmaForEncoding(ctx, resultEncoding, targetConfig); + if (!chosenKind) { + LDBG() << "inner_tiled lowering: no CPU DataTiledMMA kind available for " + "encoding/target"; + return nullptr; + } + + auto semanticsAttr = IREE::CPU::InnerTiledSemanticsAttr::get(ctx); + Operation *inner = IREE::Codegen::InnerTiledOp::create( + builder, loc, ValueRange{operands[0], operands[1]}, + ValueRange{operands[2]}, indexingMaps, iteratorTypes, chosenKind, + semanticsAttr); + return inner; +} + Operation *lowerContractionOpWithEncoding( OpBuilder &builder, linalg::LinalgOp linalgOp, ValueRange operands, IREE::Encoding::LayoutMaterializerAttr layoutAttr) { @@ -714,14 +1000,29 @@ struct CPUEncodingPackedLayoutMaterializerAttr return info; } + DictionaryAttr config = layoutAttr.getConfiguration(); + if (getEnableInnerTiledFromConfig(config)) { + return getInnerTiledEncodingInfo(type.getContext(), encoding, *cDims, + config); + } + return getMmt4dEncodingInfo(encoding, *cDims, config); + } + +private: + /// Legacy mmt4d path: enumerate candidate `TileMxNxK`s, pick one with the + /// narrow-dim-aware scoring, and apply the narrow-N→narrow-M transpose + /// trick required by the handwritten microkernels. + static MaterializeEncodingInfo + getMmt4dEncodingInfo(IREE::Encoding::EncodingAttr encoding, + const linalg::ContractionDimensions &cDims, + DictionaryAttr config) { + MaterializeEncodingInfo info; SmallVector enumeratedTileMxNxK = - enumerateCPUMatmulTiles(encoding, layoutAttr.getConfiguration()); + enumerateCPUMatmulTiles(encoding, config); if (enumeratedTileMxNxK.empty()) { return info; } auto narrowDim = IREE::Encoding::getPo2MatmulNarrowDim(encoding); - // Choose a final matmul TileMxNxK from the above-enumerated tile shapes, - // taking narrow dimensions into account. TileMxNxK chosenTileMxNxK = chooseMatmulTile(enumeratedTileMxNxK, narrowDim); FailureOr maybeEncodingInfo = @@ -731,7 +1032,7 @@ struct CPUEncodingPackedLayoutMaterializerAttr } info = std::move(maybeEncodingInfo.value()); FailureOr scalableFlags = - getScalableTileFlags(*cDims, encoding, layoutAttr.getConfiguration()); + getScalableTileFlags(cDims, encoding, config); if (succeeded(scalableFlags)) { info.scalableTiles = std::move(scalableFlags); } @@ -742,6 +1043,34 @@ struct CPUEncodingPackedLayoutMaterializerAttr } return info; } + + /// Intrinsic-first path for `iree_codegen.inner_tiled`: pick a + /// `DataTiledMMAAttr` (same decision the lowering will make) and derive the + /// packed-layout tile shape from it. No narrow-N transpose: unlike the + /// legacy mmt4d path, the inner_tiled lowering handles narrow dims natively + /// via `intrinsics_m` / `intrinsics_n`. + static MaterializeEncodingInfo getInnerTiledEncodingInfo( + MLIRContext *ctx, IREE::Encoding::EncodingAttr encoding, + const linalg::ContractionDimensions &cDims, DictionaryAttr config) { + MaterializeEncodingInfo info; + IREE::CPU::DataTiledMMAAttr mma = + chooseCpuInnerTiledMmaForEncoding(ctx, encoding, config); + if (!mma) { + return info; + } + FailureOr maybeEncodingInfo = + getEncodingInfoForMatmul(encoding, getTileMxNxK(mma)); + if (failed(maybeEncodingInfo)) { + return info; + } + info = std::move(maybeEncodingInfo.value()); + FailureOr scalableFlags = + getScalableTileFlags(cDims, encoding, config); + if (succeeded(scalableFlags)) { + info.scalableTiles = std::move(scalableFlags); + } + return info; + } }; struct CPUEncodingResolverMaterializerAttr final @@ -769,6 +1098,12 @@ struct CPUEncodingResolverMaterializerAttr final convertedOperands.drop_front(numInputs)); } if (linalg::isaContractionOpInterface(linalgOp)) { + DictionaryAttr config = layoutAttr.getConfiguration(); + if (getEnableInnerTiledFromConfig(config)) { + return lowerContractionToInnerTiled( + b, linalgOp, convertedOperands, + cast(layoutAttr)); + } return lowerContractionOpWithEncoding( b, linalgOp, convertedOperands, cast(layoutAttr)); @@ -796,6 +1131,7 @@ struct CPULayoutResolverAttr final addConfigTargetTriple(ctx, targetTriple.value(), configItems); } storeNamedAttrIfPresent(configItems, config, "ukernels"); + storeNamedAttrIfPresent(configItems, config, "enable_inner_tiled"); return CPUEncodingResolverAttr::get(ctx, DictionaryAttr::get(ctx, configItems)); } diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/test/roundtrip.mlir b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/test/roundtrip.mlir index 69f3e43a00d8..ddc0a01a7b4b 100644 --- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/test/roundtrip.mlir +++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/test/roundtrip.mlir @@ -93,6 +93,16 @@ func.func @cpu_valid_2d_encoding(%arg0: tensor<32x64xf32, #cpu_encoding>) -> ten // ----- +// CPU resolver configuration can carry enable_inner_tiled (copied from executable target during specialization). +#cpu_encoding_inner_tiled = #iree_cpu.cpu_encoding_resolver +func.func @cpu_encoding_with_inner_tiled_flag(%arg0: tensor) -> tensor { + return %arg0 : tensor +} +// CHECK-LABEL: func.func @cpu_encoding_with_inner_tiled_flag +// CHECK-SAME: tensor> + +// ----- + // VMVX encoding resolver tests. // Note: Most encoding verifier tests are done on gpu_encoding_resolver above. // This test just verifies that vmvx_encoding_resolver roundtrips correctly.