[Mlir-commits] [mlir] [mlir][x86] Add pattern to tile and unroll vector.contract (PR #229076)
Julian Oppermann
llvmlistbot at llvm.org
Fri Oct 9 01:08:43 PDT 2026
https://github.com/jopperm updated https://github.com/llvm/llvm-project/pull/229076
>From b0d4d7935e8c7b5c1008bcd3aa888689db87ce6d Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Fri, 25 Sep 2026 01:14:01 -0700
Subject: [PATCH 01/15] [mlir][x86] Add pattern to tile and unroll
vector.contract
---
.../X86/TransformOps/X86TransformOps.td | 13 +
mlir/include/mlir/Dialect/X86/Transforms.h | 5 +
.../X86/TransformOps/X86TransformOps.cpp | 5 +
.../lib/Dialect/X86/Transforms/CMakeLists.txt | 1 +
.../VectorContractMultiLevelUnroll.cpp | 796 +++++++++++
.../vector-contract-multi-level-unroll.mlir | 1239 +++++++++++++++++
6 files changed, 2059 insertions(+)
create mode 100644 mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
create mode 100644 mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
diff --git a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
index 48aabf90914e24..a1a6f99a9447ca 100644
--- a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
+++ b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
@@ -93,4 +93,17 @@ def ApplyMoveAccumulatorForContractLoopPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplyVectorContractMultiLevelUnrollPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.x86.vector_contract_multi_level_unroll",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Collect patterns that tile and unroll contractions in order to prepare them
+ for the nanokernel lowerings.
+ }];
+
+ let arguments = (ins StrAttr:$target);
+
+ let assemblyFormat = "`target` `=` $target attr-dict";
+}
+
#endif // X86_TRANSFORM_OPS
diff --git a/mlir/include/mlir/Dialect/X86/Transforms.h b/mlir/include/mlir/Dialect/X86/Transforms.h
index 8bcaf9a67feb0e..676b2f85f2c0c5 100644
--- a/mlir/include/mlir/Dialect/X86/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86/Transforms.h
@@ -115,6 +115,11 @@ void populateVectorContractToAMXDotProductPatterns(RewritePatternSet &patterns);
void populateMoveAccumulatorForContractLoopPatterns(
RewritePatternSet &patterns);
+// A set of patterns that tile and unroll contractions in order to prepare them
+// for the nanokernel lowerings.
+void populateVectorContractMultiLevelUnrollPatterns(RewritePatternSet &patterns,
+ StringRef target);
+
//===----------------------------------------------------------------------===//
/// Helpers extracted from:
/// - clang/lib/Headers/avxintrin.h
diff --git a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
index a52f08a5cd0617..ff05e38913d335 100644
--- a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
+++ b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
@@ -57,6 +57,11 @@ void mlir::transform::ApplyMoveAccumulatorForContractLoopPatternsOp::
x86::populateMoveAccumulatorForContractLoopPatterns(patterns);
}
+void mlir::transform::ApplyVectorContractMultiLevelUnrollPatternsOp::
+ populatePatterns(RewritePatternSet &patterns) {
+ x86::populateVectorContractMultiLevelUnrollPatterns(patterns, getTarget());
+}
+
//===----------------------------------------------------------------------===//
// Transform op registration
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt b/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
index 5118563c03f164..38e80c40f356ed 100644
--- a/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
@@ -8,6 +8,7 @@ add_mlir_dialect_library(MLIRX86Transforms
ShuffleVectorFMAOps.cpp
VectorContractToAMXDotProduct.cpp
MoveAccumulatorForContractLoop.cpp
+ VectorContractMultiLevelUnroll.cpp
LINK_LIBS PUBLIC
MLIRArithDialect
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
new file mode 100644
index 00000000000000..599363724e7def
--- /dev/null
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -0,0 +1,796 @@
+//===- VectorContractMultiLevelUnroll.cpp ---------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Dialect/X86/Transforms.h"
+#include "mlir/Dialect/X86/Utils/X86Utils.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
+
+#define DEBUG_TYPE "x86-vector-contract-multi-level-unroll"
+
+using namespace mlir;
+
+using llvm::all_of;
+using llvm::count;
+
+// Discardable attribute to control vector unroll patterns.
+static constexpr auto nativeShapeAttrName = "x86_vcmlu_native_shape";
+
+namespace {
+// Shared state for the pattern match and IR transformation.
+struct MLUCandidate {
+ StringRef target;
+ vector::ContractionOp contract;
+ bool isInVnniLayout = false;
+
+ // Shape of the original contraction.
+ VectorType accType;
+ VectorType lhsType;
+ VectorType rhsType;
+
+ // Shape of the register-tiled contraction. We'll insert a loop nest to cover
+ // the original shape.
+ VectorType accRegTileType;
+ VectorType lhsRegTileType;
+ VectorType rhsRegTileType;
+
+ // Shape of the native contraction (for the X86 nanokernel patterns). We'll
+ // unroll the register-tiled contraction to this shape.
+ VectorType accNativeType;
+ VectorType lhsNativeType;
+ VectorType rhsNativeType;
+ SmallVector<int64_t, 4> contractNativeShape;
+
+ // Bookkeeping for the accumulation loop and operands.
+ scf::ForOp accLoop;
+ Value accInitVal;
+
+ bool accIsZeroInit = false;
+ vector::TransferReadOp lhsRead;
+ vector::TransferReadOp rhsRead;
+
+ // Memory backing the accumulator.
+ Value accMemIn;
+ Value accMemOut;
+
+ // Reconstructed accumulation loop with native-shaped iter_args.
+ scf::ForOp slicedAccLoop;
+
+ // Keep track of operations that should be deleted after transformation
+ // (depending on which ops could be integrated into the new loop nest).
+ SmallVector<Operation *, 2> opsToDelete;
+};
+} // anonymous namespace
+
+//----------------------------------------------------------------------------//
+// Utilities //
+//----------------------------------------------------------------------------//
+
+// Return true if given vector.transfer_read/write has a default permutation
+// map, no mask and no out-of-bounds check.
+template <typename TransferOp>
+static bool isSimpleTransferOp(TransferOp op) {
+ return op.getPermutationMap().isMinorIdentity() && !op.getMask() &&
+ !op.hasOutOfBoundsDim();
+}
+
+// Create a rank-reducing subview of a memref that matches the target vector
+// type, e.g. with targetType = vector<64x32xbf16>:
+//
+// memref.subview %A[%block_m, %block_k, 0, 0] [1, 1, 64, 32] [1, 1, 1, 1]
+// : memref<?x?x64x32xbf16> to
+// memref<64x32xbf16, strided<[32, 1], offset: ?>>
+static Value makeRankReducingSubview(OpBuilder &b, Location loc, Value base,
+ ValueRange indices,
+ VectorType targetType) {
+ auto *ctx = b.getContext();
+
+ auto baseType = dyn_cast<MemRefType>(base.getType());
+ assert(baseType && "base must be a MemRefType");
+
+ int64_t baseRank = baseType.getRank();
+ int64_t targetRank = targetType.getRank();
+ assert(targetRank <= baseRank && "target shape's rank > memref's rank");
+
+ ArrayRef<int64_t> targetShape = targetType.getShape();
+
+ SmallVector<int64_t> sizesInt(baseType.getRank(), 1);
+ copy(targetShape, sizesInt.begin() + baseRank - targetRank);
+ SmallVector<int64_t> stridesInt(baseRank, 1);
+
+ auto offsets = getAsOpFoldResult(indices);
+ auto sizes = getAsIndexOpFoldResult(ctx, sizesInt);
+ auto strides = getAsIndexOpFoldResult(ctx, stridesInt);
+
+ MemRefType rankReducedType = memref::SubViewOp::inferRankReducedResultType(
+ targetShape, baseType, offsets, sizes, strides);
+ return memref::SubViewOp::create(b, loc, rankReducedType, base, offsets,
+ sizes, strides);
+}
+
+// Slice a 2D vector into smaller tiles of the given native type.
+static SmallVector<Value> slice(OpBuilder &b, Location loc, Value value,
+ VectorType nativeType) {
+ VectorType regTileType = cast<VectorType>(value.getType());
+ ArrayRef<int64_t> regTileShape = regTileType.getShape();
+ ArrayRef<int64_t> nativeShape = nativeType.getShape();
+
+ assert(regTileShape.size() == 2 && nativeShape.size() == 2 &&
+ "Only 2D shapes are supported");
+
+ int64_t regTileM = regTileShape[0];
+ int64_t regTileN = regTileShape[1];
+ int64_t nativeM = nativeShape[0];
+ int64_t nativeN = nativeShape[1];
+
+ SmallVector<Value> result;
+ result.reserve((regTileM / nativeM) * (regTileN / nativeN));
+
+ for (int64_t m = 0; m < regTileM; m += nativeM)
+ for (int64_t n = 0; n < regTileN; n += nativeN)
+ result.push_back(vector::ExtractStridedSliceOp::create(
+ b, loc, value, ArrayRef<int64_t>{m, n}, nativeShape,
+ ArrayRef<int64_t>{1, 1}));
+
+ return result;
+}
+
+// Splice smaller tiles back into a 2D vector of the given register tile type.
+static Value splice(OpBuilder &b, Location loc, ValueRange values,
+ VectorType regTileType) {
+ auto types = values.getTypes();
+ assert(!types.empty() &&
+ count(types, types.front()) == static_cast<int64_t>(types.size()) &&
+ "All values must have the same type");
+
+ VectorType nativeType = cast<VectorType>(types.front());
+
+ ArrayRef<int64_t> nativeShape = nativeType.getShape();
+ ArrayRef<int64_t> regTileShape = regTileType.getShape();
+
+ assert(nativeShape.size() == 2 && regTileShape.size() == 2 &&
+ "Only 2D shapes are supported");
+
+ int64_t nativeM = nativeShape[0];
+ int64_t nativeN = nativeShape[1];
+ int64_t regTileM = regTileShape[0];
+ int64_t regTileN = regTileShape[1];
+
+ assert((regTileM / nativeM) * (regTileN / nativeN) ==
+ static_cast<int64_t>(values.size()) &&
+ "The number of values must match the number of slices");
+
+ Value result = ub::PoisonOp::create(b, loc, regTileType);
+
+ auto valueIt = values.begin();
+ for (int64_t m = 0; m < regTileM; m += nativeM)
+ for (int64_t n = 0; n < regTileN; n += nativeN) {
+ result = vector::InsertStridedSliceOp::create(b, loc, *valueIt, result,
+ ArrayRef<int64_t>{m, n},
+ ArrayRef<int64_t>{1, 1});
+ ++valueIt;
+ }
+
+ return result;
+}
+
+static LogicalResult hasNativeShapeAttr(Operation *op) {
+ return success(op->hasAttr(nativeShapeAttrName));
+}
+
+static std::optional<SmallVector<int64_t>>
+getFromNativeShapeAttr(Operation *op) {
+ auto attr = op->getAttrOfType<DenseI64ArrayAttr>(nativeShapeAttrName);
+ if (!attr)
+ return std::nullopt;
+ return SmallVector<int64_t>(attr.asArrayRef());
+}
+
+static void setNativeShapeAttr(Operation *op, VectorType nativeType) {
+ VectorType currentType;
+ TypeSwitch<Operation *>(op)
+ .Case<vector::TransferReadOp, vector::TransferWriteOp>(
+ [&](auto transferOp) { currentType = transferOp.getVectorType(); })
+ .DefaultUnreachable("expected a vector transfer op");
+ if (currentType != nativeType)
+ op->setAttr(
+ nativeShapeAttrName,
+ DenseI64ArrayAttr::get(nativeType.getContext(), nativeType.getShape()));
+}
+
+//----------------------------------------------------------------------------//
+// Pattern match and target shape computation //
+//----------------------------------------------------------------------------//
+
+// Establish whether the candidate contraction is a canonical matrix
+// multiplication on flat or VNNI-packed vectors.
+//
+// Example (flat):
+// vector.contract {
+// kind = #vector.kind<add>, indexing_maps = [
+// affine_map<(dm, dn, dk) -> (dm, dk)>,
+// affine_map<(dm, dn, dk) -> (dk, dn)>,
+// affine_map<(dm, dn, dk) -> (dm, dn)>],
+// iterator_types = ["parallel", "parallel", "reduction"],
+// } %lhs, %rhs, %acc : !lhsType, !rhsType into !accType
+//
+// Example (VNNI-packed):
+// vector.contract {
+// kind = #vector.kind<add>, indexing_maps = [
+// affine_map<(dm, dn, dk, dvnni) -> (dm, dk, dvnni)>,
+// affine_map<(dm, dn, dk, dvnni) -> (dk, dn, dvnni)>,
+// affine_map<(dm, dn, dk, dvnni) -> (dm, dn)>],
+// iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+// } %lhs, %rhs, %acc : !lhsType, !rhsType into !accType
+static LogicalResult checkCanonicalMatmul(MLUCandidate &candidate,
+ PatternRewriter &rewriter) {
+ auto &contract = candidate.contract;
+ if (contract.getKind() != vector::CombiningKind::ADD)
+ return rewriter.notifyMatchFailure(contract,
+ "not using ADD combining kind");
+
+ VectorType accType = dyn_cast<VectorType>(contract.getAccType());
+ if (!accType || accType.getRank() != 2)
+ return rewriter.notifyMatchFailure(contract,
+ "accumulator is not a 2D vector");
+
+ VectorType lhsType = contract.getLhsType();
+ VectorType rhsType = contract.getRhsType();
+ if (lhsType.getElementType() != rhsType.getElementType())
+ return rewriter.notifyMatchFailure(
+ contract, "LHS and RHS element types do not match");
+
+ auto iteratorTypes = contract.getIteratorTypesArray();
+ auto indexingMaps = contract.getIndexingMapsArray();
+ assert(indexingMaps.size() == 3 && "expected 3 indexing maps");
+ bool isInVnniLayout = x86::isInVnniLayout(contract, indexingMaps);
+
+ auto par = vector::IteratorType::parallel;
+ auto red = vector::IteratorType::reduction;
+ auto m = rewriter.getAffineDimExpr(0);
+ auto n = rewriter.getAffineDimExpr(1);
+ auto k = rewriter.getAffineDimExpr(2);
+ auto vnni = rewriter.getAffineDimExpr(3);
+
+ if (!isInVnniLayout) {
+ if (!equal(indexingMaps[0].getResults(), ArrayRef{m, k}) ||
+ !equal(indexingMaps[1].getResults(), ArrayRef{k, n}) ||
+ !equal(indexingMaps[2].getResults(), ArrayRef{m, n}))
+ return rewriter.notifyMatchFailure(contract,
+ "indexing maps are not canonical");
+
+ if (!equal(iteratorTypes, ArrayRef{par, par, red}))
+ return rewriter.notifyMatchFailure(contract,
+ "iterator types are not canonical");
+ } else {
+ if (!equal(indexingMaps[0].getResults(), ArrayRef{m, k, vnni}) ||
+ !equal(indexingMaps[1].getResults(), ArrayRef{k, n, vnni}) ||
+ !equal(indexingMaps[2].getResults(), ArrayRef{m, n}))
+ return rewriter.notifyMatchFailure(
+ contract, "indexing maps are not canonical (VNNI)");
+
+ if (!equal(iteratorTypes, ArrayRef{par, par, red, red}))
+ return rewriter.notifyMatchFailure(
+ contract, "iterator types are not canonical (VNNI)");
+ }
+
+ candidate.isInVnniLayout = isInVnniLayout;
+ candidate.accType = accType;
+ candidate.lhsType = lhsType;
+ candidate.rhsType = rhsType;
+
+ return success();
+}
+
+// Check compatibility of datatypes with the given target ISA extension, and try
+// to determine a suitable tiling strategy for it.
+static LogicalResult checkShapesAndTypes(MLUCandidate &candidate,
+ PatternRewriter &rewriter) {
+ auto &contract = candidate.contract;
+ auto target = candidate.target;
+
+ Type accElemType = candidate.accType.getElementType();
+ Type inpElemType = candidate.lhsType.getElementType();
+
+ ArrayRef<int64_t> accShape = candidate.accType.getShape();
+ ArrayRef<int64_t> lhsShape = candidate.lhsType.getShape();
+
+ int64_t origM = accShape[0];
+ int64_t origN = accShape[1];
+ int64_t origK = lhsShape[1];
+ int64_t origVnni = candidate.isInVnniLayout ? lhsShape[2] : -1;
+
+ int64_t regTileM, regTileN, regTileK;
+ int64_t nativeM, nativeN, nativeK, vnni;
+
+ if (target.starts_with("amx")) {
+ if (!((target == "amx_bf16" && inpElemType.isBF16() &&
+ accElemType.isF32()) ||
+ (target == "amx_int8" && inpElemType.isInteger(8) &&
+ accElemType.isInteger(32))))
+ return rewriter.notifyMatchFailure(
+ contract, "unsupported combination of input and accumulator types");
+
+ nativeM = nativeN = 16;
+ vnni = 32 / inpElemType.getIntOrFloatBitWidth();
+
+ if (candidate.isInVnniLayout) {
+ nativeK = 16;
+ if (origM % nativeM != 0 || origN % nativeN != 0)
+ return rewriter.notifyMatchFailure(
+ contract, "vector shape cannot be cleanly unrolled");
+ if (origK != nativeK || origVnni != vnni)
+ return rewriter.notifyMatchFailure(
+ contract, "K dimension or VNNI factor mismatch");
+ if (origM == nativeM && origN == nativeN)
+ return rewriter.notifyMatchFailure(contract,
+ "already in native tile size");
+ // Prefer 2x2 tiling if shapes allow.
+ regTileM = origM % (2 * nativeM) == 0 ? 2 * nativeM : nativeM;
+ regTileN = origN % (2 * nativeN) == 0 ? 2 * nativeN : nativeN;
+ regTileK = nativeK;
+ } else {
+ // Online packing requires a 2x2 tile register layout.
+ nativeK = 16 * vnni;
+ regTileM = 2 * nativeM;
+ regTileN = 2 * nativeN;
+ regTileK = nativeK;
+ if (origM % regTileM != 0 || origN % regTileN != 0)
+ return rewriter.notifyMatchFailure(
+ contract, "vector shape cannot be cleanly unrolled");
+ if (origK != nativeK)
+ return rewriter.notifyMatchFailure(contract, "K dimension mismatch");
+ }
+ } else
+ return rewriter.notifyMatchFailure(contract, "unsupported target");
+
+ candidate.accRegTileType = VectorType::get({regTileM, regTileN}, accElemType);
+ candidate.accNativeType = VectorType::get({nativeM, nativeN}, accElemType);
+
+ auto getOperandType = [&](ArrayRef<int64_t> shape) {
+ if (candidate.isInVnniLayout) {
+ SmallVector<int64_t, 3> newShape(shape.begin(), shape.end());
+ newShape.push_back(vnni);
+ return VectorType::get(newShape, inpElemType);
+ }
+ return VectorType::get(shape, inpElemType);
+ };
+
+ candidate.lhsRegTileType = getOperandType({regTileM, regTileK});
+ candidate.rhsRegTileType = getOperandType({regTileK, regTileN});
+
+ candidate.lhsNativeType = getOperandType({nativeM, nativeK});
+ candidate.rhsNativeType = getOperandType({nativeK, nativeN});
+
+ candidate.contractNativeShape = {nativeM, nativeN, nativeK};
+ if (candidate.isInVnniLayout)
+ candidate.contractNativeShape.push_back(vnni);
+
+ return success();
+}
+
+// Check constraints on the accumulation loop. We're looking for an index-based
+// for-loop with a step of 1 or K, a single accumulator block argument. The body
+// shall only contain the transfer_reads for the operands, the contraction and
+// the yield operation. This restriction is in place because the pattern
+// rebuilds the accumulation loop from scratch with the known operations.
+//
+// For the operand reads, we require that they have a minor identity permutation
+// map, impose no bounds checks and don't carry a mask. In addition, they shall
+// use the induction variable as exactly on of their indices. This covers plain
+// 2D memrefs as well as block-based layouts.
+//
+// Example:
+// !accType = vector<64x128xf32>
+// !lhsType = vector<64x32xbf16>
+// !rhsType = vector<32x128xbf16>
+// ...
+// %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c)
+// -> (!accType) {
+// %a = vector.transfer_read %A[%m, %k, %c0, %c0], %c0_0
+// {in_bounds = [true, true]} : memref<?x?x64x32xbf16>, !lhsType
+// %b = vector.transfer_read %B[%n, %k, %c0, %c0], %c0_0
+// {in_bounds = [true, true]} : memref<?x?x32x128xbf16>, !rhsType
+// %d = vector.contract { ... } %a, %b, %acc
+// : !lhsType, !rhsType into !accType
+// scf.yield %d : !accType
+// }
+static LogicalResult checkAccumulationLoop(MLUCandidate &candidate,
+ PatternRewriter &rewriter) {
+ auto &contract = candidate.contract;
+
+ auto accLoop = dyn_cast_if_present<scf::ForOp>(contract->getParentOp());
+ if (!accLoop || !accLoop.getInductionVar().getType().isIndex())
+ return rewriter.notifyMatchFailure(contract,
+ "is not in an index-based for-loop");
+
+ if (auto step = getConstantIntValue(accLoop.getStep());
+ !step || (*step != 1 && *step != candidate.lhsType.getShape()[1]))
+ return rewriter.notifyMatchFailure(contract,
+ "accumulation loop step is not 1 or K");
+
+ BlockArgument accIterArg = dyn_cast<BlockArgument>(contract.getAcc());
+ if (!accIterArg || accIterArg.getOwner() != accLoop.getBody())
+ return rewriter.notifyMatchFailure(contract,
+ "accumulator is not a block argument");
+
+ if (accLoop.getNumRegionIterArgs() != 1 ||
+ accLoop.getBody()->getOperations().size() != 4)
+ return rewriter.notifyMatchFailure(
+ contract, "accumulation loop has additional iter_args or operations");
+
+ if (!contract->getResult(0).hasOneUse())
+ return rewriter.notifyMatchFailure(contract,
+ "result does not have exactly one use");
+
+ Operation *user = *contract->getResult(0).getUsers().begin();
+ if (!isa<scf::YieldOp>(user))
+ return rewriter.notifyMatchFailure(contract,
+ "accumulator is not loop-carried");
+
+ bool accIsZeroInit = isa_and_present<arith::ConstantOp>(
+ x86::traceToVectorReadLikeParentOperation(contract.getAcc()));
+
+ auto lhsRead = contract.getLhs().getDefiningOp<vector::TransferReadOp>();
+ auto rhsRead = contract.getRhs().getDefiningOp<vector::TransferReadOp>();
+ if (!lhsRead || !rhsRead)
+ return rewriter.notifyMatchFailure(
+ contract, "LHS/RHS are not vector.transfer_read ops");
+
+ auto checkTransferOp = [&](vector::TransferReadOp transferOp) {
+ if (!isSimpleTransferOp(transferOp))
+ return rewriter.notifyMatchFailure(contract, [&](Diagnostic &diag) {
+ diag << "transfer op has non-identity permutation, mask, or "
+ "out-of-bounds check: "
+ << *transferOp;
+ });
+
+ if (count(transferOp.getIndices(), accLoop.getInductionVar()) != 1)
+ return rewriter.notifyMatchFailure(
+ contract, "transfer op does not use accumulator index");
+
+ return success();
+ };
+
+ if (failed(checkTransferOp(lhsRead)) || failed(checkTransferOp(rhsRead)))
+ return failure();
+
+ candidate.accLoop = accLoop;
+ candidate.accInitVal = accLoop.getTiedLoopInit(accIterArg)->get();
+ candidate.accIsZeroInit = accIsZeroInit;
+ candidate.lhsRead = lhsRead;
+ candidate.rhsRead = rhsRead;
+
+ return success();
+}
+
+//----------------------------------------------------------------------------//
+// Transformation
+//----------------------------------------------------------------------------//
+
+// Ensure that the initial accumulator value can be sliced into register-tile
+// shaped vectors. Unless it is zero-initialized, we need it to be backed
+// by memory so that the slice can be loaded via a vector.transfer_read inside
+// the loop nest.
+//
+// If the initial value is produced by a simple vector.transfer_read, we can
+// simply derive a subview from its base and indices. Otherwise, we fall back to
+// allocating a new buffer on stack and store the initial value into it.
+static void makeAccInitializationBackedByMemory(MLUCandidate &candidate,
+ PatternRewriter &rewriter) {
+ if (candidate.accIsZeroInit)
+ return;
+
+ rewriter.setInsertionPoint(candidate.accLoop);
+ auto loc = candidate.accLoop.getLoc();
+
+ if (auto accRead =
+ candidate.accInitVal.getDefiningOp<vector::TransferReadOp>();
+ accRead && isSimpleTransferOp(accRead)) {
+ candidate.accMemIn =
+ makeRankReducingSubview(rewriter, loc, accRead.getBase(),
+ accRead.getIndices(), accRead.getVectorType());
+ return;
+ }
+
+ auto accBufferType = MemRefType::get(candidate.accType.getShape(),
+ candidate.accType.getElementType());
+ candidate.accMemIn = memref::AllocaOp::create(rewriter, loc, accBufferType);
+
+ Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ SmallVector<Value, 2> zeroIndices(candidate.accType.getRank(), c0);
+ vector::TransferWriteOp::create(rewriter, loc, candidate.accInitVal,
+ candidate.accMemIn, zeroIndices);
+}
+
+// Ensure that slices of the result of the accumulator loop can be stored, so
+// again we need it to be backed by memory.
+//
+// If the only user of the result is a simple vector.transfer_write, and its
+// base and indices are defined before the accumulator loop, we can derive a
+// subview. The transfer_write will be deleted after building the loop nest,
+// which at that point incorporates its effect.
+//
+// Otherwise, we fall back to buffering the result and rematerializing it after
+// the loop. If we already had a stack allocation for the initial value, we
+// reuse it.
+static void makeAccResultBackedByMemory(MLUCandidate &candidate,
+ PatternRewriter &rewriter) {
+ rewriter.setInsertionPoint(candidate.accLoop);
+ auto loc = candidate.accLoop.getLoc();
+
+ auto accRes = candidate.accLoop.getResult(0);
+ vector::TransferWriteOp accWrite;
+ if (accRes.hasOneUse())
+ accWrite = dyn_cast<vector::TransferWriteOp>(*accRes.getUsers().begin());
+
+ Block *accLoopBlock = candidate.accLoop->getBlock();
+ if (accWrite && isSimpleTransferOp(accWrite) &&
+ accWrite->getBlock() == accLoopBlock) {
+ auto isDefinedBeforeAccLoop = [&](Value v) {
+ Operation *defOp = v.getDefiningOp();
+ // No need for full dominance check, as we already know that accLoop and
+ // accWrite are in the same block.
+ return !defOp || defOp->getBlock() != accLoopBlock ||
+ defOp->isBeforeInBlock(candidate.accLoop);
+ };
+ if (isDefinedBeforeAccLoop(accWrite.getBase()) &&
+ all_of(accWrite.getIndices(), isDefinedBeforeAccLoop)) {
+ candidate.accMemOut = makeRankReducingSubview(
+ rewriter, loc, accWrite.getBase(), accWrite.getIndices(),
+ accWrite.getVectorType());
+ candidate.opsToDelete.push_back(accWrite);
+ candidate.opsToDelete.push_back(candidate.accLoop);
+ return;
+ }
+ }
+
+ if (!candidate.accIsZeroInit &&
+ isa<memref::AllocaOp>(candidate.accMemIn.getDefiningOp())) {
+ candidate.accMemOut = candidate.accMemIn;
+ } else {
+ auto accBufferType = MemRefType::get(candidate.accType.getShape(),
+ candidate.accType.getElementType());
+ candidate.accMemOut =
+ memref::AllocaOp::create(rewriter, loc, accBufferType);
+ }
+
+ rewriter.setInsertionPointAfter(candidate.accLoop);
+ Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ SmallVector<Value, 2> zeroIndices(candidate.accType.getRank(), c0);
+ auto accRemat = vector::TransferReadOp::create(
+ rewriter, loc, candidate.accType, candidate.accMemOut, zeroIndices,
+ /*padding=*/std::nullopt);
+
+ rewriter.replaceAllUsesWith(candidate.accLoop.getResult(0), accRemat);
+ candidate.opsToDelete.push_back(candidate.accLoop);
+}
+
+static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
+ Value ivM, Value ivN);
+static SmallVector<Value> buildKLoopBody(MLUCandidate &candidate, OpBuilder &b,
+ Location loc, Value ivM, Value ivN,
+ Value ivK, ValueRange iterArgs);
+
+// Create the loop nest for iterating over the M and N dimensions of the
+// accumulator. Construction continues in buildMNLoopBody.
+static void buildLoopNestForCandidate(MLUCandidate &candidate,
+ PatternRewriter &rewriter) {
+ rewriter.setInsertionPoint(candidate.accLoop);
+ auto loc = candidate.accLoop.getLoc();
+
+ ArrayRef<int64_t> accShape = candidate.accType.getShape();
+ ArrayRef<int64_t> accRegTileShape = candidate.accRegTileType.getShape();
+
+ auto makeIndexConst = [&](int64_t v) {
+ return arith::ConstantIndexOp::create(rewriter, loc, v);
+ };
+
+ Value c0 = makeIndexConst(0);
+ Value ubM = makeIndexConst(accShape[0]);
+ Value stepM = makeIndexConst(accRegTileShape[0]);
+ Value ubN = makeIndexConst(accShape[1]);
+ Value stepN = makeIndexConst(accRegTileShape[1]);
+
+ scf::ForOp::create(
+ rewriter, loc, c0, ubM, stepM, {},
+ [&](OpBuilder &b, Location loc, Value ivM, ValueRange iterArgs) {
+ scf::ForOp::create(
+ b, loc, c0, ubN, stepN, {},
+ [&](OpBuilder &b, Location loc, Value ivN, ValueRange iterArgs) {
+ buildMNLoopBody(candidate, b, loc, ivM, ivN);
+ scf::YieldOp::create(b, loc);
+ });
+ scf::YieldOp::create(b, loc);
+ });
+}
+
+// Create the body for M-N-loop nest according to the following scheme:
+//
+// 1. Load a accRegTileType-shaped slice of the accumulator from memory (or make
+// a zero constant).
+// 2. Split the initial value into accNativeType-shaped slices.
+// 3. Create the K-loop with the initial value slices as loop-carried arguments,
+// i.e. the iter_args have accNativeType.
+// 4. -> buildKLoopBody
+// 5. Splice the final value slices back into an accRegTileType-shaped value and
+// write it back to memory.
+//
+// The accRegTileType-shaped operations are annotated with an attribute in order
+// to unroll them with the vector-dialect's unroll patterns.
+static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
+ Value ivM, Value ivN) {
+ ArrayRef<int64_t> accRegTileShape = candidate.accRegTileType.getShape();
+
+ Value c0 = arith::ConstantIndexOp::create(b, loc, 0);
+
+ SmallVector<OpFoldResult, 2> accViewOffsets = {ivM, ivN};
+ SmallVector<OpFoldResult, 2> accViewSizes = {
+ b.getI64IntegerAttr(accRegTileShape[0]),
+ b.getI64IntegerAttr(accRegTileShape[1])};
+ SmallVector<OpFoldResult, 2> accViewStrides = {b.getI64IntegerAttr(1),
+ b.getI64IntegerAttr(1)};
+
+ Operation *accInit;
+ if (candidate.accIsZeroInit) {
+ accInit = arith::ConstantOp::create(
+ b, loc, b.getZeroAttr(candidate.accRegTileType));
+ } else {
+ Value accViewIn =
+ memref::SubViewOp::create(b, loc, candidate.accMemIn, accViewOffsets,
+ accViewSizes, accViewStrides);
+ accInit = vector::TransferReadOp::create(b, loc, candidate.accRegTileType,
+ accViewIn, ValueRange{c0, c0},
+ /*padding=*/std::nullopt);
+ setNativeShapeAttr(accInit, candidate.accNativeType);
+ }
+
+ SmallVector<Value> accInitSliced =
+ slice(b, loc, accInit->getResult(0), candidate.accNativeType);
+ candidate.slicedAccLoop = scf::ForOp::create(
+ b, loc, candidate.accLoop.getLowerBound(),
+ candidate.accLoop.getUpperBound(), candidate.accLoop.getStep(),
+ accInitSliced,
+ [&](OpBuilder &b, Location loc, Value ivK, ValueRange iterArgs) {
+ SmallVector<Value> newIterArgs =
+ buildKLoopBody(candidate, b, loc, ivM, ivN, ivK, iterArgs);
+ scf::YieldOp::create(b, loc, newIterArgs);
+ });
+
+ Value accSpliced = splice(b, loc, candidate.slicedAccLoop.getResults(),
+ candidate.accRegTileType);
+ Value accViewOut =
+ memref::SubViewOp::create(b, loc, candidate.accMemOut, accViewOffsets,
+ accViewSizes, accViewStrides);
+ Operation *accWrite = vector::TransferWriteOp::create(
+ b, loc, accSpliced, accViewOut, ValueRange{c0, c0});
+
+ setNativeShapeAttr(accWrite, candidate.accNativeType);
+}
+
+// Rebuild the accumulation loop with [lhs|rhs]RegTileType-shaped operations,
+// taking the offsets from the M-N-loop nests into account. The accumulator
+// iter_args are spliced into a accRegTileType-shaped value, and the contraction
+// result is split into accNativeType-shaped yield operands. Again, the
+// operations are annotated with an attribute for unrolling, after which the
+// splice/split logic will fold away completely.
+static SmallVector<Value> buildKLoopBody(MLUCandidate &candidate, OpBuilder &b,
+ Location loc, Value ivM, Value ivN,
+ Value ivK, ValueRange iterArgs) {
+ Value accSpliced = splice(b, loc, iterArgs, candidate.accRegTileType);
+
+ SmallVector<Value, 4> lhsOffsets = candidate.lhsRead.getIndices();
+ SmallVector<Value, 4> rhsOffsets = candidate.rhsRead.getIndices();
+ Value accLoopIV = candidate.accLoop.getInductionVar();
+ replace(lhsOffsets, accLoopIV, ivK);
+ replace(rhsOffsets, accLoopIV, ivK);
+ unsigned mPos = lhsOffsets.size() - 2 - candidate.isInVnniLayout;
+ unsigned nPos = rhsOffsets.size() - 1 - candidate.isInVnniLayout;
+ lhsOffsets[mPos] = arith::AddIOp::create(b, loc, lhsOffsets[mPos], ivM);
+ rhsOffsets[nPos] = arith::AddIOp::create(b, loc, rhsOffsets[nPos], ivN);
+ Value lhsView = makeRankReducingSubview(b, loc, candidate.lhsRead.getBase(),
+ lhsOffsets, candidate.lhsRegTileType);
+ Value rhsView = makeRankReducingSubview(b, loc, candidate.rhsRead.getBase(),
+ rhsOffsets, candidate.rhsRegTileType);
+
+ SmallVector<Value, 3> zeroIndices(2 + candidate.isInVnniLayout,
+ arith::ConstantIndexOp::create(b, loc, 0));
+ auto lhsRead = vector::TransferReadOp::create(
+ b, candidate.lhsRead.getLoc(), candidate.lhsRegTileType, lhsView,
+ zeroIndices, /*padding=*/std::nullopt);
+ auto rhsRead = vector::TransferReadOp::create(
+ b, candidate.rhsRead.getLoc(), candidate.rhsRegTileType, rhsView,
+ zeroIndices, /*padding=*/std::nullopt);
+ auto contract = vector::ContractionOp::create(
+ b, candidate.contract.getLoc(), lhsRead, rhsRead, accSpliced,
+ candidate.contract.getIndexingMaps(),
+ candidate.contract.getIteratorTypes());
+
+ setNativeShapeAttr(lhsRead, candidate.lhsNativeType);
+ setNativeShapeAttr(rhsRead, candidate.rhsNativeType);
+ contract->setAttr(nativeShapeAttrName,
+ b.getDenseI64ArrayAttr(candidate.contractNativeShape));
+
+ SmallVector<Value> accSliced =
+ slice(b, loc, contract, candidate.accNativeType);
+ return accSliced;
+}
+
+//----------------------------------------------------------------------------//
+// Pattern implementation //
+//----------------------------------------------------------------------------//
+
+namespace {
+// Prepares a contraction inside an accumulation loop for the nanokernel
+// lowerings.
+//
+// The pattern determines a register tile size for the contraction that can be
+// fully unrolled to a target-specific native shape without exceeding the
+// target's vector registers or AMX tile registers. It then introduces a
+// M-N-loop nest to iterate over such register tiles, and applies the vector
+// dialect's unroll patterns to produce native-shaped operations that the
+// nanokernel patterns can match.
+//
+// Currently, the match is deliberately strict (see the check* helpers above):
+// The contraction must be a canonical matrix multiplication, have
+// target-compatible types and shapes that can be tiled and unrolled cleanly,
+// and be embedded in an accumulation loop containing only the operand reads and
+// itself. This lets the rewrite regenerate the loop from scratch instead of
+// transforming it in place.
+struct VectorContractMultiLevelUnroll
+ : public OpRewritePattern<vector::ContractionOp> {
+ VectorContractMultiLevelUnroll(MLIRContext *context, StringRef target,
+ PatternBenefit benefit = 1)
+ : OpRewritePattern<vector::ContractionOp>(context, benefit),
+ target(target.str()) {}
+
+ LogicalResult matchAndRewrite(vector::ContractionOp contract,
+ PatternRewriter &rewriter) const override {
+ // Bail out early if this op is the result of a previous application of the
+ // pattern.
+ if (contract->hasAttr(nativeShapeAttrName))
+ return failure();
+
+ MLUCandidate candidate;
+ candidate.target = target;
+ candidate.contract = contract;
+ if (failed(checkCanonicalMatmul(candidate, rewriter)) ||
+ failed(checkShapesAndTypes(candidate, rewriter)) ||
+ failed(checkAccumulationLoop(candidate, rewriter)))
+ return failure();
+
+ makeAccInitializationBackedByMemory(candidate, rewriter);
+ makeAccResultBackedByMemory(candidate, rewriter);
+ buildLoopNestForCandidate(candidate, rewriter);
+ moveLoopInvariantCode(candidate.slicedAccLoop);
+ for (auto *op : candidate.opsToDelete)
+ rewriter.eraseOp(op);
+
+ return success();
+ }
+
+private:
+ std::string target;
+};
+} // namespace
+
+void x86::populateVectorContractMultiLevelUnrollPatterns(
+ RewritePatternSet &patterns, StringRef target) {
+ patterns.add<VectorContractMultiLevelUnroll>(patterns.getContext(), target);
+ vector::UnrollVectorOptions options;
+ options.setFilterConstraint(hasNativeShapeAttr);
+ options.setNativeShapeFn(getFromNativeShapeAttr);
+ vector::populateVectorUnrollPatterns(patterns, options);
+}
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
new file mode 100644
index 00000000000000..1c19890f91336c
--- /dev/null
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
@@ -0,0 +1,1239 @@
+// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter=entry-point=__transform_nano -canonicalize -cse -split-input-file \
+// RUN: | FileCheck %s --check-prefix=NANO --implicit-check-not=vector.contract
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @amx_bf16_flat(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<128x64xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<128x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<128x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x64xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<128x32xbf16>, vector<32x64xbf16> into vector<128x64xf32>
+ scf.yield %d : vector<128x64xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<128x64xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [128, 64] [1, 1]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC2:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC3:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]], %[[ARG2:.+]] = %[[ACC2]], %[[ARG3:.+]] = %[[ACC3]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]]] [32, 32] [1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]]] [32, 32] [1, 1]
+// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: %[[A1:.+]] = vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: %[[B1:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A0]], %[[B1]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>}
+// CHECK: %[[D2:.+]] = vector.contract {{.*}} %[[A1]], %[[B0]], %[[ARG2]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>}
+// CHECK: %[[D3:.+]] = vector.contract {{.*}} %[[A1]], %[[B1]], %[[ARG3]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>}
+// CHECK: scf.yield %[[D0]], %[[D1]], %[[D2]], %[[D3]]
+// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
+// CHECK: vector.transfer_write %[[RES]]#2, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#3, %[[ACC_VIEW]][%[[C16]], %[[C16]]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_flat(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// Zero-initialized accumulator: no reads from the destination, and the K-loop
+// starts from a native-shaped zero constant.
+func.func @amx_bf16_flat_zero_init(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c = arith.constant dense<0.0> : vector<96x128xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<96x128xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<96x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x128xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<96x32xbf16>, vector<32x128xbf16> into vector<96x128xf32>
+ scf.yield %d : vector<96x128xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<96x128xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_zero_init(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG: %[[ZERO:.+]] = arith.constant dense<0.000000e+00> : vector<16x16xf32>
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [96, 128] [1, 1]
+// CHECK-NOT: vector.transfer_read
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK-NOT: vector.transfer_read
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
+// CHECK-SAME: iter_args(%{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]])
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_flat_zero_init(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// The initial accumulator is not backed by memory, so it is spilled to a stack
+// buffer. The result is still written through a subview of the destination.
+func.func @amx_bf16_flat_init_from_arg(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %c: vector<64x96xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_init_from_arg(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[ACC_INIT:.+]]: vector<64x96xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_INIT]], %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK: %[[ACC_OUT:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
+// CHECK-NOT: memref.alloca
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for
+// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[ACC_OUT]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_flat_init_from_arg(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// The loop result feeds an epilogue instead of a plain store, so it is
+// buffered on the stack and rematerialized. The initial accumulator is still
+// read through a subview.
+func.func @amx_bf16_flat_truncf_result(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>, %D: memref<?x?xbf16>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ %trunc = arith.truncf %res : vector<64x96xf32> to vector<64x96xbf16>
+ vector.transfer_write %trunc, %D[%m, %n] {in_bounds = [true, true]} : vector<64x96xbf16>, memref<?x?xbf16>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_truncf_result(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>, %[[D:.+]]: memref<?x?xbf16>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-DAG: %[[ACC_IN:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
+// CHECK-DAG: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[ACC_IN]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for
+// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
+// CHECK: %[[TRUNC:.+]] = arith.truncf %[[ACC_RES]]
+// CHECK: vector.transfer_write %[[TRUNC]], %[[D]][%[[M]], %[[N]]]
+
+// NANO-LABEL: func.func @amx_bf16_flat_truncf_result(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// Neither the initial read nor the final write are fully in-bounds, so a
+// single stack buffer is used for both; the original transfers remain.
+func.func @amx_bf16_flat_acc_out_of_bounds(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [false, false]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [false, false]} : vector<64x96xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_acc_out_of_bounds(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-NOT: memref.subview %[[C]]
+// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]], {{.*}} : memref<?x?xf32>, vector<64x96xf32>
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_INIT]], %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.subview %[[C]]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for
+// CHECK-NOT: memref.subview %[[BUF]]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]] : vector<64x96xf32>, memref<?x?xf32>
+
+// NANO-LABEL: func.func @amx_bf16_flat_acc_out_of_bounds(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @amx_int8_flat(%A: memref<?x?xi8>, %B: memref<?x?xi8>, %C: memref<?x?xi32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c64 = arith.constant 64 : index
+ %c0_i8 = arith.constant 0 : i8
+ %c0_i32 = arith.constant 0 : i32
+ %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x128xi32>
+ %res = scf.for %k = %k_start to %k_end step %c64 iter_args(%acc = %c) -> (vector<64x128xi32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?xi8>, vector<64x64xi8>
+ %b = vector.transfer_read %B[%k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?xi8>, vector<64x128xi8>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
+ scf.yield %d : vector<64x128xi32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?xi32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_int8_flat(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xi8>, %[[B:.+]]: memref<?x?xi8>, %[[C:.+]]: memref<?x?xi32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 128] [1, 1]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C64]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]]] [32, 64] [1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]]] [64, 32] [1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_int8_flat(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+
+// M = 48 is an odd multiple of 16, so a 1x2 tile register layout is used.
+func.func @amx_bf16_vnni(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<48x64xf32>
+ %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<48x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<48x16x2xbf16>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<16x64x2xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<48x16x2xbf16>, vector<16x64x2xbf16> into vector<48x64xf32>
+ scf.yield %d : vector<48x64xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x64xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_vnni(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x2xbf16>, %[[B:.+]]: memref<?x?x2xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C48:.+]] = arith.constant 48 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [48, 64] [1, 1]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C48]] step %[[C16]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [16, 32] [1, 1]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]], 0] [16, 16, 2] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]], 0] [16, 32, 2] [1, 1, 1]
+// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: %[[B1:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A0]], %[[B1]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>}
+// CHECK: scf.yield %[[D0]], %[[D1]]
+// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_vnni(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+
+// N = 48 is an odd multiple of 16, so a 2x1 tile register layout is used.
+func.func @amx_int8_vnni(%A: memref<?x?x4xi8>, %B: memref<?x?x4xi8>, %C: memref<?x?xi32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c0_i8 = arith.constant 0 : i8
+ %c0_i32 = arith.constant 0 : i32
+ %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x48xi32>
+ %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<64x48xi32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0], %c0_i8 {in_bounds = [true, true, true]} : memref<?x?x4xi8>, vector<64x16x4xi8>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_i8 {in_bounds = [true, true, true]} : memref<?x?x4xi8>, vector<16x48x4xi8>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x16x4xi8>, vector<16x48x4xi8> into vector<64x48xi32>
+ scf.yield %d : vector<64x48xi32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x48xi32>, memref<?x?xi32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_int8_vnni(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x4xi8>, %[[B:.+]]: memref<?x?x4xi8>, %[[C:.+]]: memref<?x?xi32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C48:.+]] = arith.constant 48 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 48] [1, 1]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C48]] step %[[C16]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xi32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xi32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]], 0] [32, 16, 4] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]], 0] [16, 16, 4] [1, 1, 1]
+// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
+// CHECK: %[[A1:.+]] = vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
+// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
+// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 4>} : vector<16x16x4xi8>, vector<16x16x4xi8> into vector<16x16xi32>
+// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A1]], %[[B0]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 4>}
+// CHECK: scf.yield %[[D0]], %[[D1]]
+// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_int8_vnni(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// Blocked layout: the two leading memref dims select a block, and the block
+// has the same shape as the operand vectors.
+func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x32xbf16>, %B: memref<?x?x32x64xbf16>, %C: memref<?x?x96x64xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n, %c0, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?x96x64xf32>, vector<96x64xf32>
+ %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c) -> (vector<96x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x96x32xbf16>, vector<96x32xbf16>
+ %b = vector.transfer_read %B[%n, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x32x64xbf16>, vector<32x64xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<96x32xbf16>, vector<32x64xbf16> into vector<96x64xf32>
+ scf.yield %d : vector<96x64xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<96x64xf32>, memref<?x?x96x64xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_blocked(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x96x32xbf16>, %[[B:.+]]: memref<?x?x32x64xbf16>, %[[C:.+]]: memref<?x?x96x64xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]], 0, 0] [1, 1, 96, 64] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x96x64xf32> to memref<96x64xf32, strided<[64, 1], offset: ?>>
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C1]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], 0] [1, 1, 32, 32] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x96x32xbf16> to memref<32x32xbf16, strided<[32, 1], offset: ?>>
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], 0, %[[IV_N]]] [1, 1, 32, 32] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x32x64xbf16> to memref<32x32xbf16, strided<[64, 1], offset: ?>>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_flat_blocked(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+
+func.func @amx_bf16_vnni_blocked(%A: memref<?x?x128x16x2xbf16>, %B: memref<?x?x16x96x2xbf16>, %C: memref<?x?x128x96xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n, %c0, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?x128x96xf32>, vector<128x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c) -> (vector<128x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0, %c0, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x128x16x2xbf16>, vector<128x16x2xbf16>
+ %b = vector.transfer_read %B[%n, %k, %c0, %c0, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x16x96x2xbf16>, vector<16x96x2xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<128x16x2xbf16>, vector<16x96x2xbf16> into vector<128x96xf32>
+ scf.yield %d : vector<128x96xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<128x96xf32>, memref<?x?x128x96xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_vnni_blocked(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x128x16x2xbf16>, %[[B:.+]]: memref<?x?x16x96x2xbf16>, %[[C:.+]]: memref<?x?x128x96xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]], 0, 0] [1, 1, 128, 96] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x128x96xf32> to memref<128x96xf32, strided<[96, 1], offset: ?>>
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C1]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], 0, 0] [1, 1, 32, 16, 2] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x128x16x2xbf16> to memref<32x16x2xbf16, strided<[32, 2, 1], offset: ?>>
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], 0, %[[IV_N]], 0] [1, 1, 16, 32, 2] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x16x96x2xbf16> to memref<16x32x2xbf16, strided<[192, 2, 1], offset: ?>>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_vnni_blocked(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// The original shape equals the register tile shape, so the M-N loop nest has
+// a single iteration and folds away.
+func.func @amx_bf16_flat_single_reg_tile(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<32x32xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<32x32xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x32xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %d : vector<32x32xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<32x32xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_single_reg_tile(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [32, 32] [1, 1]
+// CHECK-NOT: memref.subview %[[ACC_VIEW]]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC2:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC3:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]], %[[ARG2:.+]] = %[[ACC2]], %[[ARG3:.+]] = %[[ACC3]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]]] [32, 32] [1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[N]]] [32, 32] [1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
+// CHECK: vector.transfer_write %[[RES]]#2, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#3, %[[ACC_VIEW]][%[[C16]], %[[C16]]]
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_flat_single_reg_tile(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+
+func.func @amx_bf16_vnni_single_reg_tile(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<32x32xf32>
+ %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<32x32xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<32x16x2xbf16>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<16x32x2xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+ scf.yield %d : vector<32x32xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<32x32xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_vnni_single_reg_tile(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x2xbf16>, %[[B:.+]]: memref<?x?x2xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [32, 32] [1, 1]
+// CHECK-NOT: memref.subview %[[ACC_VIEW]]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC2:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC3:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]], %[[ARG2:.+]] = %[[ACC2]], %[[ARG3:.+]] = %[[ACC3]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], 0] [32, 16, 2] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[N]], 0] [16, 32, 2] [1, 1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
+// CHECK: vector.transfer_write %[[RES]]#2, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#3, %[[ACC_VIEW]][%[[C16]], %[[C16]]]
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_vnni_single_reg_tile(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// RHS is transposed (N x K instead of K x N).
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @negative_transposed_operand(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%n, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<96x32xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<96x32xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @negative_transposed_operand(
+// CHECK-NOT: memref.alloca
+// CHECK: scf.for
+// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<96x32xbf16> into vector<64x96xf32>
+// CHECK-NOT: memref.alloca
+
+// NANO-LABEL: func.func @negative_transposed_operand(
+// NANO: vector.contract
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @negative_operand_read_out_of_bounds(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, false]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @negative_operand_read_out_of_bounds(
+// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+
+// NANO-LABEL: func.func @negative_operand_read_out_of_bounds(
+// NANO: vector.contract
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @negative_unsupported_type(%A: memref<?x?xi16>, %B: memref<?x?xi16>, %C: memref<?x?xi32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_i16 = arith.constant 0 : i16
+ %c0_i32 = arith.constant 0 : i32
+ %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x96xi32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xi32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_i16 {in_bounds = [true, true]} : memref<?x?xi16>, vector<64x32xi16>
+ %b = vector.transfer_read %B[%k, %n], %c0_i16 {in_bounds = [true, true]} : memref<?x?xi16>, vector<32x96xi16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xi16>, vector<32x96xi16> into vector<64x96xi32>
+ scf.yield %d : vector<64x96xi32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xi32>, memref<?x?xi32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @negative_unsupported_type(
+// CHECK: vector.contract {{.*}} : vector<64x32xi16>, vector<32x96xi16> into vector<64x96xi32>
+
+// NANO-LABEL: func.func @negative_unsupported_type(
+// NANO: vector.contract
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Flat layout requires a 2x2 tile register layout, i.e. M and N must be
+// multiples of 32.
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @negative_shape_not_unrollable(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<48x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<48x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<48x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<48x32xbf16>, vector<32x96xbf16> into vector<48x96xf32>
+ scf.yield %d : vector<48x96xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x96xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @negative_shape_not_unrollable(
+// CHECK: vector.contract {{.*}} : vector<48x32xbf16>, vector<32x96xbf16> into vector<48x96xf32>
+
+// NANO-LABEL: func.func @negative_shape_not_unrollable(
+// NANO: vector.contract
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// VNNI contraction already has the native tile shape; nothing to do.
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+
+func.func @negative_vnni_native_shape(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<16x16xf32>
+ %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<16x16xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<16x16x2xbf16>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<16x16x2xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+ scf.yield %d : vector<16x16xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<16x16xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @negative_vnni_native_shape(
+// CHECK: vector.contract {{.*}} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+// CHECK-NOT: x86_vcmlu_native_shape
+
+// NANO-LABEL: func.func @negative_vnni_native_shape(
+// Operands are read directly from a function argument; the AMX lowering
+// requires an op (e.g. memref.subview) to define the read source.
+// NANO: vector.contract
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
>From b0bb4be49b873e1915397baf4ea42b9e3fd23fd2 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Mon, 5 Oct 2026 07:48:14 -0700
Subject: [PATCH 02/15] PR feedback
---
.../X86/Transforms/VectorContractMultiLevelUnroll.cpp | 9 +++++----
1 file changed, 5 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 599363724e7def..a31d129092b750 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -236,7 +236,7 @@ static void setNativeShapeAttr(Operation *op, VectorType nativeType) {
// } %lhs, %rhs, %acc : !lhsType, !rhsType into !accType
static LogicalResult checkCanonicalMatmul(MLUCandidate &candidate,
PatternRewriter &rewriter) {
- auto &contract = candidate.contract;
+ vector::ContractionOp contract = candidate.contract;
if (contract.getKind() != vector::CombiningKind::ADD)
return rewriter.notifyMatchFailure(contract,
"not using ADD combining kind");
@@ -298,8 +298,8 @@ static LogicalResult checkCanonicalMatmul(MLUCandidate &candidate,
// to determine a suitable tiling strategy for it.
static LogicalResult checkShapesAndTypes(MLUCandidate &candidate,
PatternRewriter &rewriter) {
- auto &contract = candidate.contract;
- auto target = candidate.target;
+ vector::ContractionOp contract = candidate.contract;
+ StringRef target = candidate.target;
Type accElemType = candidate.accType.getElementType();
Type inpElemType = candidate.lhsType.getElementType();
@@ -353,8 +353,9 @@ static LogicalResult checkShapesAndTypes(MLUCandidate &candidate,
if (origK != nativeK)
return rewriter.notifyMatchFailure(contract, "K dimension mismatch");
}
- } else
+ } else {
return rewriter.notifyMatchFailure(contract, "unsupported target");
+ }
candidate.accRegTileType = VectorType::get({regTileM, regTileN}, accElemType);
candidate.accNativeType = VectorType::get({nativeM, nativeN}, accElemType);
>From 67a50af765e75492cfe11515f56e65b2656ba49d Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Mon, 5 Oct 2026 09:16:40 -0700
Subject: [PATCH 03/15] Document target field
---
.../VectorContractMultiLevelUnroll.cpp | 10 ++-
.../vector-contract-multi-level-unroll.mlir | 68 +++++++++----------
2 files changed, 42 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index a31d129092b750..232084f7429371 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -31,7 +31,13 @@ static constexpr auto nativeShapeAttrName = "x86_vcmlu_native_shape";
namespace {
// Shared state for the pattern match and IR transformation.
struct MLUCandidate {
+ // The CPU feature flag to target with the transformation, in the spelling
+ // used by llvm/include/llvm/TargetParser/X86TargetParser.def, e.g.
+ // "amx-bf16". Determines native shapes and supported data types for operands
+ // and the accumulator.
StringRef target;
+
+ // Anchor operation.
vector::ContractionOp contract;
bool isInVnniLayout = false;
@@ -316,9 +322,9 @@ static LogicalResult checkShapesAndTypes(MLUCandidate &candidate,
int64_t nativeM, nativeN, nativeK, vnni;
if (target.starts_with("amx")) {
- if (!((target == "amx_bf16" && inpElemType.isBF16() &&
+ if (!((target == "amx-bf16" && inpElemType.isBF16() &&
accElemType.isF32()) ||
- (target == "amx_int8" && inpElemType.isInteger(8) &&
+ (target == "amx-int8" && inpElemType.isInteger(8) &&
accElemType.isInteger(32))))
return rewriter.notifyMatchFailure(
contract, "unsupported combination of input and accumulator types");
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
index 1c19890f91336c..511bebc481c465 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
@@ -76,7 +76,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -84,7 +84,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -151,7 +151,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -159,7 +159,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -224,7 +224,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -232,7 +232,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -298,7 +298,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -306,7 +306,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -372,7 +372,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -380,7 +380,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -453,7 +453,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
} : !transform.any_op
transform.yield
}
@@ -461,7 +461,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -539,7 +539,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -547,7 +547,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -625,7 +625,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
} : !transform.any_op
transform.yield
}
@@ -633,7 +633,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -711,7 +711,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -719,7 +719,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -795,7 +795,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -803,7 +803,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -877,7 +877,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -885,7 +885,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -957,7 +957,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -965,7 +965,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -1012,7 +1012,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -1020,7 +1020,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -1063,7 +1063,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -1071,7 +1071,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -1114,7 +1114,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
} : !transform.any_op
transform.yield
}
@@ -1122,7 +1122,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -1167,7 +1167,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -1175,7 +1175,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -1223,7 +1223,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -1231,7 +1231,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx_bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
>From 3975f2afa935b2aedcb8926fa1cc604c33190e61 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Mon, 5 Oct 2026 09:18:59 -0700
Subject: [PATCH 04/15] Typo
---
.../Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 232084f7429371..aff05e1f1be8b1 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -396,7 +396,7 @@ static LogicalResult checkShapesAndTypes(MLUCandidate &candidate,
//
// For the operand reads, we require that they have a minor identity permutation
// map, impose no bounds checks and don't carry a mask. In addition, they shall
-// use the induction variable as exactly on of their indices. This covers plain
+// use the induction variable as exactly one of their indices. This covers plain
// 2D memrefs as well as block-based layouts.
//
// Example:
>From ed356d3dfc94d86705936d400222d04ca7f142c4 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Mon, 5 Oct 2026 10:53:42 -0700
Subject: [PATCH 05/15] Rename check...-methods
---
.../VectorContractMultiLevelUnroll.cpp | 24 ++++++++++++-------
1 file changed, 15 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index aff05e1f1be8b1..63b576a8cedda4 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -221,7 +221,8 @@ static void setNativeShapeAttr(Operation *op, VectorType nativeType) {
//----------------------------------------------------------------------------//
// Establish whether the candidate contraction is a canonical matrix
-// multiplication on flat or VNNI-packed vectors.
+// multiplication on flat or VNNI-packed vectors. If successful, the discovered
+// operand and accumulator types are recorded in the candidate struct.
//
// Example (flat):
// vector.contract {
@@ -240,7 +241,7 @@ static void setNativeShapeAttr(Operation *op, VectorType nativeType) {
// affine_map<(dm, dn, dk, dvnni) -> (dm, dn)>],
// iterator_types = ["parallel", "parallel", "reduction", "reduction"],
// } %lhs, %rhs, %acc : !lhsType, !rhsType into !accType
-static LogicalResult checkCanonicalMatmul(MLUCandidate &candidate,
+static LogicalResult matchCanonicalMatmul(MLUCandidate &candidate,
PatternRewriter &rewriter) {
vector::ContractionOp contract = candidate.contract;
if (contract.getKind() != vector::CombiningKind::ADD)
@@ -301,8 +302,10 @@ static LogicalResult checkCanonicalMatmul(MLUCandidate &candidate,
}
// Check compatibility of datatypes with the given target ISA extension, and try
-// to determine a suitable tiling strategy for it.
-static LogicalResult checkShapesAndTypes(MLUCandidate &candidate,
+// to determine a suitable tiling strategy for it. If successful, the strategy
+// is encoded in the *RegTileType and *NativeType fields of the candidate
+// struct.
+static LogicalResult matchShapesAndTypes(MLUCandidate &candidate,
PatternRewriter &rewriter) {
vector::ContractionOp contract = candidate.contract;
StringRef target = candidate.target;
@@ -399,6 +402,9 @@ static LogicalResult checkShapesAndTypes(MLUCandidate &candidate,
// use the induction variable as exactly one of their indices. This covers plain
// 2D memrefs as well as block-based layouts.
//
+// If all these conditions are met, the loop op, initial value of the
+// accumulator and the operand reads are recorded in the candidate struct.
+//
// Example:
// !accType = vector<64x128xf32>
// !lhsType = vector<64x32xbf16>
@@ -414,7 +420,7 @@ static LogicalResult checkShapesAndTypes(MLUCandidate &candidate,
// : !lhsType, !rhsType into !accType
// scf.yield %d : !accType
// }
-static LogicalResult checkAccumulationLoop(MLUCandidate &candidate,
+static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
PatternRewriter &rewriter) {
auto &contract = candidate.contract;
@@ -750,7 +756,7 @@ namespace {
// dialect's unroll patterns to produce native-shaped operations that the
// nanokernel patterns can match.
//
-// Currently, the match is deliberately strict (see the check* helpers above):
+// Currently, the match is deliberately strict (see the helpers above):
// The contraction must be a canonical matrix multiplication, have
// target-compatible types and shapes that can be tiled and unrolled cleanly,
// and be embedded in an accumulation loop containing only the operand reads and
@@ -773,9 +779,9 @@ struct VectorContractMultiLevelUnroll
MLUCandidate candidate;
candidate.target = target;
candidate.contract = contract;
- if (failed(checkCanonicalMatmul(candidate, rewriter)) ||
- failed(checkShapesAndTypes(candidate, rewriter)) ||
- failed(checkAccumulationLoop(candidate, rewriter)))
+ if (failed(matchCanonicalMatmul(candidate, rewriter)) ||
+ failed(matchShapesAndTypes(candidate, rewriter)) ||
+ failed(matchAccumulationLoop(candidate, rewriter)))
return failure();
makeAccInitializationBackedByMemory(candidate, rewriter);
>From 93aa3481ddad0328062bb2b6595611f47f41cb6f Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 7 Oct 2026 00:39:27 -0700
Subject: [PATCH 06/15] Don't specify number of inline elements
---
.../Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 63b576a8cedda4..03f6de80e23f60 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -57,7 +57,7 @@ struct MLUCandidate {
VectorType accNativeType;
VectorType lhsNativeType;
VectorType rhsNativeType;
- SmallVector<int64_t, 4> contractNativeShape;
+ SmallVector<int64_t> contractNativeShape;
// Bookkeeping for the accumulation loop and operands.
scf::ForOp accLoop;
@@ -76,7 +76,7 @@ struct MLUCandidate {
// Keep track of operations that should be deleted after transformation
// (depending on which ops could be integrated into the new loop nest).
- SmallVector<Operation *, 2> opsToDelete;
+ SmallVector<Operation *> opsToDelete;
};
} // anonymous namespace
>From 32d3eb4a01d56c7720812e963441f6ee5f7a7d20 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 7 Oct 2026 00:44:35 -0700
Subject: [PATCH 07/15] Use VectorTransferOpInterface
---
.../VectorContractMultiLevelUnroll.cpp | 24 ++++++++-----------
1 file changed, 10 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 03f6de80e23f60..96daf55d0bed28 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -86,8 +86,7 @@ struct MLUCandidate {
// Return true if given vector.transfer_read/write has a default permutation
// map, no mask and no out-of-bounds check.
-template <typename TransferOp>
-static bool isSimpleTransferOp(TransferOp op) {
+static bool isSimpleTransferOp(VectorTransferOpInterface op) {
return op.getPermutationMap().isMinorIdentity() && !op.getMask() &&
!op.hasOutOfBoundsDim();
}
@@ -204,13 +203,9 @@ getFromNativeShapeAttr(Operation *op) {
return SmallVector<int64_t>(attr.asArrayRef());
}
-static void setNativeShapeAttr(Operation *op, VectorType nativeType) {
- VectorType currentType;
- TypeSwitch<Operation *>(op)
- .Case<vector::TransferReadOp, vector::TransferWriteOp>(
- [&](auto transferOp) { currentType = transferOp.getVectorType(); })
- .DefaultUnreachable("expected a vector transfer op");
- if (currentType != nativeType)
+static void setNativeShapeAttr(VectorTransferOpInterface op,
+ VectorType nativeType) {
+ if (op.getVectorType() != nativeType)
op->setAttr(
nativeShapeAttrName,
DenseI64ArrayAttr::get(nativeType.getContext(), nativeType.getShape()));
@@ -664,10 +659,11 @@ static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
Value accViewIn =
memref::SubViewOp::create(b, loc, candidate.accMemIn, accViewOffsets,
accViewSizes, accViewStrides);
- accInit = vector::TransferReadOp::create(b, loc, candidate.accRegTileType,
- accViewIn, ValueRange{c0, c0},
- /*padding=*/std::nullopt);
- setNativeShapeAttr(accInit, candidate.accNativeType);
+ auto accInitRead = vector::TransferReadOp::create(
+ b, loc, candidate.accRegTileType, accViewIn, ValueRange{c0, c0},
+ /*padding=*/std::nullopt);
+ setNativeShapeAttr(accInitRead, candidate.accNativeType);
+ accInit = accInitRead;
}
SmallVector<Value> accInitSliced =
@@ -687,7 +683,7 @@ static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
Value accViewOut =
memref::SubViewOp::create(b, loc, candidate.accMemOut, accViewOffsets,
accViewSizes, accViewStrides);
- Operation *accWrite = vector::TransferWriteOp::create(
+ auto accWrite = vector::TransferWriteOp::create(
b, loc, accSpliced, accViewOut, ValueRange{c0, c0});
setNativeShapeAttr(accWrite, candidate.accNativeType);
>From e574792c8b09bace24812c77008b54231fd4a298 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 7 Oct 2026 02:03:29 -0700
Subject: [PATCH 08/15] Add test case with existing loop nest
---
.../vector-contract-multi-level-unroll.mlir | 107 ++++++++++++++++++
1 file changed, 107 insertions(+)
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
index 511bebc481c465..789bf11472a6e4 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
@@ -470,6 +470,113 @@ module attributes {transform.with_named_sequence} {
// -----
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// The accumulation loop sits inside an existing batch-M-N loop nest.
+func.func @amx_int8_flat_batched_loop_nest(%A: memref<?x?x?xi8>, %B: memref<?x?x?xi8>, %C: memref<?x?x?xi32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c0_i8 = arith.constant 0 : i8
+ %c0_i32 = arith.constant 0 : i32
+ %batch = memref.dim %C, %c0 : memref<?x?x?xi32>
+ %M = memref.dim %C, %c1 : memref<?x?x?xi32>
+ %N = memref.dim %C, %c2 : memref<?x?x?xi32>
+ %K = memref.dim %A, %c2 : memref<?x?x?xi8>
+ scf.for %bi = %c0 to %batch step %c1 {
+ scf.for %m = %c0 to %M step %c64 {
+ scf.for %n = %c0 to %N step %c128 {
+ %c = vector.transfer_read %C[%bi, %m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?x?xi32>, vector<64x128xi32>
+ %res = scf.for %k = %c0 to %K step %c64 iter_args(%acc = %c) -> (vector<64x128xi32>) {
+ %a = vector.transfer_read %A[%bi, %m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x64xi8>
+ %b = vector.transfer_read %B[%bi, %k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x128xi8>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
+ scf.yield %d : vector<64x128xi32>
+ }
+ vector.transfer_write %res, %C[%bi, %m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?x?xi32>
+ }
+ }
+ }
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_int8_flat_batched_loop_nest(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x?xi8>, %[[B:.+]]: memref<?x?x?xi8>, %[[C:.+]]: memref<?x?x?xi32>)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[C]], %[[C0]]
+// CHECK-DAG: %[[M:.+]] = memref.dim %[[C]], %[[C1]]
+// CHECK-DAG: %[[N:.+]] = memref.dim %[[C]], %[[C2]]
+// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
+// CHECK-NOT: memref.alloca
+// CHECK: scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]] {
+// CHECK: scf.for %[[OFF_M:.+]] = %[[C0]] to %[[M]] step %[[C64]] {
+// CHECK: scf.for %[[OFF_N:.+]] = %[[C0]] to %[[N]] step %[[C128]] {
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[IV_B]], %[[OFF_M]], %[[OFF_N]]] [1, 64, 128] [1, 1, 1]
+// CHECK-SAME: memref<?x?x?xi32> to memref<64x128xi32, strided<[?, 1], offset: ?>>
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
+// CHECK: %[[A_OFF_M:.+]] = arith.addi %[[OFF_M]], %[[IV_M]]
+// CHECK: %[[B_OFF_N:.+]] = arith.addi %[[OFF_N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C64]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[A_OFF_M]], %[[IV_K]]] [1, 32, 64] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[B_OFF_N]]] [1, 64, 32] [1, 1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: }
+// CHECK: }
+// CHECK: }
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_int8_flat_batched_loop_nest(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
>From 3c4571ba3115572e565a7d05ef9b39f2d36c53b7 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 7 Oct 2026 03:18:00 -0700
Subject: [PATCH 09/15] Add test for (bad) batch loop lowering
---
.../vector-contract-multi-level-unroll.mlir | 106 ++++++++++++++++++
1 file changed, 106 insertions(+)
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
index 789bf11472a6e4..7d17f81d8fd643 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
@@ -577,6 +577,112 @@ module attributes {transform.with_named_sequence} {
// -----
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// The accumulator is loop-carried through an enclosing batch loop, so it is
+// neither read from nor written to memory directly.
+//
+// This test is currently intended to demonstrate successful lowering using
+// the fallback stack-buffer path. A more desirable lowering would keep the
+// accumulator on the tile registers across both loops.
+func.func @amx_int8_flat_batch_reduction(%A: memref<?x?x?xi8>, %B: memref<?x?x?xi8>, %C: memref<?x?xi32>,
+ %m: index, %n: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c64 = arith.constant 64 : index
+ %c0_i8 = arith.constant 0 : i8
+ %c0_i32 = arith.constant 0 : i32
+ %batch = memref.dim %A, %c0 : memref<?x?x?xi8>
+ %K = memref.dim %A, %c2 : memref<?x?x?xi8>
+ %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x128xi32>
+ %res = scf.for %b = %c0 to %batch step %c1 iter_args(%acc_b = %c) -> (vector<64x128xi32>) {
+ %res_b = scf.for %k = %c0 to %K step %c64 iter_args(%acc = %acc_b) -> (vector<64x128xi32>) {
+ %a = vector.transfer_read %A[%b, %m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x64xi8>
+ %bv = vector.transfer_read %B[%b, %k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x128xi8>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %bv, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
+ scf.yield %d : vector<64x128xi32>
+ }
+ scf.yield %res_b : vector<64x128xi32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?xi32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_int8_flat_batch_reduction(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x?xi8>, %[[B:.+]]: memref<?x?x?xi8>, %[[C:.+]]: memref<?x?xi32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[A]], %[[C0]]
+// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.subview
+// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]]
+// CHECK: %[[ACC_RES:.+]] = scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]] iter_args(%[[ACC_B:.+]] = %[[ACC_INIT]])
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x128xi32>
+// CHECK: vector.transfer_write %[[ACC_B]], %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK-NOT: memref.alloca
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C64]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[OFF_M]], %[[IV_K]]] [1, 32, 64] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[OFF_N]]] [1, 64, 32] [1, 1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-NOT: memref.subview %[[BUF]]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK: %[[ACC_NEXT:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x128xi32>
+// CHECK: scf.yield %[[ACC_NEXT]]
+// CHECK: }
+// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]]
+
+// NANO-LABEL: func.func @amx_int8_flat_batch_reduction(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
>From a07bf0c794b71d1320405bfceff2096eb491811a Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 7 Oct 2026 05:21:25 -0700
Subject: [PATCH 10/15] Playing with loop order
---
.../vector-contract-multi-level-unroll.mlir | 196 ++++++++++++++++--
1 file changed, 182 insertions(+), 14 deletions(-)
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
index 7d17f81d8fd643..7222af1e6241ac 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
@@ -683,6 +683,113 @@ module attributes {transform.with_named_sequence} {
// -----
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// Loop order K-batch, accumulating over both loops: the inner batch loop
+// carries the accumulator, its IV appears exactly once in each operand read and
+// its step is 1. It is a valid accumulation loop, indistinguishable from a
+// K-loop over a blocked layout, and is treated as such.
+//
+// Again, spilling the accumulator value to memory in the outer loop is correct,
+// but not ideal.
+func.func @amx_bf16_flat_k_batch_order(%A: memref<?x?x?xbf16>, %B: memref<?x?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %batch = memref.dim %A, %c0 : memref<?x?x?xbf16>
+ %K = memref.dim %A, %c2 : memref<?x?x?xbf16>
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %c0 to %K step %c32 iter_args(%acc_k = %c) -> (vector<64x96xf32>) {
+ %res_k = scf.for %bi = %c0 to %batch step %c1 iter_args(%acc = %acc_k) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%bi, %m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%bi, %k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ scf.yield %res_k : vector<64x96xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_k_batch_order(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x?xbf16>, %[[B:.+]]: memref<?x?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[A]], %[[C0]]
+// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.subview
+// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]]
+// CHECK: %[[ACC_RES:.+]] = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C32]] iter_args(%[[ACC_K:.+]] = %[[ACC_INIT]])
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_K]], %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK-NOT: memref.alloca
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[OFF_M]], %[[IV_K]]] [1, 32, 32] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[OFF_N]]] [1, 32, 32] [1, 1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-NOT: memref.subview %[[BUF]]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK: %[[ACC_NEXT:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
+// CHECK: scf.yield %[[ACC_NEXT]]
+// CHECK: }
+// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]]
+
+// NANO-LABEL: func.func @amx_bf16_flat_k_batch_order(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
@@ -1213,13 +1320,74 @@ func.func @negative_transposed_operand(%A: memref<?x?xbf16>, %B: memref<?x?xbf16
}
// CHECK-LABEL: func.func @negative_transposed_operand(
-// CHECK-NOT: memref.alloca
-// CHECK: scf.for
-// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<96x32xbf16> into vector<64x96xf32>
-// CHECK-NOT: memref.alloca
+// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<96x32xbf16> into vector<64x96xf32>
+
+// NANO-LABEL: func.func @negative_transposed_operand(
+// NANO: vector.contract
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// Loop order M-K-N, accumulating over both K and N: the IV of the innermost
+// loop carrying the accumulator (N) is not present in both operand loads'
+// indices.
+func.func @negative_k_loop_not_innermost(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c32 = arith.constant 32 : index
+ %c64 = arith.constant 64 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %M = memref.dim %C, %c0 : memref<?x?xf32>
+ %N = memref.dim %C, %c1 : memref<?x?xf32>
+ %K = memref.dim %A, %c1 : memref<?x?xbf16>
+ scf.for %m = %c0 to %M step %c64 {
+ %c = vector.transfer_read %C[%m, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x32xf32>
+ %res = scf.for %k = %c0 to %K step %c32 iter_args(%acc_k = %c) -> (vector<64x32xf32>) {
+ %res_k = scf.for %n = %c0 to %N step %c32 iter_args(%acc = %acc_k) -> (vector<64x32xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x32xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x32xbf16> into vector<64x32xf32>
+ scf.yield %d : vector<64x32xf32>
+ }
+ scf.yield %res_k : vector<64x32xf32>
+ }
+ vector.transfer_write %res, %C[%m, %c0] {in_bounds = [true, true]} : vector<64x32xf32>, memref<?x?xf32>
+ }
+ func.return
+}
+
+// CHECK-LABEL: func.func @negative_k_loop_not_innermost(
+// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<32x32xbf16> into vector<64x32xf32>
-// NANO-LABEL: func.func @negative_transposed_operand(
-// NANO: vector.contract
+// NANO-LABEL: func.func @negative_k_loop_not_innermost(
+// NANO: vector.contract
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -1269,8 +1437,8 @@ func.func @negative_operand_read_out_of_bounds(%A: memref<?x?xbf16>, %B: memref<
// CHECK-LABEL: func.func @negative_operand_read_out_of_bounds(
// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
-// NANO-LABEL: func.func @negative_operand_read_out_of_bounds(
-// NANO: vector.contract
+// NANO-LABEL: func.func @negative_operand_read_out_of_bounds(
+// NANO: vector.contract
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -1320,8 +1488,8 @@ func.func @negative_unsupported_type(%A: memref<?x?xi16>, %B: memref<?x?xi16>, %
// CHECK-LABEL: func.func @negative_unsupported_type(
// CHECK: vector.contract {{.*}} : vector<64x32xi16>, vector<32x96xi16> into vector<64x96xi32>
-// NANO-LABEL: func.func @negative_unsupported_type(
-// NANO: vector.contract
+// NANO-LABEL: func.func @negative_unsupported_type(
+// NANO: vector.contract
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -1373,8 +1541,8 @@ func.func @negative_shape_not_unrollable(%A: memref<?x?xbf16>, %B: memref<?x?xbf
// CHECK-LABEL: func.func @negative_shape_not_unrollable(
// CHECK: vector.contract {{.*}} : vector<48x32xbf16>, vector<32x96xbf16> into vector<48x96xf32>
-// NANO-LABEL: func.func @negative_shape_not_unrollable(
-// NANO: vector.contract
+// NANO-LABEL: func.func @negative_shape_not_unrollable(
+// NANO: vector.contract
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -1427,10 +1595,10 @@ func.func @negative_vnni_native_shape(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xb
// CHECK: vector.contract {{.*}} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
// CHECK-NOT: x86_vcmlu_native_shape
-// NANO-LABEL: func.func @negative_vnni_native_shape(
+// NANO-LABEL: func.func @negative_vnni_native_shape(
// Operands are read directly from a function argument; the AMX lowering
// requires an op (e.g. memref.subview) to define the read source.
-// NANO: vector.contract
+// NANO: vector.contract
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
>From 67687cfb3708ac4d1dc04908a80a92d8ebe56671 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 7 Oct 2026 05:46:06 -0700
Subject: [PATCH 11/15] Don't try to trace the accumulator init over multiple
loop levels
---
.../VectorContractMultiLevelUnroll.cpp | 18 +++++++++++++++---
1 file changed, 15 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 96daf55d0bed28..d2ef799c454edd 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -84,6 +84,18 @@ struct MLUCandidate {
// Utilities //
//----------------------------------------------------------------------------//
+// Return true if the argument is an all-zero vector.
+static bool isZeroVector(Value v) {
+ auto constOp = v.getDefiningOp<arith::ConstantOp>();
+ if (!constOp)
+ return false;
+ auto denseAttr = dyn_cast<DenseElementsAttr>(constOp.getValue());
+ if (!denseAttr || !denseAttr.isSplat())
+ return false;
+
+ return isZeroIntegerOrFloat(denseAttr.getSplatValue<Attribute>());
+}
+
// Return true if given vector.transfer_read/write has a default permutation
// map, no mask and no out-of-bounds check.
static bool isSimpleTransferOp(VectorTransferOpInterface op) {
@@ -448,8 +460,8 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return rewriter.notifyMatchFailure(contract,
"accumulator is not loop-carried");
- bool accIsZeroInit = isa_and_present<arith::ConstantOp>(
- x86::traceToVectorReadLikeParentOperation(contract.getAcc()));
+ Value accInitVal = accLoop.getTiedLoopInit(accIterArg)->get();
+ bool accIsZeroInit = isZeroVector(accInitVal);
auto lhsRead = contract.getLhs().getDefiningOp<vector::TransferReadOp>();
auto rhsRead = contract.getRhs().getDefiningOp<vector::TransferReadOp>();
@@ -476,7 +488,7 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return failure();
candidate.accLoop = accLoop;
- candidate.accInitVal = accLoop.getTiedLoopInit(accIterArg)->get();
+ candidate.accInitVal = accInitVal;
candidate.accIsZeroInit = accIsZeroInit;
candidate.lhsRead = lhsRead;
candidate.rhsRead = rhsRead;
>From e86d7661ea60965fb28966e36070c1522fe48a5b Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Thu, 8 Oct 2026 06:57:58 -0700
Subject: [PATCH 12/15] Check whether it's safe to use a view for the result
write-back
---
.../VectorContractMultiLevelUnroll.cpp | 27 +++++--
.../vector-contract-multi-level-unroll.mlir | 77 +++++++++++++++++++
2 files changed, 99 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index d2ef799c454edd..5a3b1bad44ac29 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -559,6 +559,22 @@ static void makeAccResultBackedByMemory(MLUCandidate &candidate,
Block *accLoopBlock = candidate.accLoop->getBlock();
if (accWrite && isSimpleTransferOp(accWrite) &&
accWrite->getBlock() == accLoopBlock) {
+
+ Value accWriteBase = accWrite.getBase();
+ ValueRange accWriteIndices = accWrite.getIndices();
+
+ auto isSafeToUseView = [&]() {
+ if (candidate.accIsZeroInit)
+ return true;
+ auto accMemInView = candidate.accMemIn.getDefiningOp<memref::SubViewOp>();
+ if (!accMemInView)
+ return true;
+
+ return accMemInView.getSource() == accWriteBase &&
+ isEqualConstantIntOrValueArray(accMemInView.getMixedOffsets(),
+ getAsOpFoldResult(accWriteIndices));
+ };
+
auto isDefinedBeforeAccLoop = [&](Value v) {
Operation *defOp = v.getDefiningOp();
// No need for full dominance check, as we already know that accLoop and
@@ -566,11 +582,12 @@ static void makeAccResultBackedByMemory(MLUCandidate &candidate,
return !defOp || defOp->getBlock() != accLoopBlock ||
defOp->isBeforeInBlock(candidate.accLoop);
};
- if (isDefinedBeforeAccLoop(accWrite.getBase()) &&
- all_of(accWrite.getIndices(), isDefinedBeforeAccLoop)) {
- candidate.accMemOut = makeRankReducingSubview(
- rewriter, loc, accWrite.getBase(), accWrite.getIndices(),
- accWrite.getVectorType());
+
+ if (isSafeToUseView() && isDefinedBeforeAccLoop(accWriteBase) &&
+ all_of(accWriteIndices, isDefinedBeforeAccLoop)) {
+ candidate.accMemOut =
+ makeRankReducingSubview(rewriter, loc, accWriteBase, accWriteIndices,
+ accWrite.getVectorType());
candidate.opsToDelete.push_back(accWrite);
candidate.opsToDelete.push_back(candidate.accLoop);
return;
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
index 7222af1e6241ac..4097b4bed7a13b 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
@@ -393,6 +393,83 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+// The initial read and the final write target the same memref but at different
+// offsets, so the regions partially overlap. Writing the result tiles through
+// a view would clobber initial values not yet read by later M/N iterations, so
+// the initial value is still read via a subview, but the result goes through a
+// stack buffer and the original write remains.
+func.func @amx_bf16_flat_acc_partial_overlap(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %m_off = arith.addi %m, %c32 : index
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ vector.transfer_write %res, %C[%m_off, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_acc_partial_overlap(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-DAG: %[[M_OFF:.+]] = arith.addi %[[M]], %[[C32]]
+// CHECK-DAG: %[[ACC_IN:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
+// CHECK-DAG: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK-NOT: memref.subview %[[C]]
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[ACC_IN]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for
+// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M_OFF]], %[[N]]] {{.*}} : vector<64x96xf32>, memref<?x?xf32>
+
+// NANO-LABEL: func.func @amx_bf16_flat_acc_partial_overlap(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
func.func @amx_int8_flat(%A: memref<?x?xi8>, %B: memref<?x?xi8>, %C: memref<?x?xi32>,
%m: index, %n: index, %k_start: index, %k_end: index) {
%c64 = arith.constant 64 : index
>From 5b1fb7b1ad94dbb1f51474be1b9c1e75d322412c Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Thu, 8 Oct 2026 07:13:07 -0700
Subject: [PATCH 13/15] Support 1x2 tile layout
---
.../VectorContractMultiLevelUnroll.cpp | 10 +-
.../vector-contract-multi-level-unroll.mlir | 107 ++++++++++++++++--
2 files changed, 102 insertions(+), 15 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 5a3b1bad44ac29..02c4ae1add2662 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -358,16 +358,16 @@ static LogicalResult matchShapesAndTypes(MLUCandidate &candidate,
regTileN = origN % (2 * nativeN) == 0 ? 2 * nativeN : nativeN;
regTileK = nativeK;
} else {
- // Online packing requires a 2x2 tile register layout.
+ // Online packing requires a contract pair in the N-dimension.
nativeK = 16 * vnni;
- regTileM = 2 * nativeM;
- regTileN = 2 * nativeN;
- regTileK = nativeK;
- if (origM % regTileM != 0 || origN % regTileN != 0)
+ if (origM % nativeM != 0 || origN % (2 * nativeN) != 0)
return rewriter.notifyMatchFailure(
contract, "vector shape cannot be cleanly unrolled");
if (origK != nativeK)
return rewriter.notifyMatchFailure(contract, "K dimension mismatch");
+ regTileM = origM % (2 * nativeM) == 0 ? 2 * nativeM : nativeM;
+ regTileN = 2 * nativeN;
+ regTileK = nativeK;
}
} else {
return rewriter.notifyMatchFailure(contract, "unsupported target");
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
index 4097b4bed7a13b..b76541abfd080f 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
@@ -551,6 +551,93 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+// M = 48 is an odd multiple of 16, so a 1x2 tile register layout is used.
+// Online packing of the flat RHS needs the pair of contracts along N, which
+// this layout still provides.
+func.func @amx_bf16_flat_1x2_tile(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<48x64xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<48x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<48x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x64xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<48x32xbf16>, vector<32x64xbf16> into vector<48x64xf32>
+ scf.yield %d : vector<48x64xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x64xf32>, memref<?x?xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @amx_bf16_flat_1x2_tile(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C48:.+]] = arith.constant 48 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [48, 64] [1, 1]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C48]] step %[[C16]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [16, 32] [1, 1]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]]] [16, 32] [1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]]] [32, 32] [1, 1]
+// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: %[[B1:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A0]], %[[B1]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>}
+// CHECK: scf.yield %[[D0]], %[[D1]]
+// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+
+// NANO-LABEL: func.func @amx_bf16_flat_1x2_tile(
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
// The accumulation loop sits inside an existing batch-M-N loop nest.
func.func @amx_int8_flat_batched_loop_nest(%A: memref<?x?x?xi8>, %B: memref<?x?x?xi8>, %C: memref<?x?x?xi32>) {
%c0 = arith.constant 0 : index
@@ -1589,8 +1676,8 @@ module attributes {transform.with_named_sequence} {
// -----
-// Flat layout requires a 2x2 tile register layout, i.e. M and N must be
-// multiples of 32.
+// N = 48 would require a 2x1 tile register layout, but online packing of the
+// flat RHS needs a pair of contracts along N, i.e. N must be a multiple of 32.
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
@@ -1600,23 +1687,23 @@ func.func @negative_shape_not_unrollable(%A: memref<?x?xbf16>, %B: memref<?x?xbf
%c32 = arith.constant 32 : index
%c0_0 = arith.constant 0.0 : bf16
%c0_0_f32 = arith.constant 0.0 : f32
- %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<48x96xf32>
- %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<48x96xf32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<48x32xbf16>
- %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x48xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x48xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x48xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
iterator_types = ["parallel", "parallel", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<48x32xbf16>, vector<32x96xbf16> into vector<48x96xf32>
- scf.yield %d : vector<48x96xf32>
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x48xbf16> into vector<64x48xf32>
+ scf.yield %d : vector<64x48xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x96xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x48xf32>, memref<?x?xf32>
func.return
}
// CHECK-LABEL: func.func @negative_shape_not_unrollable(
-// CHECK: vector.contract {{.*}} : vector<48x32xbf16>, vector<32x96xbf16> into vector<48x96xf32>
+// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<32x48xbf16> into vector<64x48xf32>
// NANO-LABEL: func.func @negative_shape_not_unrollable(
// NANO: vector.contract
>From c7c895ff3e292fe800629e860f65a794b147e54c Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Thu, 8 Oct 2026 07:39:36 -0700
Subject: [PATCH 14/15] Reordered tests
---
.../vector-contract-multi-level-unroll.mlir | 1410 ++++++++---------
1 file changed, 705 insertions(+), 705 deletions(-)
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
index b76541abfd080f..479ad0d51c43b9 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll.mlir
@@ -93,59 +93,70 @@ module attributes {transform.with_named_sequence} {
// -----
-#map = affine_map<(d0, d1, d2) -> (d0, d2)>
-#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
-#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
-// Zero-initialized accumulator: no reads from the destination, and the K-loop
-// starts from a native-shaped zero constant.
-func.func @amx_bf16_flat_zero_init(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c32 = arith.constant 32 : index
+// M = 48 is an odd multiple of 16, so a 1x2 tile register layout is used.
+func.func @amx_bf16_vnni(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
%c0_0 = arith.constant 0.0 : bf16
- %c = arith.constant dense<0.0> : vector<96x128xf32>
- %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<96x128xf32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<96x32xbf16>
- %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x128xbf16>
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<48x64xf32>
+ %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<48x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<48x16x2xbf16>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<16x64x2xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction"],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<96x32xbf16>, vector<32x128xbf16> into vector<96x128xf32>
- scf.yield %d : vector<96x128xf32>
+ %a, %b, %acc : vector<48x16x2xbf16>, vector<16x64x2xbf16> into vector<48x64xf32>
+ scf.yield %d : vector<48x64xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<96x128xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x64xf32>, memref<?x?xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_zero_init(
-// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-LABEL: func.func @amx_bf16_vnni(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x2xbf16>, %[[B:.+]]: memref<?x?x2xbf16>, %[[C:.+]]: memref<?x?xf32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
-// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
-// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
-// CHECK-DAG: %[[ZERO:.+]] = arith.constant dense<0.000000e+00> : vector<16x16xf32>
+// CHECK-DAG: %[[C48:.+]] = arith.constant 48 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-NOT: memref.alloca
-// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [96, 128] [1, 1]
-// CHECK-NOT: vector.transfer_read
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
-// CHECK-NOT: vector.transfer_read
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
-// CHECK-SAME: iter_args(%{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]])
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
-// CHECK: scf.yield
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [48, 64] [1, 1]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C48]] step %[[C16]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [16, 32] [1, 1]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]], 0] [16, 16, 2] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]], 0] [16, 32, 2] [1, 1, 1]
+// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: %[[B1:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A0]], %[[B1]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>}
+// CHECK: scf.yield %[[D0]], %[[D1]]
// CHECK: }
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
// CHECK: }
// CHECK: }
+// CHECK-NOT: vector.transfer_read
// CHECK-NOT: vector.transfer_write
// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_flat_zero_init(
+// NANO-LABEL: func.func @amx_bf16_vnni(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -172,53 +183,66 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// The initial accumulator is not backed by memory, so it is spilled to a stack
-// buffer. The result is still written through a subview of the destination.
-func.func @amx_bf16_flat_init_from_arg(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
- %c: vector<64x96xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c32 = arith.constant 32 : index
+// Blocked layout: the two leading memref dims select a block, and the block
+// has the same shape as the operand vectors.
+func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x32xbf16>, %B: memref<?x?x32x64xbf16>, %C: memref<?x?x96x64xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
%c0_0 = arith.constant 0.0 : bf16
- %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
- %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n, %c0, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?x96x64xf32>, vector<96x64xf32>
+ %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c) -> (vector<96x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x96x32xbf16>, vector<96x32xbf16>
+ %b = vector.transfer_read %B[%n, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x32x64xbf16>, vector<32x64xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
iterator_types = ["parallel", "parallel", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
- scf.yield %d : vector<64x96xf32>
+ %a, %b, %acc : vector<96x32xbf16>, vector<32x64xbf16> into vector<96x64xf32>
+ scf.yield %d : vector<96x64xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<96x64xf32>, memref<?x?x96x64xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_init_from_arg(
-// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
-// CHECK-SAME: %[[ACC_INIT:.+]]: vector<64x96xf32>,
+// CHECK-LABEL: func.func @amx_bf16_flat_blocked(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x96x32xbf16>, %[[B:.+]]: memref<?x?x32x64xbf16>, %[[C:.+]]: memref<?x?x96x64xf32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
-// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
-// CHECK: vector.transfer_write %[[ACC_INIT]], %[[BUF]][%[[C0]], %[[C0]]]
-// CHECK: %[[ACC_OUT:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
// CHECK-NOT: memref.alloca
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
-// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
-// CHECK: %[[RES:.+]]:4 = scf.for
-// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[ACC_OUT]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]], 0, 0] [1, 1, 96, 64] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x96x64xf32> to memref<96x64xf32, strided<[64, 1], offset: ?>>
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C1]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], 0] [1, 1, 32, 32] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x96x32xbf16> to memref<32x32xbf16, strided<[32, 1], offset: ?>>
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], 0, %[[IV_N]]] [1, 1, 32, 32] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x32x64xbf16> to memref<32x32xbf16, strided<[64, 1], offset: ?>>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
// CHECK: }
// CHECK-NOT: vector.transfer_read
// CHECK-NOT: vector.transfer_write
// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_flat_init_from_arg(
+// NANO-LABEL: func.func @amx_bf16_flat_blocked(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -241,58 +265,68 @@ module attributes {transform.with_named_sequence} {
// -----
-#map = affine_map<(d0, d1, d2) -> (d0, d2)>
-#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
-#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
-// The loop result feeds an epilogue instead of a plain store, so it is
-// buffered on the stack and rematerialized. The initial accumulator is still
-// read through a subview.
-func.func @amx_bf16_flat_truncf_result(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>, %D: memref<?x?xbf16>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c32 = arith.constant 32 : index
+func.func @amx_bf16_vnni_blocked(%A: memref<?x?x128x16x2xbf16>, %B: memref<?x?x16x96x2xbf16>, %C: memref<?x?x128x96xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
%c0_0 = arith.constant 0.0 : bf16
%c0_0_f32 = arith.constant 0.0 : f32
- %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
- %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
- %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %c = vector.transfer_read %C[%m, %n, %c0, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?x128x96xf32>, vector<128x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c) -> (vector<128x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0, %c0, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x128x16x2xbf16>, vector<128x16x2xbf16>
+ %b = vector.transfer_read %B[%n, %k, %c0, %c0, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x16x96x2xbf16>, vector<16x96x2xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction"],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
- scf.yield %d : vector<64x96xf32>
+ %a, %b, %acc : vector<128x16x2xbf16>, vector<16x96x2xbf16> into vector<128x96xf32>
+ scf.yield %d : vector<128x96xf32>
}
- %trunc = arith.truncf %res : vector<64x96xf32> to vector<64x96xbf16>
- vector.transfer_write %trunc, %D[%m, %n] {in_bounds = [true, true]} : vector<64x96xbf16>, memref<?x?xbf16>
+ vector.transfer_write %res, %C[%m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<128x96xf32>, memref<?x?x128x96xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_truncf_result(
-// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>, %[[D:.+]]: memref<?x?xbf16>,
+// CHECK-LABEL: func.func @amx_bf16_vnni_blocked(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x128x16x2xbf16>, %[[B:.+]]: memref<?x?x16x96x2xbf16>, %[[C:.+]]: memref<?x?x128x96xf32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
-// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
-// CHECK-DAG: %[[ACC_IN:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
-// CHECK-DAG: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
-// CHECK-NOT: vector.transfer_write
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]], 0, 0] [1, 1, 128, 96] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x128x96xf32> to memref<128x96xf32, strided<[96, 1], offset: ?>>
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
-// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[ACC_IN]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
-// CHECK: %[[RES:.+]]:4 = scf.for
-// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C1]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], 0, 0] [1, 1, 32, 16, 2] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x128x16x2xbf16> to memref<32x16x2xbf16, strided<[32, 2, 1], offset: ?>>
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], 0, %[[IV_N]], 0] [1, 1, 16, 32, 2] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x16x96x2xbf16> to memref<16x32x2xbf16, strided<[192, 2, 1], offset: ?>>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
// CHECK: }
-// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
-// CHECK: %[[TRUNC:.+]] = arith.truncf %[[ACC_RES]]
-// CHECK: vector.transfer_write %[[TRUNC]], %[[D]][%[[M]], %[[N]]]
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_flat_truncf_result(
+// NANO-LABEL: func.func @amx_bf16_vnni_blocked(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -319,60 +353,67 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// Neither the initial read nor the final write are fully in-bounds, so a
-// single stack buffer is used for both; the original transfers remain.
-func.func @amx_bf16_flat_acc_out_of_bounds(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c32 = arith.constant 32 : index
- %c0_0 = arith.constant 0.0 : bf16
- %c0_0_f32 = arith.constant 0.0 : f32
- %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [false, false]} : memref<?x?xf32>, vector<64x96xf32>
- %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
- %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+func.func @amx_int8_flat(%A: memref<?x?xi8>, %B: memref<?x?xi8>, %C: memref<?x?xi32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c64 = arith.constant 64 : index
+ %c0_i8 = arith.constant 0 : i8
+ %c0_i32 = arith.constant 0 : i32
+ %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x128xi32>
+ %res = scf.for %k = %k_start to %k_end step %c64 iter_args(%acc = %c) -> (vector<64x128xi32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?xi8>, vector<64x64xi8>
+ %b = vector.transfer_read %B[%k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?xi8>, vector<64x128xi8>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
iterator_types = ["parallel", "parallel", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
- scf.yield %d : vector<64x96xf32>
+ %a, %b, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
+ scf.yield %d : vector<64x128xi32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [false, false]} : vector<64x96xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?xi32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_acc_out_of_bounds(
-// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-LABEL: func.func @amx_int8_flat(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xi8>, %[[B:.+]]: memref<?x?xi8>, %[[C:.+]]: memref<?x?xi32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
-// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
-// CHECK-NOT: memref.subview %[[C]]
-// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]], {{.*}} : memref<?x?xf32>, vector<64x96xf32>
-// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
-// CHECK: vector.transfer_write %[[ACC_INIT]], %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK-NOT: memref.alloca
-// CHECK-NOT: memref.subview %[[C]]
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 128] [1, 1]
// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
-// CHECK: %[[RES:.+]]:4 = scf.for
-// CHECK-NOT: memref.subview %[[BUF]]
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C64]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]]] [32, 64] [1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]]] [64, 32] [1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
+// CHECK: scf.yield
+// CHECK: }
// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
// CHECK: }
-// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
-// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]] : vector<64x96xf32>, memref<?x?xf32>
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_flat_acc_out_of_bounds(
+// NANO-LABEL: func.func @amx_int8_flat(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
} : !transform.any_op
transform.yield
}
@@ -380,7 +421,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -389,67 +430,76 @@ module attributes {transform.with_named_sequence} {
// -----
-#map = affine_map<(d0, d1, d2) -> (d0, d2)>
-#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
-#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
-// The initial read and the final write target the same memref but at different
-// offsets, so the regions partially overlap. Writing the result tiles through
-// a view would clobber initial values not yet read by later M/N iterations, so
-// the initial value is still read via a subview, but the result goes through a
-// stack buffer and the original write remains.
-func.func @amx_bf16_flat_acc_partial_overlap(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c32 = arith.constant 32 : index
- %c0_0 = arith.constant 0.0 : bf16
- %c0_0_f32 = arith.constant 0.0 : f32
- %m_off = arith.addi %m, %c32 : index
- %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
- %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
- %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+// N = 48 is an odd multiple of 16, so a 2x1 tile register layout is used.
+func.func @amx_int8_vnni(%A: memref<?x?x4xi8>, %B: memref<?x?x4xi8>, %C: memref<?x?xi32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c0_i8 = arith.constant 0 : i8
+ %c0_i32 = arith.constant 0 : i32
+ %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x48xi32>
+ %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<64x48xi32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0], %c0_i8 {in_bounds = [true, true, true]} : memref<?x?x4xi8>, vector<64x16x4xi8>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_i8 {in_bounds = [true, true, true]} : memref<?x?x4xi8>, vector<16x48x4xi8>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction"],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
- scf.yield %d : vector<64x96xf32>
+ %a, %b, %acc : vector<64x16x4xi8>, vector<16x48x4xi8> into vector<64x48xi32>
+ scf.yield %d : vector<64x48xi32>
}
- vector.transfer_write %res, %C[%m_off, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x48xi32>, memref<?x?xi32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_acc_partial_overlap(
-// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-LABEL: func.func @amx_int8_vnni(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x4xi8>, %[[B:.+]]: memref<?x?x4xi8>, %[[C:.+]]: memref<?x?xi32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C48:.+]] = arith.constant 48 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
-// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
-// CHECK-DAG: %[[M_OFF:.+]] = arith.addi %[[M]], %[[C32]]
-// CHECK-DAG: %[[ACC_IN:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
-// CHECK-DAG: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
-// CHECK-NOT: memref.subview %[[C]]
-// CHECK-NOT: vector.transfer_write
+// CHECK-NOT: memref.alloca
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 48] [1, 1]
// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
-// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[ACC_IN]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
-// CHECK: %[[RES:.+]]:4 = scf.for
-// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C48]] step %[[C16]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xi32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xi32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]], 0] [32, 16, 4] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]], 0] [16, 16, 4] [1, 1, 1]
+// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
+// CHECK: %[[A1:.+]] = vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
+// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
+// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 4>} : vector<16x16x4xi8>, vector<16x16x4xi8> into vector<16x16xi32>
+// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A1]], %[[B0]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 4>}
+// CHECK: scf.yield %[[D0]], %[[D1]]
+// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
// CHECK: }
// CHECK: }
-// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
-// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M_OFF]], %[[N]]] {{.*}} : vector<64x96xf32>, memref<?x?xf32>
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_flat_acc_partial_overlap(
+// NANO-LABEL: func.func @amx_int8_vnni(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
} : !transform.any_op
transform.yield
}
@@ -457,7 +507,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -470,67 +520,61 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-func.func @amx_int8_flat(%A: memref<?x?xi8>, %B: memref<?x?xi8>, %C: memref<?x?xi32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c64 = arith.constant 64 : index
- %c0_i8 = arith.constant 0 : i8
- %c0_i32 = arith.constant 0 : i32
- %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x128xi32>
- %res = scf.for %k = %k_start to %k_end step %c64 iter_args(%acc = %c) -> (vector<64x128xi32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?xi8>, vector<64x64xi8>
- %b = vector.transfer_read %B[%k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?xi8>, vector<64x128xi8>
+// Zero-initialized accumulator: no reads from the destination, and the K-loop
+// starts from a native-shaped zero constant.
+func.func @amx_bf16_flat_zero_init(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c = arith.constant dense<0.0> : vector<96x128xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<96x128xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<96x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x128xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
iterator_types = ["parallel", "parallel", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
- scf.yield %d : vector<64x128xi32>
+ %a, %b, %acc : vector<96x32xbf16>, vector<32x128xbf16> into vector<96x128xf32>
+ scf.yield %d : vector<96x128xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?xi32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<96x128xf32>, memref<?x?xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_int8_flat(
-// CHECK-SAME: %[[A:.+]]: memref<?x?xi8>, %[[B:.+]]: memref<?x?xi8>, %[[C:.+]]: memref<?x?xi32>,
+// CHECK-LABEL: func.func @amx_bf16_flat_zero_init(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
-// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG: %[[ZERO:.+]] = arith.constant dense<0.000000e+00> : vector<16x16xf32>
// CHECK-NOT: memref.alloca
-// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 128] [1, 1]
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [96, 128] [1, 1]
+// CHECK-NOT: vector.transfer_read
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
-// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
-// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C64]]
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]]] [32, 64] [1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]]] [64, 32] [1, 1]
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
+// CHECK-NOT: vector.transfer_read
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
+// CHECK-SAME: iter_args(%{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]], %{{.+}} = %[[ZERO]])
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
// CHECK: scf.yield
// CHECK: }
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
// CHECK: }
-// CHECK-NOT: vector.transfer_read
// CHECK-NOT: vector.transfer_write
// CHECK: return
-// NANO-LABEL: func.func @amx_int8_flat(
+// NANO-LABEL: func.func @amx_bf16_flat_zero_init(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -538,7 +582,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -551,67 +595,53 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// M = 48 is an odd multiple of 16, so a 1x2 tile register layout is used.
-// Online packing of the flat RHS needs the pair of contracts along N, which
-// this layout still provides.
-func.func @amx_bf16_flat_1x2_tile(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
+// The initial accumulator is not backed by memory, so it is spilled to a stack
+// buffer. The result is still written through a subview of the destination.
+func.func @amx_bf16_flat_init_from_arg(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %c: vector<64x96xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
%c32 = arith.constant 32 : index
%c0_0 = arith.constant 0.0 : bf16
- %c0_0_f32 = arith.constant 0.0 : f32
- %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<48x64xf32>
- %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<48x64xf32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<48x32xbf16>
- %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x64xbf16>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
iterator_types = ["parallel", "parallel", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<48x32xbf16>, vector<32x64xbf16> into vector<48x64xf32>
- scf.yield %d : vector<48x64xf32>
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x64xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_1x2_tile(
+// CHECK-LABEL: func.func @amx_bf16_flat_init_from_arg(
// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[ACC_INIT:.+]]: vector<64x96xf32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
-// CHECK-DAG: %[[C48:.+]] = arith.constant 48 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_INIT]], %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK: %[[ACC_OUT:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
// CHECK-NOT: memref.alloca
-// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [48, 64] [1, 1]
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C48]] step %[[C16]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [16, 32] [1, 1]
-// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
-// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
-// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
-// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]])
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]]] [16, 32] [1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]]] [32, 32] [1, 1]
-// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
-// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
-// CHECK: %[[B1:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
-// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
-// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A0]], %[[B1]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>}
-// CHECK: scf.yield %[[D0]], %[[D1]]
-// CHECK: }
-// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
-// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for
+// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[ACC_OUT]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
// CHECK: }
// CHECK: }
// CHECK-NOT: vector.transfer_read
// CHECK-NOT: vector.transfer_write
// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_flat_1x2_tile(
+// NANO-LABEL: func.func @amx_bf16_flat_init_from_arg(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -638,93 +668,60 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// The accumulation loop sits inside an existing batch-M-N loop nest.
-func.func @amx_int8_flat_batched_loop_nest(%A: memref<?x?x?xi8>, %B: memref<?x?x?xi8>, %C: memref<?x?x?xi32>) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- %c2 = arith.constant 2 : index
- %c64 = arith.constant 64 : index
- %c128 = arith.constant 128 : index
- %c0_i8 = arith.constant 0 : i8
- %c0_i32 = arith.constant 0 : i32
- %batch = memref.dim %C, %c0 : memref<?x?x?xi32>
- %M = memref.dim %C, %c1 : memref<?x?x?xi32>
- %N = memref.dim %C, %c2 : memref<?x?x?xi32>
- %K = memref.dim %A, %c2 : memref<?x?x?xi8>
- scf.for %bi = %c0 to %batch step %c1 {
- scf.for %m = %c0 to %M step %c64 {
- scf.for %n = %c0 to %N step %c128 {
- %c = vector.transfer_read %C[%bi, %m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?x?xi32>, vector<64x128xi32>
- %res = scf.for %k = %c0 to %K step %c64 iter_args(%acc = %c) -> (vector<64x128xi32>) {
- %a = vector.transfer_read %A[%bi, %m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x64xi8>
- %b = vector.transfer_read %B[%bi, %k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x128xi8>
- %d = vector.contract {
- indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction"],
- kind = #vector.kind<add>}
- %a, %b, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
- scf.yield %d : vector<64x128xi32>
- }
- vector.transfer_write %res, %C[%bi, %m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?x?xi32>
- }
- }
+// The loop result feeds an epilogue instead of a plain store, so it is
+// buffered on the stack and rematerialized. The initial accumulator is still
+// read through a subview.
+func.func @amx_bf16_flat_truncf_result(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>, %D: memref<?x?xbf16>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
}
+ %trunc = arith.truncf %res : vector<64x96xf32> to vector<64x96xbf16>
+ vector.transfer_write %trunc, %D[%m, %n] {in_bounds = [true, true]} : vector<64x96xbf16>, memref<?x?xbf16>
func.return
}
-// CHECK-LABEL: func.func @amx_int8_flat_batched_loop_nest(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x?xi8>, %[[B:.+]]: memref<?x?x?xi8>, %[[C:.+]]: memref<?x?x?xi32>)
+// CHECK-LABEL: func.func @amx_bf16_flat_truncf_result(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>, %[[D:.+]]: memref<?x?xbf16>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
-// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
-// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
-// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[C]], %[[C0]]
-// CHECK-DAG: %[[M:.+]] = memref.dim %[[C]], %[[C1]]
-// CHECK-DAG: %[[N:.+]] = memref.dim %[[C]], %[[C2]]
-// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
-// CHECK-NOT: memref.alloca
-// CHECK: scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]] {
-// CHECK: scf.for %[[OFF_M:.+]] = %[[C0]] to %[[M]] step %[[C64]] {
-// CHECK: scf.for %[[OFF_N:.+]] = %[[C0]] to %[[N]] step %[[C128]] {
-// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[IV_B]], %[[OFF_M]], %[[OFF_N]]] [1, 64, 128] [1, 1, 1]
-// CHECK-SAME: memref<?x?x?xi32> to memref<64x128xi32, strided<[?, 1], offset: ?>>
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
-// CHECK: %[[A_OFF_M:.+]] = arith.addi %[[OFF_M]], %[[IV_M]]
-// CHECK: %[[B_OFF_N:.+]] = arith.addi %[[OFF_N]], %[[IV_N]]
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C64]]
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[A_OFF_M]], %[[IV_K]]] [1, 32, 64] [1, 1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[B_OFF_N]]] [1, 64, 32] [1, 1, 1]
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
-// CHECK: scf.yield
-// CHECK: }
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
-// CHECK: }
-// CHECK: }
-// CHECK-NOT: vector.transfer_read
-// CHECK-NOT: vector.transfer_write
-// CHECK: }
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-DAG: %[[ACC_IN:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
+// CHECK-DAG: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[ACC_IN]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for
+// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
// CHECK: }
// CHECK: }
-// CHECK: return
+// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
+// CHECK: %[[TRUNC:.+]] = arith.truncf %[[ACC_RES]]
+// CHECK: vector.transfer_write %[[TRUNC]], %[[D]][%[[M]], %[[N]]]
-// NANO-LABEL: func.func @amx_int8_flat_batched_loop_nest(
+// NANO-LABEL: func.func @amx_bf16_flat_truncf_result(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -732,7 +729,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -745,92 +742,60 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// The accumulator is loop-carried through an enclosing batch loop, so it is
-// neither read from nor written to memory directly.
-//
-// This test is currently intended to demonstrate successful lowering using
-// the fallback stack-buffer path. A more desirable lowering would keep the
-// accumulator on the tile registers across both loops.
-func.func @amx_int8_flat_batch_reduction(%A: memref<?x?x?xi8>, %B: memref<?x?x?xi8>, %C: memref<?x?xi32>,
- %m: index, %n: index) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- %c2 = arith.constant 2 : index
- %c64 = arith.constant 64 : index
- %c0_i8 = arith.constant 0 : i8
- %c0_i32 = arith.constant 0 : i32
- %batch = memref.dim %A, %c0 : memref<?x?x?xi8>
- %K = memref.dim %A, %c2 : memref<?x?x?xi8>
- %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x128xi32>
- %res = scf.for %b = %c0 to %batch step %c1 iter_args(%acc_b = %c) -> (vector<64x128xi32>) {
- %res_b = scf.for %k = %c0 to %K step %c64 iter_args(%acc = %acc_b) -> (vector<64x128xi32>) {
- %a = vector.transfer_read %A[%b, %m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x64xi8>
- %bv = vector.transfer_read %B[%b, %k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x128xi8>
- %d = vector.contract {
- indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction"],
- kind = #vector.kind<add>}
- %a, %bv, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
- scf.yield %d : vector<64x128xi32>
- }
- scf.yield %res_b : vector<64x128xi32>
+// Neither the initial read nor the final write are fully in-bounds, so a
+// single stack buffer is used for both; the original transfers remain.
+func.func @amx_bf16_flat_acc_out_of_bounds(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [false, false]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?xi32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [false, false]} : vector<64x96xf32>, memref<?x?xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_int8_flat_batch_reduction(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x?xi8>, %[[B:.+]]: memref<?x?x?xi8>, %[[C:.+]]: memref<?x?xi32>,
-// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index)
+// CHECK-LABEL: func.func @amx_bf16_flat_acc_out_of_bounds(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
-// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
-// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
-// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[A]], %[[C0]]
-// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-NOT: memref.subview %[[C]]
+// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]], {{.*}} : memref<?x?xf32>, vector<64x96xf32>
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_INIT]], %[[BUF]][%[[C0]], %[[C0]]]
// CHECK-NOT: memref.alloca
-// CHECK-NOT: memref.subview
-// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]]
-// CHECK: %[[ACC_RES:.+]] = scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]] iter_args(%[[ACC_B:.+]] = %[[ACC_INIT]])
-// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x128xi32>
-// CHECK: vector.transfer_write %[[ACC_B]], %[[BUF]][%[[C0]], %[[C0]]]
-// CHECK-NOT: memref.alloca
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
-// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
-// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C64]]
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[OFF_M]], %[[IV_K]]] [1, 32, 64] [1, 1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[OFF_N]]] [1, 64, 32] [1, 1, 1]
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
-// CHECK: scf.yield
-// CHECK: }
-// CHECK-NOT: memref.subview %[[BUF]]
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
-// CHECK: }
+// CHECK-NOT: memref.subview %[[C]]
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for
+// CHECK-NOT: memref.subview %[[BUF]]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
-// CHECK: %[[ACC_NEXT:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x128xi32>
-// CHECK: scf.yield %[[ACC_NEXT]]
// CHECK: }
-// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]]
+// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]] : vector<64x96xf32>, memref<?x?xf32>
-// NANO-LABEL: func.func @amx_int8_flat_batch_reduction(
+// NANO-LABEL: func.func @amx_bf16_flat_acc_out_of_bounds(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -838,7 +803,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -851,87 +816,57 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// Loop order K-batch, accumulating over both loops: the inner batch loop
-// carries the accumulator, its IV appears exactly once in each operand read and
-// its step is 1. It is a valid accumulation loop, indistinguishable from a
-// K-loop over a blocked layout, and is treated as such.
-//
-// Again, spilling the accumulator value to memory in the outer loop is correct,
-// but not ideal.
-func.func @amx_bf16_flat_k_batch_order(%A: memref<?x?x?xbf16>, %B: memref<?x?x?xbf16>, %C: memref<?x?xf32>,
- %m: index, %n: index) {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- %c2 = arith.constant 2 : index
+// The initial read and the final write target the same memref but at different
+// offsets, so the regions partially overlap. Writing the result tiles through
+// a view would clobber initial values not yet read by later M/N iterations, so
+// the initial value is still read via a subview, but the result goes through a
+// stack buffer and the original write remains.
+func.func @amx_bf16_flat_acc_partial_overlap(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
%c32 = arith.constant 32 : index
%c0_0 = arith.constant 0.0 : bf16
%c0_0_f32 = arith.constant 0.0 : f32
- %batch = memref.dim %A, %c0 : memref<?x?x?xbf16>
- %K = memref.dim %A, %c2 : memref<?x?x?xbf16>
+ %m_off = arith.addi %m, %c32 : index
%c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
- %res = scf.for %k = %c0 to %K step %c32 iter_args(%acc_k = %c) -> (vector<64x96xf32>) {
- %res_k = scf.for %bi = %c0 to %batch step %c1 iter_args(%acc = %acc_k) -> (vector<64x96xf32>) {
- %a = vector.transfer_read %A[%bi, %m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?x?xbf16>, vector<64x32xbf16>
- %b = vector.transfer_read %B[%bi, %k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?x?xbf16>, vector<32x96xbf16>
- %d = vector.contract {
- indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction"],
- kind = #vector.kind<add>}
- %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
- scf.yield %d : vector<64x96xf32>
- }
- scf.yield %res_k : vector<64x96xf32>
- }
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
- func.return
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ vector.transfer_write %res, %C[%m_off, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
+ func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_k_batch_order(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x?xbf16>, %[[B:.+]]: memref<?x?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
-// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index)
+// CHECK-LABEL: func.func @amx_bf16_flat_acc_partial_overlap(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
-// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
-// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[A]], %[[C0]]
-// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
-// CHECK-NOT: memref.alloca
-// CHECK-NOT: memref.subview
-// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]]
-// CHECK: %[[ACC_RES:.+]] = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C32]] iter_args(%[[ACC_K:.+]] = %[[ACC_INIT]])
-// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
-// CHECK: vector.transfer_write %[[ACC_K]], %[[BUF]][%[[C0]], %[[C0]]]
-// CHECK-NOT: memref.alloca
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
-// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
-// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]]
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[OFF_M]], %[[IV_K]]] [1, 32, 32] [1, 1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[OFF_N]]] [1, 32, 32] [1, 1, 1]
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
-// CHECK: scf.yield
-// CHECK: }
-// CHECK-NOT: memref.subview %[[BUF]]
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
-// CHECK: }
+// CHECK-DAG: %[[M_OFF:.+]] = arith.addi %[[M]], %[[C32]]
+// CHECK-DAG: %[[ACC_IN:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 96] [1, 1]
+// CHECK-DAG: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK-NOT: memref.subview %[[C]]
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[IN_VIEW:.+]] = memref.subview %[[ACC_IN]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[IN_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for
+// CHECK: %[[OUT_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[OUT_VIEW]]
// CHECK: }
-// CHECK: %[[ACC_NEXT:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
-// CHECK: scf.yield %[[ACC_NEXT]]
// CHECK: }
-// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]]
+// CHECK: %[[ACC_RES:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M_OFF]], %[[N]]] {{.*}} : vector<64x96xf32>, memref<?x?xf32>
-// NANO-LABEL: func.func @amx_bf16_flat_k_batch_order(
+// NANO-LABEL: func.func @amx_bf16_flat_acc_partial_overlap(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -954,34 +889,35 @@ module attributes {transform.with_named_sequence} {
// -----
-#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
-#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
-#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
// M = 48 is an odd multiple of 16, so a 1x2 tile register layout is used.
-func.func @amx_bf16_vnni(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: memref<?x?xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c0 = arith.constant 0 : index
- %c16 = arith.constant 16 : index
+// Online packing of the flat RHS needs the pair of contracts along N, which
+// this layout still provides.
+func.func @amx_bf16_flat_1x2_tile(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
%c0_0 = arith.constant 0.0 : bf16
%c0_0_f32 = arith.constant 0.0 : f32
%c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<48x64xf32>
- %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<48x64xf32>) {
- %a = vector.transfer_read %A[%m, %k, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<48x16x2xbf16>
- %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<16x64x2xbf16>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<48x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<48x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x64xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+ iterator_types = ["parallel", "parallel", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<48x16x2xbf16>, vector<16x64x2xbf16> into vector<48x64xf32>
+ %a, %b, %acc : vector<48x32xbf16>, vector<32x64xbf16> into vector<48x64xf32>
scf.yield %d : vector<48x64xf32>
}
vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x64xf32>, memref<?x?xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_vnni(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x2xbf16>, %[[B:.+]]: memref<?x?x2xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-LABEL: func.func @amx_bf16_flat_1x2_tile(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
@@ -997,16 +933,16 @@ func.func @amx_bf16_vnni(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: mem
// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
-// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
+// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]])
// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]], 0] [16, 16, 2] [1, 1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]], 0] [16, 32, 2] [1, 1, 1]
-// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: %[[B1:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
-// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A0]], %[[B1]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>}
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]]] [16, 32] [1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]]] [32, 32] [1, 1]
+// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: %[[B1:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A0]], %[[B1]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>}
// CHECK: scf.yield %[[D0]], %[[D1]]
// CHECK: }
// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
@@ -1017,7 +953,7 @@ func.func @amx_bf16_vnni(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: mem
// CHECK-NOT: vector.transfer_write
// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_vnni(
+// NANO-LABEL: func.func @amx_bf16_flat_1x2_tile(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -1040,76 +976,72 @@ module attributes {transform.with_named_sequence} {
// -----
-#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
-#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
-#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// N = 48 is an odd multiple of 16, so a 2x1 tile register layout is used.
-func.func @amx_int8_vnni(%A: memref<?x?x4xi8>, %B: memref<?x?x4xi8>, %C: memref<?x?xi32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c0 = arith.constant 0 : index
- %c16 = arith.constant 16 : index
- %c0_i8 = arith.constant 0 : i8
- %c0_i32 = arith.constant 0 : i32
- %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x48xi32>
- %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<64x48xi32>) {
- %a = vector.transfer_read %A[%m, %k, %c0], %c0_i8 {in_bounds = [true, true, true]} : memref<?x?x4xi8>, vector<64x16x4xi8>
- %b = vector.transfer_read %B[%k, %n, %c0], %c0_i8 {in_bounds = [true, true, true]} : memref<?x?x4xi8>, vector<16x48x4xi8>
+// The original shape equals the register tile shape, so the M-N loop nest has
+// a single iteration and folds away.
+func.func @amx_bf16_flat_single_reg_tile(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c32 = arith.constant 32 : index
+ %c0_0 = arith.constant 0.0 : bf16
+ %c0_0_f32 = arith.constant 0.0 : f32
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<32x32xf32>
+ %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<32x32xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x32xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x32xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+ iterator_types = ["parallel", "parallel", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<64x16x4xi8>, vector<16x48x4xi8> into vector<64x48xi32>
- scf.yield %d : vector<64x48xi32>
+ %a, %b, %acc : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %d : vector<32x32xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x48xi32>, memref<?x?xi32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<32x32xf32>, memref<?x?xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_int8_vnni(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x4xi8>, %[[B:.+]]: memref<?x?x4xi8>, %[[C:.+]]: memref<?x?xi32>,
+// CHECK-LABEL: func.func @amx_bf16_flat_single_reg_tile(
+// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
-// CHECK-DAG: %[[C48:.+]] = arith.constant 48 : index
-// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-NOT: memref.alloca
-// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [64, 48] [1, 1]
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C48]] step %[[C16]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
-// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xi32>
-// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xi32>
-// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
-// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
-// CHECK: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
-// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]])
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[OFF_M]], %[[IV_K]], 0] [32, 16, 4] [1, 1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[OFF_N]], 0] [16, 16, 4] [1, 1, 1]
-// CHECK: %[[A0:.+]] = vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
-// CHECK: %[[A1:.+]] = vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
-// CHECK: %[[B0:.+]] = vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x4xi8>
-// CHECK: %[[D0:.+]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ARG0]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 4>} : vector<16x16x4xi8>, vector<16x16x4xi8> into vector<16x16xi32>
-// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A1]], %[[B0]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 4>}
-// CHECK: scf.yield %[[D0]], %[[D1]]
-// CHECK: }
-// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
-// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
-// CHECK: }
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [32, 32] [1, 1]
+// CHECK-NOT: memref.subview %[[ACC_VIEW]]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC2:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC3:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]], %[[ARG2:.+]] = %[[ACC2]], %[[ARG3:.+]] = %[[ACC3]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]]] [32, 32] [1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[N]]] [32, 32] [1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
+// CHECK: vector.transfer_write %[[RES]]#2, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#3, %[[ACC_VIEW]][%[[C16]], %[[C16]]]
// CHECK-NOT: vector.transfer_read
// CHECK-NOT: vector.transfer_write
// CHECK: return
-// NANO-LABEL: func.func @amx_int8_vnni(
+// NANO-LABEL: func.func @amx_bf16_flat_single_reg_tile(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
} : !transform.any_op
transform.yield
}
@@ -1117,7 +1049,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -1126,70 +1058,64 @@ module attributes {transform.with_named_sequence} {
// -----
-#map = affine_map<(d0, d1, d2) -> (d0, d2)>
-#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
-#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
-// Blocked layout: the two leading memref dims select a block, and the block
-// has the same shape as the operand vectors.
-func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x32xbf16>, %B: memref<?x?x32x64xbf16>, %C: memref<?x?x96x64xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
+func.func @amx_bf16_vnni_single_reg_tile(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
%c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
%c0_0 = arith.constant 0.0 : bf16
%c0_0_f32 = arith.constant 0.0 : f32
- %c = vector.transfer_read %C[%m, %n, %c0, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?x96x64xf32>, vector<96x64xf32>
- %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c) -> (vector<96x64xf32>) {
- %a = vector.transfer_read %A[%m, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x96x32xbf16>, vector<96x32xbf16>
- %b = vector.transfer_read %B[%n, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x32x64xbf16>, vector<32x64xbf16>
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<32x32xf32>
+ %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<32x32xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<32x16x2xbf16>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<16x32x2xbf16>
%d = vector.contract {
indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction"],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
kind = #vector.kind<add>}
- %a, %b, %acc : vector<96x32xbf16>, vector<32x64xbf16> into vector<96x64xf32>
- scf.yield %d : vector<96x64xf32>
+ %a, %b, %acc : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+ scf.yield %d : vector<32x32xf32>
}
- vector.transfer_write %res, %C[%m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<96x64xf32>, memref<?x?x96x64xf32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<32x32xf32>, memref<?x?xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_blocked(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x96x32xbf16>, %[[B:.+]]: memref<?x?x32x64xbf16>, %[[C:.+]]: memref<?x?x96x64xf32>,
+// CHECK-LABEL: func.func @amx_bf16_vnni_single_reg_tile(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x2xbf16>, %[[B:.+]]: memref<?x?x2xbf16>, %[[C:.+]]: memref<?x?xf32>,
// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
-// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
-// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
-// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
// CHECK-NOT: memref.alloca
-// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]], 0, 0] [1, 1, 96, 64] [1, 1, 1, 1]
-// CHECK-SAME: memref<?x?x96x64xf32> to memref<96x64xf32, strided<[64, 1], offset: ?>>
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C1]]
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], 0] [1, 1, 32, 32] [1, 1, 1, 1]
-// CHECK-SAME: memref<?x?x96x32xbf16> to memref<32x32xbf16, strided<[32, 1], offset: ?>>
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], 0, %[[IV_N]]] [1, 1, 32, 32] [1, 1, 1, 1]
-// CHECK-SAME: memref<?x?x32x64xbf16> to memref<32x32xbf16, strided<[64, 1], offset: ?>>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
-// CHECK: scf.yield
-// CHECK: }
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
-// CHECK: }
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [32, 32] [1, 1]
+// CHECK-NOT: memref.subview %[[ACC_VIEW]]
+// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC2:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[ACC3:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C16]]], {{.*}} vector<16x16xf32>
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]], %[[ARG2:.+]] = %[[ACC2]], %[[ARG3:.+]] = %[[ACC3]])
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], 0] [32, 16, 2] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[N]], 0] [16, 32, 2] [1, 1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
// CHECK: }
+// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
+// CHECK: vector.transfer_write %[[RES]]#2, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
+// CHECK: vector.transfer_write %[[RES]]#3, %[[ACC_VIEW]][%[[C16]], %[[C16]]]
// CHECK-NOT: vector.transfer_read
// CHECK-NOT: vector.transfer_write
// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_flat_blocked(
+// NANO-LABEL: func.func @amx_bf16_vnni_single_reg_tile(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -1212,74 +1138,97 @@ module attributes {transform.with_named_sequence} {
// -----
-#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
-#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
-#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-func.func @amx_bf16_vnni_blocked(%A: memref<?x?x128x16x2xbf16>, %B: memref<?x?x16x96x2xbf16>, %C: memref<?x?x128x96xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
+// The accumulation loop sits inside an existing batch-M-N loop nest.
+func.func @amx_int8_flat_batched_loop_nest(%A: memref<?x?x?xi8>, %B: memref<?x?x?xi8>, %C: memref<?x?x?xi32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
- %c0_0 = arith.constant 0.0 : bf16
- %c0_0_f32 = arith.constant 0.0 : f32
- %c = vector.transfer_read %C[%m, %n, %c0, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?x128x96xf32>, vector<128x96xf32>
- %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c) -> (vector<128x96xf32>) {
- %a = vector.transfer_read %A[%m, %k, %c0, %c0, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x128x16x2xbf16>, vector<128x16x2xbf16>
- %b = vector.transfer_read %B[%n, %k, %c0, %c0, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x16x96x2xbf16>, vector<16x96x2xbf16>
- %d = vector.contract {
- indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction", "reduction"],
- kind = #vector.kind<add>}
- %a, %b, %acc : vector<128x16x2xbf16>, vector<16x96x2xbf16> into vector<128x96xf32>
- scf.yield %d : vector<128x96xf32>
+ %c2 = arith.constant 2 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c0_i8 = arith.constant 0 : i8
+ %c0_i32 = arith.constant 0 : i32
+ %batch = memref.dim %C, %c0 : memref<?x?x?xi32>
+ %M = memref.dim %C, %c1 : memref<?x?x?xi32>
+ %N = memref.dim %C, %c2 : memref<?x?x?xi32>
+ %K = memref.dim %A, %c2 : memref<?x?x?xi8>
+ scf.for %bi = %c0 to %batch step %c1 {
+ scf.for %m = %c0 to %M step %c64 {
+ scf.for %n = %c0 to %N step %c128 {
+ %c = vector.transfer_read %C[%bi, %m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?x?xi32>, vector<64x128xi32>
+ %res = scf.for %k = %c0 to %K step %c64 iter_args(%acc = %c) -> (vector<64x128xi32>) {
+ %a = vector.transfer_read %A[%bi, %m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x64xi8>
+ %b = vector.transfer_read %B[%bi, %k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x128xi8>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
+ scf.yield %d : vector<64x128xi32>
+ }
+ vector.transfer_write %res, %C[%bi, %m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?x?xi32>
+ }
+ }
}
- vector.transfer_write %res, %C[%m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<128x96xf32>, memref<?x?x128x96xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_vnni_blocked(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x128x16x2xbf16>, %[[B:.+]]: memref<?x?x16x96x2xbf16>, %[[C:.+]]: memref<?x?x128x96xf32>,
-// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-LABEL: func.func @amx_int8_flat_batched_loop_nest(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x?xi8>, %[[B:.+]]: memref<?x?x?xi8>, %[[C:.+]]: memref<?x?x?xi32>)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
-// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[C]], %[[C0]]
+// CHECK-DAG: %[[M:.+]] = memref.dim %[[C]], %[[C1]]
+// CHECK-DAG: %[[N:.+]] = memref.dim %[[C]], %[[C2]]
+// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
// CHECK-NOT: memref.alloca
-// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]], 0, 0] [1, 1, 128, 96] [1, 1, 1, 1]
-// CHECK-SAME: memref<?x?x128x96xf32> to memref<128x96xf32, strided<[96, 1], offset: ?>>
-// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
-// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
-// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C1]]
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], 0, 0] [1, 1, 32, 16, 2] [1, 1, 1, 1, 1]
-// CHECK-SAME: memref<?x?x128x16x2xbf16> to memref<32x16x2xbf16, strided<[32, 2, 1], offset: ?>>
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], 0, %[[IV_N]], 0] [1, 1, 16, 32, 2] [1, 1, 1, 1, 1]
-// CHECK-SAME: memref<?x?x16x96x2xbf16> to memref<16x32x2xbf16, strided<[192, 2, 1], offset: ?>>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
-// CHECK: scf.yield
+// CHECK: scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]] {
+// CHECK: scf.for %[[OFF_M:.+]] = %[[C0]] to %[[M]] step %[[C64]] {
+// CHECK: scf.for %[[OFF_N:.+]] = %[[C0]] to %[[N]] step %[[C128]] {
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[IV_B]], %[[OFF_M]], %[[OFF_N]]] [1, 64, 128] [1, 1, 1]
+// CHECK-SAME: memref<?x?x?xi32> to memref<64x128xi32, strided<[?, 1], offset: ?>>
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[ACC_BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
+// CHECK: %[[A_OFF_M:.+]] = arith.addi %[[OFF_M]], %[[IV_M]]
+// CHECK: %[[B_OFF_N:.+]] = arith.addi %[[OFF_N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C64]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[A_OFF_M]], %[[IV_K]]] [1, 32, 64] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[B_OFF_N]]] [1, 64, 32] [1, 1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
// CHECK: }
-// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
// CHECK: }
-// CHECK-NOT: vector.transfer_read
-// CHECK-NOT: vector.transfer_write
// CHECK: return
-// NANO-LABEL: func.func @amx_bf16_vnni_blocked(
+// NANO-LABEL: func.func @amx_int8_flat_batched_loop_nest(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
} : !transform.any_op
transform.yield
}
@@ -1287,7 +1236,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -1300,68 +1249,92 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-// The original shape equals the register tile shape, so the M-N loop nest has
-// a single iteration and folds away.
-func.func @amx_bf16_flat_single_reg_tile(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
- %c32 = arith.constant 32 : index
- %c0_0 = arith.constant 0.0 : bf16
- %c0_0_f32 = arith.constant 0.0 : f32
- %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<32x32xf32>
- %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<32x32xf32>) {
- %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x32xbf16>
- %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x32xbf16>
- %d = vector.contract {
- indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction"],
- kind = #vector.kind<add>}
- %a, %b, %acc : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
- scf.yield %d : vector<32x32xf32>
+// The accumulator is loop-carried through an enclosing batch loop, so it is
+// neither read from nor written to memory directly.
+//
+// This test is currently intended to demonstrate successful lowering using
+// the fallback stack-buffer path. A more desirable lowering would keep the
+// accumulator on the tile registers across both loops.
+func.func @amx_int8_flat_batch_reduction(%A: memref<?x?x?xi8>, %B: memref<?x?x?xi8>, %C: memref<?x?xi32>,
+ %m: index, %n: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c64 = arith.constant 64 : index
+ %c0_i8 = arith.constant 0 : i8
+ %c0_i32 = arith.constant 0 : i32
+ %batch = memref.dim %A, %c0 : memref<?x?x?xi8>
+ %K = memref.dim %A, %c2 : memref<?x?x?xi8>
+ %c = vector.transfer_read %C[%m, %n], %c0_i32 {in_bounds = [true, true]} : memref<?x?xi32>, vector<64x128xi32>
+ %res = scf.for %b = %c0 to %batch step %c1 iter_args(%acc_b = %c) -> (vector<64x128xi32>) {
+ %res_b = scf.for %k = %c0 to %K step %c64 iter_args(%acc = %acc_b) -> (vector<64x128xi32>) {
+ %a = vector.transfer_read %A[%b, %m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x64xi8>
+ %bv = vector.transfer_read %B[%b, %k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?xi8>, vector<64x128xi8>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %bv, %acc : vector<64x64xi8>, vector<64x128xi8> into vector<64x128xi32>
+ scf.yield %d : vector<64x128xi32>
+ }
+ scf.yield %res_b : vector<64x128xi32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<32x32xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?xi32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_flat_single_reg_tile(
-// CHECK-SAME: %[[A:.+]]: memref<?x?xbf16>, %[[B:.+]]: memref<?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
-// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-LABEL: func.func @amx_int8_flat_batch_reduction(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x?xi8>, %[[B:.+]]: memref<?x?x?xi8>, %[[C:.+]]: memref<?x?xi32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[A]], %[[C0]]
+// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
// CHECK-NOT: memref.alloca
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [32, 32] [1, 1]
-// CHECK-NOT: memref.subview %[[ACC_VIEW]]
-// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[ACC2:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[ACC3:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C16]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C32]]
-// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]], %[[ARG2:.+]] = %[[ACC2]], %[[ARG3:.+]] = %[[ACC3]])
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]]] [32, 32] [1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[N]]] [32, 32] [1, 1]
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
-// CHECK: scf.yield
+// CHECK-NOT: memref.subview
+// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]]
+// CHECK: %[[ACC_RES:.+]] = scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]] iter_args(%[[ACC_B:.+]] = %[[ACC_INIT]])
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x128xi32>
+// CHECK: vector.transfer_write %[[ACC_B]], %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK-NOT: memref.alloca
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C128]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xi32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C64]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[OFF_M]], %[[IV_K]]] [1, 32, 64] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[OFF_N]]] [1, 64, 32] [1, 1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x64xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<64x16xi8>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<64x16xi8>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 64>} : vector<16x64xi8>, vector<64x16xi8> into vector<16x16xi32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-NOT: memref.subview %[[BUF]]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK: %[[ACC_NEXT:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x128xi32>
+// CHECK: scf.yield %[[ACC_NEXT]]
// CHECK: }
-// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
-// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
-// CHECK: vector.transfer_write %[[RES]]#2, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
-// CHECK: vector.transfer_write %[[RES]]#3, %[[ACC_VIEW]][%[[C16]], %[[C16]]]
-// CHECK-NOT: vector.transfer_read
-// CHECK-NOT: vector.transfer_write
-// CHECK: return
+// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]]
-// NANO-LABEL: func.func @amx_bf16_flat_single_reg_tile(
+// NANO-LABEL: func.func @amx_int8_flat_batch_reduction(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
} : !transform.any_op
transform.yield
}
@@ -1369,7 +1342,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_nano(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-bf16"
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll target = "amx-int8"
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -1378,64 +1351,91 @@ module attributes {transform.with_named_sequence} {
// -----
-#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
-#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
-#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-func.func @amx_bf16_vnni_single_reg_tile(%A: memref<?x?x2xbf16>, %B: memref<?x?x2xbf16>, %C: memref<?x?xf32>,
- %m: index, %n: index, %k_start: index, %k_end: index) {
+// Loop order K-batch, accumulating over both loops: the inner batch loop
+// carries the accumulator, its IV appears exactly once in each operand read and
+// its step is 1. It is a valid accumulation loop, indistinguishable from a
+// K-loop over a blocked layout, and is treated as such.
+//
+// Again, spilling the accumulator value to memory in the outer loop is correct,
+// but not ideal.
+func.func @amx_bf16_flat_k_batch_order(%A: memref<?x?x?xbf16>, %B: memref<?x?x?xbf16>, %C: memref<?x?xf32>,
+ %m: index, %n: index) {
%c0 = arith.constant 0 : index
- %c16 = arith.constant 16 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c32 = arith.constant 32 : index
%c0_0 = arith.constant 0.0 : bf16
%c0_0_f32 = arith.constant 0.0 : f32
- %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<32x32xf32>
- %res = scf.for %k = %k_start to %k_end step %c16 iter_args(%acc = %c) -> (vector<32x32xf32>) {
- %a = vector.transfer_read %A[%m, %k, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<32x16x2xbf16>
- %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<16x32x2xbf16>
- %d = vector.contract {
- indexing_maps = [#map, #map1, #map2],
- iterator_types = ["parallel", "parallel", "reduction", "reduction"],
- kind = #vector.kind<add>}
- %a, %b, %acc : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
- scf.yield %d : vector<32x32xf32>
+ %batch = memref.dim %A, %c0 : memref<?x?x?xbf16>
+ %K = memref.dim %A, %c2 : memref<?x?x?xbf16>
+ %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<64x96xf32>
+ %res = scf.for %k = %c0 to %K step %c32 iter_args(%acc_k = %c) -> (vector<64x96xf32>) {
+ %res_k = scf.for %bi = %c0 to %batch step %c1 iter_args(%acc = %acc_k) -> (vector<64x96xf32>) {
+ %a = vector.transfer_read %A[%bi, %m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?x?xbf16>, vector<64x32xbf16>
+ %b = vector.transfer_read %B[%bi, %k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?x?xbf16>, vector<32x96xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+ scf.yield %d : vector<64x96xf32>
+ }
+ scf.yield %res_k : vector<64x96xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<32x32xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x96xf32>, memref<?x?xf32>
func.return
}
-// CHECK-LABEL: func.func @amx_bf16_vnni_single_reg_tile(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x2xbf16>, %[[B:.+]]: memref<?x?x2xbf16>, %[[C:.+]]: memref<?x?xf32>,
-// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index, %[[K_START:.+]]: index, %[[K_END:.+]]: index)
+// CHECK-LABEL: func.func @amx_bf16_flat_k_batch_order(
+// CHECK-SAME: %[[A:.+]]: memref<?x?x?xbf16>, %[[B:.+]]: memref<?x?x?xbf16>, %[[C:.+]]: memref<?x?xf32>,
+// CHECK-SAME: %[[M:.+]]: index, %[[N:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C96:.+]] = arith.constant 96 : index
+// CHECK-DAG: %[[BATCH:.+]] = memref.dim %[[A]], %[[C0]]
+// CHECK-DAG: %[[K:.+]] = memref.dim %[[A]], %[[C2]]
// CHECK-NOT: memref.alloca
-// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[C]][%[[M]], %[[N]]] [32, 32] [1, 1]
-// CHECK-NOT: memref.subview %[[ACC_VIEW]]
-// CHECK: %[[ACC0:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[ACC1:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[ACC2:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[ACC3:.+]] = vector.transfer_read %[[ACC_VIEW]][%[[C16]], %[[C16]]], {{.*}} vector<16x16xf32>
-// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C16]]
-// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[ACC0]], %[[ARG1:.+]] = %[[ACC1]], %[[ARG2:.+]] = %[[ACC2]], %[[ARG3:.+]] = %[[ACC3]])
-// CHECK-NOT: arith.addi
-// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], 0] [32, 16, 2] [1, 1, 1]
-// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_K]], %[[N]], 0] [16, 32, 2] [1, 1, 1]
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]], %[[C0]]], {{.*}} vector<16x16x2xbf16>
-// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 16, 2>} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
-// CHECK: scf.yield
+// CHECK-NOT: memref.subview
+// CHECK: %[[ACC_INIT:.+]] = vector.transfer_read %[[C]][%[[M]], %[[N]]]
+// CHECK: %[[ACC_RES:.+]] = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C32]] iter_args(%[[ACC_K:.+]] = %[[ACC_INIT]])
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<64x96xf32>
+// CHECK: vector.transfer_write %[[ACC_K]], %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK-NOT: memref.alloca
+// CHECK: scf.for %[[IV_M:.+]] = %[[C0]] to %[[C64]] step %[[C32]] {
+// CHECK: scf.for %[[IV_N:.+]] = %[[C0]] to %[[C96]] step %[[C32]] {
+// CHECK: %[[ACC_VIEW:.+]] = memref.subview %[[BUF]][%[[IV_M]], %[[IV_N]]] [32, 32] [1, 1]
+// CHECK-COUNT-4: vector.transfer_read %[[ACC_VIEW]]{{.*}} vector<16x16xf32>
+// CHECK: %[[OFF_M:.+]] = arith.addi %[[M]], %[[IV_M]]
+// CHECK: %[[OFF_N:.+]] = arith.addi %[[N]], %[[IV_N]]
+// CHECK: %[[RES:.+]]:4 = scf.for %[[IV_B:.+]] = %[[C0]] to %[[BATCH]] step %[[C1]]
+// CHECK-NOT: arith.addi
+// CHECK: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[OFF_M]], %[[IV_K]]] [1, 32, 32] [1, 1, 1]
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[IV_K]], %[[OFF_N]]] [1, 32, 32] [1, 1, 1]
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[A_VIEW]][%[[C16]], %[[C0]]], {{.*}} vector<16x32xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C0]]], {{.*}} vector<32x16xbf16>
+// CHECK: vector.transfer_read %[[B_VIEW]][%[[C0]], %[[C16]]], {{.*}} vector<32x16xbf16>
+// CHECK-COUNT-4: vector.contract {{.*}} {x86_vcmlu_native_shape = array<i64: 16, 16, 32>} : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+// CHECK: scf.yield
+// CHECK: }
+// CHECK-NOT: memref.subview %[[BUF]]
+// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
+// CHECK: }
+// CHECK: }
+// CHECK: %[[ACC_NEXT:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]], {{.*}} vector<64x96xf32>
+// CHECK: scf.yield %[[ACC_NEXT]]
// CHECK: }
-// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
-// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
-// CHECK: vector.transfer_write %[[RES]]#2, %[[ACC_VIEW]][%[[C16]], %[[C0]]]
-// CHECK: vector.transfer_write %[[RES]]#3, %[[ACC_VIEW]][%[[C16]], %[[C16]]]
-// CHECK-NOT: vector.transfer_read
-// CHECK-NOT: vector.transfer_write
-// CHECK: return
+// CHECK: vector.transfer_write %[[ACC_RES]], %[[C]][%[[M]], %[[N]]]
-// NANO-LABEL: func.func @amx_bf16_vnni_single_reg_tile(
+// NANO-LABEL: func.func @amx_bf16_flat_k_batch_order(
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
>From ad0e1b500e930fba102ba439ba004ddfa0047719 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Fri, 9 Oct 2026 01:06:32 -0700
Subject: [PATCH 15/15] Add target enum
---
mlir/include/mlir/Dialect/X86/CMakeLists.txt | 7 +++-
.../X86/TransformOps/X86TransformOps.td | 1 +
mlir/include/mlir/Dialect/X86/Transforms.h | 3 +-
mlir/include/mlir/Dialect/X86/X86Dialect.h | 1 +
mlir/include/mlir/Dialect/X86/X86Enums.h | 16 +++++++++
mlir/include/mlir/Dialect/X86/X86Enums.td | 34 ++++++++++++++++++
mlir/lib/Dialect/X86/IR/CMakeLists.txt | 1 +
mlir/lib/Dialect/X86/IR/X86Dialect.cpp | 2 ++
.../X86/TransformOps/X86TransformOps.cpp | 10 +++++-
.../VectorContractMultiLevelUnroll.cpp | 35 +++++++++++--------
10 files changed, 92 insertions(+), 18 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/X86/X86Enums.h
create mode 100644 mlir/include/mlir/Dialect/X86/X86Enums.td
diff --git a/mlir/include/mlir/Dialect/X86/CMakeLists.txt b/mlir/include/mlir/Dialect/X86/CMakeLists.txt
index 190573e043b173..a16d545e5020f4 100644
--- a/mlir/include/mlir/Dialect/X86/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/X86/CMakeLists.txt
@@ -1,7 +1,12 @@
add_mlir_dialect(X86 x86)
add_mlir_doc(X86 X86 Dialects/ -gen-dialect-doc -dialect=x86)
+set(LLVM_TARGET_DEFINITIONS X86Enums.td)
+mlir_tablegen(X86Enums.h.inc -gen-enum-decls)
+mlir_tablegen(X86Enums.cpp.inc -gen-enum-defs)
+add_mlir_dialect_tablegen_target(MLIRX86EnumsIncGen)
+
add_mlir_interface(X86Interfaces)
-add_dependencies(MLIRX86IncGen MLIRX86InterfacesIncGen)
+add_dependencies(MLIRX86IncGen MLIRX86InterfacesIncGen MLIRX86EnumsIncGen)
add_subdirectory(TransformOps)
diff --git a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
index a1a6f99a9447ca..5f55d97ec94109 100644
--- a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
+++ b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
@@ -104,6 +104,7 @@ def ApplyVectorContractMultiLevelUnrollPatternsOp : Op<Transform_Dialect,
let arguments = (ins StrAttr:$target);
let assemblyFormat = "`target` `=` $target attr-dict";
+ let hasVerifier = 1;
}
#endif // X86_TRANSFORM_OPS
diff --git a/mlir/include/mlir/Dialect/X86/Transforms.h b/mlir/include/mlir/Dialect/X86/Transforms.h
index 676b2f85f2c0c5..05c9ba2c1b6d19 100644
--- a/mlir/include/mlir/Dialect/X86/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86/Transforms.h
@@ -9,6 +9,7 @@
#ifndef MLIR_DIALECT_X86_TRANSFORMS_H
#define MLIR_DIALECT_X86_TRANSFORMS_H
+#include "mlir/Dialect/X86/X86Enums.h"
#include "mlir/IR/Value.h"
namespace mlir {
@@ -118,7 +119,7 @@ void populateMoveAccumulatorForContractLoopPatterns(
// A set of patterns that tile and unroll contractions in order to prepare them
// for the nanokernel lowerings.
void populateVectorContractMultiLevelUnrollPatterns(RewritePatternSet &patterns,
- StringRef target);
+ MLUTarget target);
//===----------------------------------------------------------------------===//
/// Helpers extracted from:
diff --git a/mlir/include/mlir/Dialect/X86/X86Dialect.h b/mlir/include/mlir/Dialect/X86/X86Dialect.h
index 6b1358b31e666d..c9643cfc71c52a 100644
--- a/mlir/include/mlir/Dialect/X86/X86Dialect.h
+++ b/mlir/include/mlir/Dialect/X86/X86Dialect.h
@@ -16,6 +16,7 @@
#include "mlir/Bytecode/BytecodeOpInterface.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/X86/X86Enums.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
diff --git a/mlir/include/mlir/Dialect/X86/X86Enums.h b/mlir/include/mlir/Dialect/X86/X86Enums.h
new file mode 100644
index 00000000000000..c65f00531920dd
--- /dev/null
+++ b/mlir/include/mlir/Dialect/X86/X86Enums.h
@@ -0,0 +1,16 @@
+//===- X86Enums.h - X86 Dialect Enums ---------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_X86_X86ENUMS_H
+#define MLIR_DIALECT_X86_X86ENUMS_H
+
+#include "mlir/IR/Dialect.h"
+
+#include "mlir/Dialect/X86/X86Enums.h.inc"
+
+#endif // MLIR_DIALECT_X86_X86ENUMS_H
diff --git a/mlir/include/mlir/Dialect/X86/X86Enums.td b/mlir/include/mlir/Dialect/X86/X86Enums.td
new file mode 100644
index 00000000000000..2fa059509db068
--- /dev/null
+++ b/mlir/include/mlir/Dialect/X86/X86Enums.td
@@ -0,0 +1,34 @@
+//===-- X86Enums.td - X86 dialect enum definitions ---------*- tablegen -*-===//
+//
+// 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 defines the enums of the X86 dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef X86_ENUMS
+#define X86_ENUMS
+
+include "mlir/IR/EnumAttr.td"
+
+//===----------------------------------------------------------------------===//
+// MLUTarget
+//===----------------------------------------------------------------------===//
+
+// The target CPU feature for the vector.contract multi-level unroll pattern
+// determines native shapes and supported data types for operands and the
+// accumulator, and corresponds to one of the nanokernel patterns.
+def MLUTarget : I32EnumAttr<"MLUTarget",
+ "Target CPU feature for the vector.contract multi-level unroll pattern", [
+ I32EnumAttrCase<"AMX_BF16", 0, "amx-bf16">, //-> VectorContractToAMXDotProduct
+ I32EnumAttrCase<"AMX_INT8", 1, "amx-int8">, //-> VectorContractToAMXDotProduct
+]> {
+ let cppNamespace = "::mlir::x86";
+ let genSpecializedAttr = 0;
+}
+
+#endif // X86_ENUMS
diff --git a/mlir/lib/Dialect/X86/IR/CMakeLists.txt b/mlir/lib/Dialect/X86/IR/CMakeLists.txt
index 442e419e6b289d..85c87c31361b57 100644
--- a/mlir/lib/Dialect/X86/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86/IR/CMakeLists.txt
@@ -7,6 +7,7 @@ add_mlir_dialect_library(MLIRX86Dialect
DEPENDS
MLIRX86IncGen
MLIRX86InterfacesIncGen
+ MLIRX86EnumsIncGen
LINK_LIBS PUBLIC
MLIRIR
diff --git a/mlir/lib/Dialect/X86/IR/X86Dialect.cpp b/mlir/lib/Dialect/X86/IR/X86Dialect.cpp
index 2ef061a27ddad6..1c435c2f248694 100644
--- a/mlir/lib/Dialect/X86/IR/X86Dialect.cpp
+++ b/mlir/lib/Dialect/X86/IR/X86Dialect.cpp
@@ -27,6 +27,8 @@ using namespace mlir;
#include "mlir/Dialect/X86/X86Dialect.cpp.inc"
+#include "mlir/Dialect/X86/X86Enums.cpp.inc"
+
void x86::X86Dialect::initialize() {
addTypes<
#define GET_TYPEDEF_LIST
diff --git a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
index ff05e38913d335..ec37851e9ad59f 100644
--- a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
+++ b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
@@ -59,7 +59,15 @@ void mlir::transform::ApplyMoveAccumulatorForContractLoopPatternsOp::
void mlir::transform::ApplyVectorContractMultiLevelUnrollPatternsOp::
populatePatterns(RewritePatternSet &patterns) {
- x86::populateVectorContractMultiLevelUnrollPatterns(patterns, getTarget());
+ x86::populateVectorContractMultiLevelUnrollPatterns(
+ patterns, *x86::symbolizeMLUTarget(getTarget()));
+}
+
+LogicalResult
+mlir::transform::ApplyVectorContractMultiLevelUnrollPatternsOp::verify() {
+ if (!x86::symbolizeMLUTarget(getTarget()))
+ return emitError("unsupported target '") << getTarget() << "'";
+ return success();
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 02c4ae1add2662..e450bc20bf9dc2 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -31,11 +31,14 @@ static constexpr auto nativeShapeAttrName = "x86_vcmlu_native_shape";
namespace {
// Shared state for the pattern match and IR transformation.
struct MLUCandidate {
- // The CPU feature flag to target with the transformation, in the spelling
- // used by llvm/include/llvm/TargetParser/X86TargetParser.def, e.g.
- // "amx-bf16". Determines native shapes and supported data types for operands
- // and the accumulator.
- StringRef target;
+ // The target determines native shapes and supported data types for operands
+ // and the accumulator, and corresponds to one of the nanokernel patterns.
+ //
+ // For example, x86::MLUTarget::AMX_INT8 instructs this pattern to tile and
+ // unroll the contraction in such a way that the VectorContractToAMXDotProduct
+ // pattern can be applied (which then will lower the smaller contraction(s) to
+ // sequences of x86.amx.tile_{load,muli,store} operations).
+ x86::MLUTarget target;
// Anchor operation.
vector::ContractionOp contract;
@@ -315,7 +318,7 @@ static LogicalResult matchCanonicalMatmul(MLUCandidate &candidate,
static LogicalResult matchShapesAndTypes(MLUCandidate &candidate,
PatternRewriter &rewriter) {
vector::ContractionOp contract = candidate.contract;
- StringRef target = candidate.target;
+ x86::MLUTarget target = candidate.target;
Type accElemType = candidate.accType.getElementType();
Type inpElemType = candidate.lhsType.getElementType();
@@ -331,10 +334,12 @@ static LogicalResult matchShapesAndTypes(MLUCandidate &candidate,
int64_t regTileM, regTileN, regTileK;
int64_t nativeM, nativeN, nativeK, vnni;
- if (target.starts_with("amx")) {
- if (!((target == "amx-bf16" && inpElemType.isBF16() &&
+ switch (target) {
+ case x86::MLUTarget::AMX_BF16:
+ case x86::MLUTarget::AMX_INT8: {
+ if (!((target == x86::MLUTarget::AMX_BF16 && inpElemType.isBF16() &&
accElemType.isF32()) ||
- (target == "amx-int8" && inpElemType.isInteger(8) &&
+ (target == x86::MLUTarget::AMX_INT8 && inpElemType.isInteger(8) &&
accElemType.isInteger(32))))
return rewriter.notifyMatchFailure(
contract, "unsupported combination of input and accumulator types");
@@ -369,8 +374,8 @@ static LogicalResult matchShapesAndTypes(MLUCandidate &candidate,
regTileN = 2 * nativeN;
regTileK = nativeK;
}
- } else {
- return rewriter.notifyMatchFailure(contract, "unsupported target");
+ break;
+ }
}
candidate.accRegTileType = VectorType::get({regTileM, regTileN}, accElemType);
@@ -789,10 +794,10 @@ namespace {
// transforming it in place.
struct VectorContractMultiLevelUnroll
: public OpRewritePattern<vector::ContractionOp> {
- VectorContractMultiLevelUnroll(MLIRContext *context, StringRef target,
+ VectorContractMultiLevelUnroll(MLIRContext *context, x86::MLUTarget target,
PatternBenefit benefit = 1)
: OpRewritePattern<vector::ContractionOp>(context, benefit),
- target(target.str()) {}
+ target(target) {}
LogicalResult matchAndRewrite(vector::ContractionOp contract,
PatternRewriter &rewriter) const override {
@@ -820,12 +825,12 @@ struct VectorContractMultiLevelUnroll
}
private:
- std::string target;
+ x86::MLUTarget target;
};
} // namespace
void x86::populateVectorContractMultiLevelUnrollPatterns(
- RewritePatternSet &patterns, StringRef target) {
+ RewritePatternSet &patterns, MLUTarget target) {
patterns.add<VectorContractMultiLevelUnroll>(patterns.getContext(), target);
vector::UnrollVectorOptions options;
options.setFilterConstraint(hasNativeShapeAttr);
More information about the Mlir-commits
mailing list