[Mlir-commits] [mlir] [mlir][x86] Support multi-level unroll of contractions with non-native K extent (PR #234593)
Julian Oppermann
llvmlistbot at llvm.org
Wed Oct 28 05:49:34 PDT 2026
https://github.com/jopperm updated https://github.com/llvm/llvm-project/pull/234593
>From 7e6cd8080a81d2867f2610e8ba60ff8f9610e667 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 21 Oct 2026 12:26:31 -0700
Subject: [PATCH 1/4] [mlir][x86] Support multi-level unroll of contractions
with non-native K extent
---
.../VectorContractMultiLevelUnroll.cpp | 235 +++++++++---
...r-contract-multi-level-unroll-for-amx.mlir | 358 +++++++++++-------
2 files changed, 394 insertions(+), 199 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 5a7ad285446b34..94168aa54a8570 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -23,7 +23,6 @@
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";
@@ -39,6 +38,7 @@ struct MLUCandidate {
VectorType accType;
VectorType lhsType;
VectorType rhsType;
+ int64_t contractK;
// Shape of the register-tiled contraction. We'll insert a loop nest to cover
// the original shape.
@@ -61,12 +61,19 @@ struct MLUCandidate {
vector::TransferReadOp lhsRead;
vector::TransferReadOp rhsRead;
+ // Distinguishes whether the IV of the accumulation loop is used to index a
+ // vector dimension ("offset") or a non-vector dimension ("block") on the
+ // operand transfer reads.
+ enum class IndexKind { Unknown, Offset, Block };
+ IndexKind operandIndexKind = IndexKind::Unknown;
+
// Memory backing the accumulator.
Value accMemIn;
Value accMemOut;
- // Reconstructed accumulation loop with native-shaped iter_args.
+ // Reconstructed accumulation loop(s) with native-shaped iter_args.
scf::ForOp slicedAccLoop;
+ scf::ForOp slicedInnerAccLoop;
// Keep track of operations that should be deleted after transformation
// (depending on which ops could be integrated into the new loop nest).
@@ -90,9 +97,10 @@ using MatchTypesAndShapeFn = std::function<LogicalResult(
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);
+static SmallVector<Value> buildKLLoopBody(MLUCandidate &candidate, OpBuilder &b,
+ Location loc, Value ivM, Value ivN,
+ Value ivK, Value ivL,
+ ValueRange iterArgs);
//----------------------------------------------------------------------------//
// Utilities //
@@ -254,7 +262,8 @@ static void setNativeShapeAttr(VectorTransferOpInterface op,
// Establish whether the candidate contraction is a canonical matrix
// multiplication on flat or VNNI-packed vectors. If successful, the discovered
-// operand and accumulator types are recorded in the candidate struct.
+// operand/accumulator types and the extent of the main reduction dimension are
+// recorded in the candidate struct.
//
// Example (flat):
// vector.contract {
@@ -329,6 +338,7 @@ static LogicalResult matchCanonicalMatmul(MLUCandidate &candidate,
candidate.accType = accType;
candidate.lhsType = lhsType;
candidate.rhsType = rhsType;
+ candidate.contractK = lhsType.getShape()[1];
return success();
}
@@ -379,10 +389,13 @@ matchTypesAndShape(MLUCandidate &candidate,
return rewriter.notifyMatchFailure(
contract, "contraction is already in native shape");
- if (!all_of(drop_begin(*maybeTileShapeRatio, 2), isOne) ||
- !all_of(drop_begin(*maybeUnrollShapeRatio, 2), isOne))
+ if (!all_of(drop_begin(*maybeUnrollShapeRatio, 2), isOne))
return rewriter.notifyMatchFailure(
- contract, "cannot tile/unroll on the reduction dimensions");
+ contract, "cannot unroll in the reduction dimensions");
+
+ if (candidate.isInVnniLayout && (maybeTileShapeRatio->back() != 1))
+ return rewriter.notifyMatchFailure(contract,
+ "cannot tile in the VNNI dimension");
int64_t regTileM = contractRegTileShape[0];
int64_t regTileN = contractRegTileShape[1];
@@ -417,20 +430,23 @@ matchTypesAndShape(MLUCandidate &candidate,
}
// 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-loop with a constant step and 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 accumulation loops 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 one of their indices. This covers plain
-// 2D memrefs as well as block-based layouts.
+// use the induction variable as exactly one of their indices. We check the
+// loop's step based on the index kind of the reads: IndexKind::Offset requires
+// a step equal to the contraction's K extent, whereas IndexKind::Block requires
+// unit steps.
//
// If all these conditions are met, the loop op, initial value of the
-// accumulator and the operand reads are recorded in the candidate struct.
+// accumulator and the operand reads (with their index kind) are recorded in the
+// candidate struct.
//
-// Example:
+// Example (IndexKind::Block):
// !accType = vector<64x128xf32>
// !lhsType = vector<64x32xbf16>
// !rhsType = vector<32x128xbf16>
@@ -454,11 +470,6 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return rewriter.notifyMatchFailure(contract,
"is not in an index-based for-loop");
- if (auto step = getConstantIntValue(accLoop.getStep());
- !step || (*step != 1 && *step != candidate.contractNativeShape[2]))
- 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,
@@ -487,22 +498,78 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return rewriter.notifyMatchFailure(
contract, "LHS/RHS are not vector.transfer_read ops");
- auto checkTransferOp = [&](vector::TransferReadOp transferOp) {
+ auto matchOperandRead = [&](vector::TransferReadOp transferOp, bool isLhs) {
+ using IK = MLUCandidate::IndexKind;
+
if (!isSimpleTransferOp(transferOp))
- return rewriter.notifyMatchFailure(contract, [&](Diagnostic &diag) {
- diag << "transfer op has non-identity permutation, mask, or "
- "out-of-bounds check: "
- << *transferOp;
- });
+ return rewriter.notifyMatchFailure(
+ contract, "operand read has non-identity permutation, mask, or "
+ "out-of-bounds check");
+
+ auto indices = transferOp.getIndices();
+ auto ivIt =
+ std::find(indices.begin(), indices.end(), accLoop.getInductionVar());
+ if (ivIt == indices.end())
+ return rewriter.notifyMatchFailure(
+ contract, "operand read does not use loop index");
+ if (std::find(ivIt + 1, indices.end(), accLoop.getInductionVar()) !=
+ indices.end())
+ return rewriter.notifyMatchFailure(
+ contract, "operand read uses loop index more than once");
+
+ unsigned firstVecDim =
+ indices.size() - transferOp.getVectorType().getRank();
+ unsigned ivPos = std::distance(indices.begin(), ivIt);
+ unsigned kPos = firstVecDim + isLhs;
+ if (ivPos >= firstVecDim) {
+ if (ivPos != kPos)
+ return rewriter.notifyMatchFailure(
+ contract,
+ "loop index is not in expected K position in operand read");
+ if (candidate.operandIndexKind != IK::Unknown &&
+ candidate.operandIndexKind != IK::Offset)
+ return rewriter.notifyMatchFailure(contract,
+ "inconsistent operand index kind");
+ candidate.operandIndexKind = IK::Offset;
+ } else {
+ // Temporary limitation: having a zero index in the K position means that
+ // we can apply an offset without introducing additional operations in the
+ // accumulation loop.
+ if (!isZeroInteger(indices[kPos]))
+ return rewriter.notifyMatchFailure(
+ contract, "operand read K index must be zero for block index kind");
+ if (candidate.operandIndexKind != IK::Unknown &&
+ candidate.operandIndexKind != IK::Block)
+ return rewriter.notifyMatchFailure(contract,
+ "inconsistent operand index kind");
+ candidate.operandIndexKind = IK::Block;
+ }
+ assert(candidate.operandIndexKind != IK::Unknown &&
+ "operand index kind must be determined");
+
+ if (candidate.isInVnniLayout && !isZeroInteger(indices.back()))
+ return rewriter.notifyMatchFailure(
+ contract, "operand read VNNI dimension offset must be zero");
- if (count(transferOp.getIndices(), accLoop.getInductionVar()) != 1)
+ auto maybeStep = getConstantIntValue(accLoop.getStep());
+ if (!maybeStep)
+ return rewriter.notifyMatchFailure(
+ contract, "accumulation loop step is not constant");
+
+ if (candidate.operandIndexKind == IK::Offset &&
+ *maybeStep != candidate.contractK)
+ return rewriter.notifyMatchFailure(
+ contract,
+ "loop step must match contraction's K extent for offset indexing");
+ if (candidate.operandIndexKind == IK::Block && *maybeStep != 1)
return rewriter.notifyMatchFailure(
- contract, "transfer op does not use accumulator index");
+ contract, "loop step must be 1 for block index kind");
return success();
};
- if (failed(checkTransferOp(lhsRead)) || failed(checkTransferOp(rhsRead)))
+ if (failed(matchOperandRead(lhsRead, /*isLhs=*/true)) ||
+ failed(matchOperandRead(rhsRead, /*isLhs=*/false)))
return failure();
candidate.accLoop = accLoop;
@@ -676,8 +743,12 @@ static void buildLoopNestForCandidate(MLUCandidate &candidate,
// 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
+// i.e. the iter_args have accNativeType. Tiling along the K-dimension
+// depends on the operand reads' index kind:
+// IndexKind::Offset: Use the same bounds but smaller steps/more iterations.
+// IndexKind::Block: Nest an additional L-loop to move over the block in
+// steps.
+// 4. -> buildKLLoopBody
// 5. Splice the final value slices back into an accRegTileType-shaped value and
// write it back to memory.
//
@@ -685,6 +756,8 @@ static void buildLoopNestForCandidate(MLUCandidate &candidate,
// to unroll them with the vector-dialect's unroll patterns.
static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
Value ivM, Value ivN) {
+ using IK = MLUCandidate::IndexKind;
+
ArrayRef<int64_t> accRegTileShape = candidate.accRegTileType.getShape();
Value c0 = arith::ConstantIndexOp::create(b, loc, 0);
@@ -713,14 +786,36 @@ static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
SmallVector<Value> accInitSliced =
slice(b, loc, accInit->getResult(0), candidate.accNativeType);
+
+ Value lbK = candidate.accLoop.getLowerBound();
+ Value ubK = candidate.accLoop.getUpperBound();
+ // Note: Using _native_ shape here as we never unroll along the reduction
+ // dimension, so the K extents of the register-tile and native shapes match.
+ int64_t stepKInt = candidate.operandIndexKind == IK::Offset
+ ? candidate.contractNativeShape[2]
+ : 1;
+ Value stepK = arith::ConstantIndexOp::create(b, loc, stepKInt);
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);
+ b, loc, lbK, ubK, stepK, accInitSliced,
+ [&](OpBuilder &b, Location loc, Value ivK, ValueRange iterArgsK) {
+ if (candidate.operandIndexKind == IK::Offset) {
+ SmallVector<Value> newIterArgs = buildKLLoopBody(
+ candidate, b, loc, ivM, ivN, ivK, /*ivL=*/Value{}, iterArgsK);
+ scf::YieldOp::create(b, loc, newIterArgs);
+ return;
+ }
+
+ Value ubL = arith::ConstantIndexOp::create(b, loc, candidate.contractK);
+ Value stepL = arith::ConstantIndexOp::create(
+ b, loc, candidate.contractNativeShape[2]);
+ candidate.slicedInnerAccLoop = scf::ForOp::create(
+ b, loc, c0, ubL, stepL, iterArgsK,
+ [&](OpBuilder &b, Location loc, Value ivL, ValueRange iterArgsL) {
+ SmallVector<Value> newIterArgs = buildKLLoopBody(
+ candidate, b, loc, ivM, ivN, ivK, ivL, iterArgsL);
+ scf::YieldOp::create(b, loc, newIterArgs);
+ });
+ scf::YieldOp::create(b, loc, candidate.slicedInnerAccLoop.getResults());
});
Value accSpliced = splice(b, loc, candidate.slicedAccLoop.getResults(),
@@ -735,29 +830,41 @@ static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
}
// 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) {
+// replacing the induction variable and taking the offsets from the
+// M-N-(L)-loop nest 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> buildKLLoopBody(MLUCandidate &candidate, OpBuilder &b,
+ Location loc, Value ivM, Value ivN,
+ Value ivK, Value ivL,
+ ValueRange iterArgs) {
Value accSpliced = splice(b, loc, iterArgs, candidate.accRegTileType);
- SmallVector<Value, 4> lhsOffsets = candidate.lhsRead.getIndices();
- SmallVector<Value, 4> rhsOffsets = candidate.rhsRead.getIndices();
+ SmallVector<Value, 4> lhsIndices = candidate.lhsRead.getIndices();
+ SmallVector<Value, 4> rhsIndices = 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);
+ replace(lhsIndices, accLoopIV, ivK);
+ replace(rhsIndices, accLoopIV, ivK);
+
+ unsigned mPos = lhsIndices.size() - 2 - candidate.isInVnniLayout;
+ unsigned nPos = rhsIndices.size() - 1 - candidate.isInVnniLayout;
+ lhsIndices[mPos] = arith::AddIOp::create(b, loc, lhsIndices[mPos], ivM);
+ rhsIndices[nPos] = arith::AddIOp::create(b, loc, rhsIndices[nPos], ivN);
+
+ if (ivL) {
+ unsigned lPosL = lhsIndices.size() - 1 - candidate.isInVnniLayout;
+ unsigned lPosR = rhsIndices.size() - 2 - candidate.isInVnniLayout;
+ lhsIndices[lPosL] = arith::AddIOp::create(b, loc, lhsIndices[lPosL], ivL);
+ rhsIndices[lPosR] = arith::AddIOp::create(b, loc, rhsIndices[lPosR], ivL);
+ }
+
Value lhsView = makeRankReducingSubview(b, loc, candidate.lhsRead.getBase(),
- lhsOffsets, candidate.lhsRegTileType);
+ lhsIndices, candidate.lhsRegTileType);
Value rhsView = makeRankReducingSubview(b, loc, candidate.rhsRead.getBase(),
- rhsOffsets, candidate.rhsRegTileType);
+ rhsIndices, candidate.rhsRegTileType);
SmallVector<Value, 3> zeroIndices(2 + candidate.isInVnniLayout,
arith::ConstantIndexOp::create(b, loc, 0));
@@ -790,12 +897,12 @@ 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
+// The pattern determines a register-tile shape 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.
+// target's vector registers or AMX tile registers. It then introduces a 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 helpers above):
// The contraction must be a canonical matrix multiplication, have
@@ -830,6 +937,8 @@ struct VectorContractMultiLevelUnroll
makeAccInitializationBackedByMemory(candidate, rewriter);
makeAccResultBackedByMemory(candidate, rewriter);
buildLoopNestForCandidate(candidate, rewriter);
+ if (candidate.slicedInnerAccLoop)
+ moveLoopInvariantCode(candidate.slicedInnerAccLoop);
moveLoopInvariantCode(candidate.slicedAccLoop);
for (auto *op : candidate.opsToDelete)
rewriter.eraseOp(op);
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
index 48f0892ab957db..0eadcf1a114177 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
@@ -8,18 +8,18 @@
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
+ %c64 = arith.constant 64 : 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>
+ %res = scf.for %k = %k_start to %k_end step %c64 iter_args(%acc = %c) -> (vector<128x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<128x64xbf16>
+ %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<64x64xbf16>
%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>
+ %a, %b, %acc : vector<128x64xbf16>, vector<64x64xbf16> into vector<128x64xf32>
scf.yield %d : vector<128x64xf32>
}
vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<128x64xf32>, memref<?x?xf32>
@@ -97,22 +97,22 @@ module attributes {transform.with_named_sequence} {
#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.
+// M = 48 is an odd multiple of 16, so a 1x2 grid of tile registers 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) {
+ %c48 = arith.constant 48 : 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>
+ %res = scf.for %k = %k_start to %k_end step %c48 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<48x48x2xbf16>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x2xbf16>, vector<48x64x2xbf16>
%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>
+ %a, %b, %acc : vector<48x48x2xbf16>, vector<48x64x2xbf16> into vector<48x64xf32>
scf.yield %d : vector<48x64xf32>
}
vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x64xf32>, memref<?x?xf32>
@@ -185,7 +185,7 @@ module attributes {transform.with_named_sequence} {
// 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>,
+func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x128xbf16>, %B: memref<?x?x128x64xbf16>, %C: memref<?x?x96x64xf32>,
%m: index, %n: index, %k_start: index, %k_end: index) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
@@ -193,13 +193,13 @@ func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x32xbf16>, %B: memref<?x?x32x6
%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>
+ %a = vector.transfer_read %A[%m, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x96x128xbf16>, vector<96x128xbf16>
+ %b = vector.transfer_read %B[%n, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x128x64xbf16>, vector<128x64xbf16>
%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>
+ %a, %b, %acc : vector<96x128xbf16>, vector<128x64xbf16> 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>
@@ -207,7 +207,7 @@ func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x32xbf16>, %B: memref<?x?x32x6
}
// CHECK-LABEL: func.func @amx_bf16_flat_blocked(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x96x32xbf16>, %[[B:.+]]: memref<?x?x32x64xbf16>, %[[C:.+]]: memref<?x?x96x64xf32>,
+// CHECK-SAME: %[[A:.+]]: memref<?x?x96x128xbf16>, %[[B:.+]]: memref<?x?x128x64xbf16>, %[[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
@@ -215,6 +215,7 @@ func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x32xbf16>, %B: memref<?x?x32x6
// 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-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: ?>>
@@ -223,11 +224,14 @@ func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x32xbf16>, %B: memref<?x?x32x6
// 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-SAME: iter_args(%[[K0:.+]] = %{{.+}}, %[[K1:.+]] = %{{.+}}, %[[K2:.+]] = %{{.+}}, %[[K3:.+]] = %{{.+}})
+// CHECK: %[[INNER:.+]]:4 = scf.for %[[IV_L:.+]] = %[[C0]] to %[[C128]] step %[[C32]]
+// CHECK-SAME: iter_args(%{{.+}} = %[[K0]], %{{.+}} = %[[K1]], %{{.+}} = %[[K2]], %{{.+}} = %[[K3]])
// 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: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], %[[IV_L]]] [1, 1, 32, 32] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x96x128xbf16> to memref<32x32xbf16, strided<[128, 1], offset: ?>>
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], %[[IV_L]], %[[IV_N]]] [1, 1, 32, 32] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x128x64xbf16> 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>
@@ -235,6 +239,8 @@ func.func @amx_bf16_flat_blocked(%A: memref<?x?x96x32xbf16>, %B: memref<?x?x32x6
// 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: scf.yield %[[INNER]]#0, %[[INNER]]#1, %[[INNER]]#2, %[[INNER]]#3
+// CHECK: }
// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
// CHECK: }
@@ -269,7 +275,7 @@ module attributes {transform.with_named_sequence} {
#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>,
+func.func @amx_bf16_vnni_blocked(%A: memref<?x?x128x32x2xbf16>, %B: memref<?x?x32x96x2xbf16>, %C: memref<?x?x128x96xf32>,
%m: index, %n: index, %k_start: index, %k_end: index) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
@@ -277,13 +283,13 @@ func.func @amx_bf16_vnni_blocked(%A: memref<?x?x128x16x2xbf16>, %B: memref<?x?x1
%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>
+ %a = vector.transfer_read %A[%m, %k, %c0, %c0, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x128x32x2xbf16>, vector<128x32x2xbf16>
+ %b = vector.transfer_read %B[%n, %k, %c0, %c0, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x32x96x2xbf16>, vector<32x96x2xbf16>
%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>
+ %a, %b, %acc : vector<128x32x2xbf16>, vector<32x96x2xbf16> 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>
@@ -291,7 +297,7 @@ func.func @amx_bf16_vnni_blocked(%A: memref<?x?x128x16x2xbf16>, %B: memref<?x?x1
}
// CHECK-LABEL: func.func @amx_bf16_vnni_blocked(
-// CHECK-SAME: %[[A:.+]]: memref<?x?x128x16x2xbf16>, %[[B:.+]]: memref<?x?x16x96x2xbf16>, %[[C:.+]]: memref<?x?x128x96xf32>,
+// CHECK-SAME: %[[A:.+]]: memref<?x?x128x32x2xbf16>, %[[B:.+]]: memref<?x?x32x96x2xbf16>, %[[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
@@ -307,11 +313,14 @@ func.func @amx_bf16_vnni_blocked(%A: memref<?x?x128x16x2xbf16>, %B: memref<?x?x1
// 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-SAME: iter_args(%[[K0:.+]] = %{{.+}}, %[[K1:.+]] = %{{.+}}, %[[K2:.+]] = %{{.+}}, %[[K3:.+]] = %{{.+}})
+// CHECK: %[[INNER:.+]]:4 = scf.for %[[IV_L:.+]] = %[[C0]] to %[[C32]] step %[[C16]]
+// CHECK-SAME: iter_args(%{{.+}} = %[[K0]], %{{.+}} = %[[K1]], %{{.+}} = %[[K2]], %{{.+}} = %[[K3]])
// 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: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], %[[IV_L]], 0] [1, 1, 32, 16, 2] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x128x32x2xbf16> to memref<32x16x2xbf16, strided<[64, 2, 1], offset: ?>>
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], %[[IV_L]], %[[IV_N]], 0] [1, 1, 16, 32, 2] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x32x96x2xbf16> 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>
@@ -319,6 +328,8 @@ func.func @amx_bf16_vnni_blocked(%A: memref<?x?x128x16x2xbf16>, %B: memref<?x?x1
// 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: scf.yield %[[INNER]]#0, %[[INNER]]#1, %[[INNER]]#2, %[[INNER]]#3
+// CHECK: }
// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
// CHECK: }
@@ -355,18 +366,18 @@ module attributes {transform.with_named_sequence} {
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
+ %c192 = arith.constant 192 : 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>
+ %res = scf.for %k = %k_start to %k_end step %c192 iter_args(%acc = %c) -> (vector<64x128xi32>) {
+ %a = vector.transfer_read %A[%m, %k], %c0_i8 {in_bounds = [true, true]} : memref<?x?xi8>, vector<64x192xi8>
+ %b = vector.transfer_read %B[%k, %n], %c0_i8 {in_bounds = [true, true]} : memref<?x?xi8>, vector<192x128xi8>
%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>
+ %a, %b, %acc : vector<64x192xi8>, vector<192x128xi8> into vector<64x128xi32>
scf.yield %d : vector<64x128xi32>
}
vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?xi32>
@@ -434,22 +445,22 @@ module attributes {transform.with_named_sequence} {
#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.
+// N = 48 is an odd multiple of 16, so a 2x1 grid of tile registers 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) {
+ %c64 = arith.constant 64 : 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>
+ %res = scf.for %k = %k_start to %k_end step %c64 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<64x64x4xi8>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_i8 {in_bounds = [true, true, true]} : memref<?x?x4xi8>, vector<64x48x4xi8>
%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>
+ %a, %b, %acc : vector<64x64x4xi8>, vector<64x48x4xi8> into vector<64x48xi32>
scf.yield %d : vector<64x48xi32>
}
vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<64x48xi32>, memref<?x?xi32>
@@ -1042,51 +1053,56 @@ 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.
+// M = 48 is an odd multiple of 16, so a 1x2 grid of tile registers 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>,
+func.func @amx_bf16_flat_1x2_tile(%A: memref<?x?x48x96xbf16>, %B: memref<?x?x96x64xbf16>, %C: memref<?x?x48x64xf32>,
%m: index, %n: index, %k_start: index, %k_end: index) {
- %c32 = arith.constant 32 : 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<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>
+ %c = vector.transfer_read %C[%m, %n, %c0, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?x48x64xf32>, vector<48x64xf32>
+ %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c) -> (vector<48x64xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x48x96xbf16>, vector<48x96xbf16>
+ %b = vector.transfer_read %B[%n, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x96x64xbf16>, vector<96x64xbf16>
%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>
+ %a, %b, %acc : vector<48x96xbf16>, vector<96x64xbf16> into vector<48x64xf32>
scf.yield %d : vector<48x64xf32>
}
- vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<48x64xf32>, memref<?x?xf32>
+ vector.transfer_write %res, %C[%m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<48x64xf32>, memref<?x?x48x64xf32>
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: %[[A:.+]]: memref<?x?x48x96xbf16>, %[[B:.+]]: memref<?x?x96x64xbf16>, %[[C:.+]]: memref<?x?x48x64xf32>,
// 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: %[[C48:.+]] = arith.constant 48 : 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]]] [48, 64] [1, 1]
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[M]], %[[N]], 0, 0] [1, 1, 48, 64] [1, 1, 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: %[[RES:.+]]:2 = scf.for %[[IV_K:.+]] = %[[K_START]] to %[[K_END]] step %[[C1]]
+// CHECK-SAME: iter_args(%[[K0:.+]] = %[[ACC0]], %[[K1:.+]] = %[[ACC1]])
+// CHECK: %[[INNER:.+]]:2 = scf.for %[[IV_L:.+]] = %[[C0]] to %[[C96]] step %[[C32]]
+// CHECK-SAME: iter_args(%[[ARG0:.+]] = %[[K0]], %[[ARG1:.+]] = %[[K1]])
// 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: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[M]], %[[IV_K]], %[[IV_M]], %[[IV_L]]] [1, 1, 16, 32] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x48x96xbf16> to memref<16x32xbf16, strided<[96, 1], offset: ?>>
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[N]], %[[IV_K]], %[[IV_L]], %[[IV_N]]] [1, 1, 32, 32] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x96x64xbf16> to memref<32x32xbf16, strided<[64, 1], offset: ?>>
// 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>
@@ -1094,6 +1110,8 @@ func.func @amx_bf16_flat_1x2_tile(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C
// CHECK: %[[D1:.+]] = vector.contract {{.*}} %[[A0]], %[[B1]], %[[ARG1]] {x86_vcmlu_native_shape = array<i64: 16, 16, 32>}
// CHECK: scf.yield %[[D0]], %[[D1]]
// CHECK: }
+// CHECK: scf.yield %[[INNER]]#0, %[[INNER]]#1
+// CHECK: }
// CHECK: vector.transfer_write %[[RES]]#0, %[[ACC_VIEW]][%[[C0]], %[[C0]]]
// CHECK: vector.transfer_write %[[RES]]#1, %[[ACC_VIEW]][%[[C0]], %[[C16]]]
// CHECK: }
@@ -1292,33 +1310,31 @@ module attributes {transform.with_named_sequence} {
#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>) {
+func.func @amx_int8_flat_batched_loop_nest(%A: memref<?x?x?x64x192xi8>, %B: memref<?x?x?x192x128xi8>, %C: memref<?x?x?x64x128xi32>) {
%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>
+ %batch = memref.dim %C, %c0 : memref<?x?x?x64x128xi32>
+ %M = memref.dim %C, %c1 : memref<?x?x?x64x128xi32>
+ %N = memref.dim %C, %c2 : memref<?x?x?x64x128xi32>
+ %K = memref.dim %A, %c2 : memref<?x?x?x64x192xi8>
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>
+ scf.for %m = %c0 to %M step %c1 {
+ scf.for %n = %c0 to %N step %c1 {
+ %c = vector.transfer_read %C[%bi, %m, %n, %c0, %c0], %c0_i32 {in_bounds = [true, true]} : memref<?x?x?x64x128xi32>, vector<64x128xi32>
+ %res = scf.for %k = %c0 to %K step %c1 iter_args(%acc = %c) -> (vector<64x128xi32>) {
+ %a = vector.transfer_read %A[%bi, %m, %k, %c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?x64x192xi8>, vector<64x192xi8>
+ %b = vector.transfer_read %B[%bi, %n, %k, %c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<?x?x?x192x128xi8>, vector<192x128xi8>
%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>
+ %a, %b, %acc : vector<64x192xi8>, vector<192x128xi8> 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[%bi, %m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<64x128xi32>, memref<?x?x?x64x128xi32>
}
}
}
@@ -1326,7 +1342,7 @@ func.func @amx_int8_flat_batched_loop_nest(%A: memref<?x?x?xi8>, %B: memref<?x?x
}
// 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-SAME: %[[A:.+]]: memref<?x?x?x64x192xi8>, %[[B:.+]]: memref<?x?x?x192x128xi8>, %[[C:.+]]: memref<?x?x?x64x128xi32>)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
@@ -1334,26 +1350,30 @@ func.func @amx_int8_flat_batched_loop_nest(%A: memref<?x?x?xi8>, %B: memref<?x?x
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG: %[[C192:.+]] = arith.constant 192 : 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 %[[OFF_M:.+]] = %[[C0]] to %[[M]] step %[[C1]] {
+// CHECK: scf.for %[[OFF_N:.+]] = %[[C0]] to %[[N]] step %[[C1]] {
+// CHECK: %[[ACC_BUF:.+]] = memref.subview %[[C]][%[[IV_B]], %[[OFF_M]], %[[OFF_N]], 0, 0] [1, 1, 1, 64, 128] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x?x64x128xi32> to memref<64x128xi32, strided<[128, 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: %[[RES:.+]]:4 = scf.for %[[IV_K:.+]] = %[[C0]] to %[[K]] step %[[C1]]
+// CHECK-SAME: iter_args(%[[K0:.+]] = %{{.+}}, %[[K1:.+]] = %{{.+}}, %[[K2:.+]] = %{{.+}}, %[[K3:.+]] = %{{.+}})
+// CHECK: %[[INNER:.+]]:4 = scf.for %[[IV_L:.+]] = %[[C0]] to %[[C192]] step %[[C64]]
+// CHECK-SAME: iter_args(%{{.+}} = %[[K0]], %{{.+}} = %[[K1]], %{{.+}} = %[[K2]], %{{.+}} = %[[K3]])
// 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: %[[A_VIEW:.+]] = memref.subview %[[A]][%[[IV_B]], %[[OFF_M]], %[[IV_K]], %[[IV_M]], %[[IV_L]]] [1, 1, 1, 32, 64] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x?x64x192xi8> to memref<32x64xi8, strided<[192, 1], offset: ?>>
+// CHECK: %[[B_VIEW:.+]] = memref.subview %[[B]][%[[IV_B]], %[[OFF_N]], %[[IV_K]], %[[IV_L]], %[[IV_N]]] [1, 1, 1, 64, 32] [1, 1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x?x192x128xi8> to memref<64x32xi8, strided<[128, 1], offset: ?>>
// 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>
@@ -1361,6 +1381,8 @@ func.func @amx_int8_flat_batched_loop_nest(%A: memref<?x?x?xi8>, %B: memref<?x?x
// 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: scf.yield %[[INNER]]#0, %[[INNER]]#1, %[[INNER]]#2, %[[INNER]]#3
+// CHECK: }
// CHECK-COUNT-4: vector.transfer_write %[[RES]]#{{[0-3]}}, %[[ACC_VIEW]]
// CHECK: }
// CHECK: }
@@ -1504,14 +1526,9 @@ 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>,
+// Loop order K-batch, accumulating over both loops: the inner loops doesn't
+// satisfy the condition for block indexing (last two indices are non-zero).
+func.func @negative_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
@@ -1539,52 +1556,11 @@ func.func @amx_bf16_flat_k_batch_order(%A: memref<?x?x?xbf16>, %B: memref<?x?x?x
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]]]
+// CHECK-LABEL: func.func @negative_k_batch_order(
+// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
-// NANO-LABEL: func.func @amx_bf16_flat_k_batch_order(
+// NANO-LABEL: func.func @negative_k_batch_order(
+// NANO: vector.contract
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -1825,7 +1801,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// N = 48 would require a 2x1 tile register layout, but online packing of the
+// N = 48 would require a 2x1 grid of tile registers, 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)>
@@ -1931,3 +1907,113 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// The bf16 VNNI dimension has extent 4, while the native extent is 2, requiring
+// tiling in the VNNI dimension.
+#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_dimension_tiling(%A: memref<?x?x4xbf16>, %B: memref<?x?x4xbf16>, %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?x4xbf16>, vector<32x16x4xbf16>
+ %b = vector.transfer_read %B[%k, %n, %c0], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x4xbf16>, vector<16x32x4xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<32x16x4xbf16>, vector<16x32x4xbf16> 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 @negative_vnni_dimension_tiling(
+// CHECK: vector.contract {{.*}} : vector<32x16x4xbf16>, vector<16x32x4xbf16> into vector<32x32xf32>
+// CHECK-NOT: x86_vcmlu_native_shape
+
+// NANO-LABEL: func.func @negative_vnni_dimension_tiling(
+// 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_for_amx
+ } : !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_for_amx
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// The RHS read has a non-zero offset in the VNNI dimension.
+#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_nonzero_index(%A: memref<?x?x2xbf16>, %B: memref<?x?x3xbf16>, %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_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, %c1], %c0_0 {in_bounds = [true, true, true]} : memref<?x?x3xbf16>, 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 @negative_vnni_nonzero_index(
+// CHECK: vector.contract {{.*}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+// CHECK-NOT: x86_vcmlu_native_shape
+
+// NANO-LABEL: func.func @negative_vnni_nonzero_index(
+// 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_for_amx
+ } : !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_for_amx
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
>From 31b03afca25761154e87ee4640f7446fed6846f6 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 28 Oct 2026 02:50:09 -0700
Subject: [PATCH 2/4] Rework without IndexKind concept
---
.../VectorContractMultiLevelUnroll.cpp | 170 +++++++++---------
...r-contract-multi-level-unroll-for-amx.mlir | 60 ++++++-
2 files changed, 138 insertions(+), 92 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 94168aa54a8570..2d70561b2b6b3a 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -23,6 +23,7 @@
using namespace mlir;
using llvm::all_of;
+using llvm::drop_begin;
// Discardable attribute to control vector unroll patterns.
static constexpr auto nativeShapeAttrName = "x86_vcmlu_native_shape";
@@ -61,12 +62,6 @@ struct MLUCandidate {
vector::TransferReadOp lhsRead;
vector::TransferReadOp rhsRead;
- // Distinguishes whether the IV of the accumulation loop is used to index a
- // vector dimension ("offset") or a non-vector dimension ("block") on the
- // operand transfer reads.
- enum class IndexKind { Unknown, Offset, Block };
- IndexKind operandIndexKind = IndexKind::Unknown;
-
// Memory backing the accumulator.
Value accMemIn;
Value accMemOut;
@@ -430,23 +425,20 @@ matchTypesAndShape(MLUCandidate &candidate,
}
// Check constraints on the accumulation loop. We're looking for an index-based
-// for-loop with a constant step and 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 accumulation loops from scratch with the known operations.
+// 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 one of their indices. We check the
-// loop's step based on the index kind of the reads: IndexKind::Offset requires
-// a step equal to the contraction's K extent, whereas IndexKind::Block requires
-// unit steps.
+// 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 (with their index kind) are recorded in the
-// candidate struct.
+// accumulator and the operand reads are recorded in the candidate struct.
//
-// Example (IndexKind::Block):
+// Example:
// !accType = vector<64x128xf32>
// !lhsType = vector<64x32xbf16>
// !rhsType = vector<32x128xbf16>
@@ -470,6 +462,11 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return rewriter.notifyMatchFailure(contract,
"is not in an index-based for-loop");
+ if (auto step = getConstantIntValue(accLoop.getStep());
+ !step || (*step != 1 && *step != candidate.contractK))
+ 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,
@@ -498,15 +495,18 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return rewriter.notifyMatchFailure(
contract, "LHS/RHS are not vector.transfer_read ops");
- auto matchOperandRead = [&](vector::TransferReadOp transferOp, bool isLhs) {
- using IK = MLUCandidate::IndexKind;
-
+ auto matchOperandRead = [&](vector::TransferReadOp transferOp,
+ unsigned kVecDim) {
if (!isSimpleTransferOp(transferOp))
return rewriter.notifyMatchFailure(
contract, "operand read has non-identity permutation, mask, or "
"out-of-bounds check");
auto indices = transferOp.getIndices();
+ if (candidate.isInVnniLayout && !isZeroInteger(indices.back()))
+ return rewriter.notifyMatchFailure(
+ contract, "operand read VNNI dimension offset must be zero");
+
auto ivIt =
std::find(indices.begin(), indices.end(), accLoop.getInductionVar());
if (ivIt == indices.end())
@@ -517,59 +517,38 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return rewriter.notifyMatchFailure(
contract, "operand read uses loop index more than once");
- unsigned firstVecDim =
- indices.size() - transferOp.getVectorType().getRank();
+ // The following situation is a limitation in the implementation, so detect
+ // and bail out:
+ //
+ // If the loop index selects a block in a blocked layout, and tiling of the
+ // contraction dimension is required, we have to construct a new innermost
+ // tiling loop over the result vector's K dimension (called "L-loop" in the
+ // transformation later). Now if the operand read has a non-zero offset in
+ // that dimension, we would have to insert an addition into the new loop,
+ // which currently would prevent application of the nanokernel pattern.
+ //
+ // Sketch:
+ // non-zero offset in the vector's K dimension
+ // IV used in non-vector dimension, |
+ // selects a block |
+ // | |
+ // v v
+ // vector.transfer_read %A[%block_M, %IV, %c0, %offset_k] ...
+ // : memref<?x?x?x?xi8>, vector<64x64xi8>
+ unsigned firstVecDimPos = transferOp.getLeadingShapedRank();
unsigned ivPos = std::distance(indices.begin(), ivIt);
- unsigned kPos = firstVecDim + isLhs;
- if (ivPos >= firstVecDim) {
- if (ivPos != kPos)
- return rewriter.notifyMatchFailure(
- contract,
- "loop index is not in expected K position in operand read");
- if (candidate.operandIndexKind != IK::Unknown &&
- candidate.operandIndexKind != IK::Offset)
- return rewriter.notifyMatchFailure(contract,
- "inconsistent operand index kind");
- candidate.operandIndexKind = IK::Offset;
- } else {
- // Temporary limitation: having a zero index in the K position means that
- // we can apply an offset without introducing additional operations in the
- // accumulation loop.
- if (!isZeroInteger(indices[kPos]))
- return rewriter.notifyMatchFailure(
- contract, "operand read K index must be zero for block index kind");
- if (candidate.operandIndexKind != IK::Unknown &&
- candidate.operandIndexKind != IK::Block)
- return rewriter.notifyMatchFailure(contract,
- "inconsistent operand index kind");
- candidate.operandIndexKind = IK::Block;
- }
- assert(candidate.operandIndexKind != IK::Unknown &&
- "operand index kind must be determined");
-
- if (candidate.isInVnniLayout && !isZeroInteger(indices.back()))
- return rewriter.notifyMatchFailure(
- contract, "operand read VNNI dimension offset must be zero");
-
- auto maybeStep = getConstantIntValue(accLoop.getStep());
- if (!maybeStep)
- return rewriter.notifyMatchFailure(
- contract, "accumulation loop step is not constant");
-
- if (candidate.operandIndexKind == IK::Offset &&
- *maybeStep != candidate.contractK)
- return rewriter.notifyMatchFailure(
- contract,
- "loop step must match contraction's K extent for offset indexing");
- if (candidate.operandIndexKind == IK::Block && *maybeStep != 1)
- return rewriter.notifyMatchFailure(
- contract, "loop step must be 1 for block index kind");
+ if (ivPos < firstVecDimPos &&
+ candidate.contractK > candidate.contractNativeShape[2] &&
+ !isZeroInteger(indices[firstVecDimPos + kVecDim]))
+ return rewriter.notifyMatchFailure(contract,
+ "operand read from blocked layout has "
+ "offset in vector's K dimension");
return success();
};
- if (failed(matchOperandRead(lhsRead, /*isLhs=*/true)) ||
- failed(matchOperandRead(rhsRead, /*isLhs=*/false)))
+ if (failed(matchOperandRead(lhsRead, /*[m, k] -> kVecDim=*/1)) ||
+ failed(matchOperandRead(rhsRead, /*[k, n] -> kVecDim=*/0)))
return failure();
candidate.accLoop = accLoop;
@@ -744,10 +723,9 @@ static void buildLoopNestForCandidate(MLUCandidate &candidate,
// 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. Tiling along the K-dimension
-// depends on the operand reads' index kind:
-// IndexKind::Offset: Use the same bounds but smaller steps/more iterations.
-// IndexKind::Block: Nest an additional L-loop to move over the block in
-// steps.
+// depends on the step:
+// K: Use the same bounds but smaller steps/more iterations.
+// 1: Nest an additional L-loop to traverse the block in smaller steps.
// 4. -> buildKLLoopBody
// 5. Splice the final value slices back into an accRegTileType-shaped value and
// write it back to memory.
@@ -756,8 +734,6 @@ static void buildLoopNestForCandidate(MLUCandidate &candidate,
// to unroll them with the vector-dialect's unroll patterns.
static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
Value ivM, Value ivN) {
- using IK = MLUCandidate::IndexKind;
-
ArrayRef<int64_t> accRegTileShape = candidate.accRegTileType.getShape();
Value c0 = arith::ConstantIndexOp::create(b, loc, 0);
@@ -787,29 +763,49 @@ static void buildMNLoopBody(MLUCandidate &candidate, OpBuilder &b, Location loc,
SmallVector<Value> accInitSliced =
slice(b, loc, accInit->getResult(0), candidate.accNativeType);
+ // Iteration space for accumulation loop (nest).
Value lbK = candidate.accLoop.getLowerBound();
Value ubK = candidate.accLoop.getUpperBound();
+ Value stepK = candidate.accLoop.getStep();
+ Value lbL;
+ Value ubL;
+ Value stepL;
+
+ // If we need to tile the K dimension, the simplest way is to adjust the
+ // step to the native K extent. However, if the original accumulation loop
+ // uses unit steps, we have to insert an inner loop to handle the tiling
+ // explicitly.
+ //
// Note: Using _native_ shape here as we never unroll along the reduction
// dimension, so the K extents of the register-tile and native shapes match.
- int64_t stepKInt = candidate.operandIndexKind == IK::Offset
- ? candidate.contractNativeShape[2]
- : 1;
- Value stepK = arith::ConstantIndexOp::create(b, loc, stepKInt);
+ int64_t contractNativeK = candidate.contractNativeShape[2];
+ bool needToTileK = candidate.contractK > contractNativeK;
+ bool needLLoop = false;
+ if (needToTileK) {
+ needLLoop = *getConstantIntValue(stepK) % contractNativeK != 0;
+ Value cContractNativeK =
+ arith::ConstantIndexOp::create(b, loc, contractNativeK);
+ if (needLLoop) {
+ lbL = c0;
+ ubL = arith::ConstantIndexOp::create(b, loc, candidate.contractK);
+ stepL = cContractNativeK;
+ } else {
+ stepK = cContractNativeK;
+ }
+ }
+
candidate.slicedAccLoop = scf::ForOp::create(
b, loc, lbK, ubK, stepK, accInitSliced,
[&](OpBuilder &b, Location loc, Value ivK, ValueRange iterArgsK) {
- if (candidate.operandIndexKind == IK::Offset) {
+ if (!needLLoop) {
SmallVector<Value> newIterArgs = buildKLLoopBody(
candidate, b, loc, ivM, ivN, ivK, /*ivL=*/Value{}, iterArgsK);
scf::YieldOp::create(b, loc, newIterArgs);
return;
}
- Value ubL = arith::ConstantIndexOp::create(b, loc, candidate.contractK);
- Value stepL = arith::ConstantIndexOp::create(
- b, loc, candidate.contractNativeShape[2]);
candidate.slicedInnerAccLoop = scf::ForOp::create(
- b, loc, c0, ubL, stepL, iterArgsK,
+ b, loc, lbL, ubL, stepL, iterArgsK,
[&](OpBuilder &b, Location loc, Value ivL, ValueRange iterArgsL) {
SmallVector<Value> newIterArgs = buildKLLoopBody(
candidate, b, loc, ivM, ivN, ivK, ivL, iterArgsL);
@@ -857,8 +853,12 @@ static SmallVector<Value> buildKLLoopBody(MLUCandidate &candidate, OpBuilder &b,
if (ivL) {
unsigned lPosL = lhsIndices.size() - 1 - candidate.isInVnniLayout;
unsigned lPosR = rhsIndices.size() - 2 - candidate.isInVnniLayout;
- lhsIndices[lPosL] = arith::AddIOp::create(b, loc, lhsIndices[lPosL], ivL);
- rhsIndices[lPosR] = arith::AddIOp::create(b, loc, rhsIndices[lPosR], ivL);
+ // Checked in matchAccumulationLoop.
+ assert(isZeroInteger(lhsIndices[lPosL]) &&
+ isZeroInteger(rhsIndices[lPosR]) &&
+ "Expected zero offsets for L dimension");
+ lhsIndices[lPosL] = ivL;
+ rhsIndices[lPosR] = ivL;
}
Value lhsView = makeRankReducingSubview(b, loc, candidate.lhsRead.getBase(),
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
index 0eadcf1a114177..410d7261da7408 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
@@ -1526,9 +1526,14 @@ 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 loops doesn't
-// satisfy the condition for block indexing (last two indices are non-zero).
-func.func @negative_k_batch_order(%A: memref<?x?x?xbf16>, %B: memref<?x?x?xbf16>, %C: memref<?x?xf32>,
+// 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
@@ -1556,11 +1561,52 @@ func.func @negative_k_batch_order(%A: memref<?x?x?xbf16>, %B: memref<?x?x?xbf16>
func.return
}
-// CHECK-LABEL: func.func @negative_k_batch_order(
-// CHECK: vector.contract {{.*}} : vector<64x32xbf16>, vector<32x96xbf16> into vector<64x96xf32>
+// 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 @negative_k_batch_order(
-// NANO: vector.contract
+// 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 123b2a0f42f39dfafc39443f8a9fbfb330800d35 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 28 Oct 2026 05:27:39 -0700
Subject: [PATCH 3/4] Test for unsupported situation
---
...r-contract-multi-level-unroll-for-amx.mlir | 55 +++++++++++++++++++
1 file changed, 55 insertions(+)
diff --git a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
index 410d7261da7408..15945d52bb11d5 100644
--- a/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-multi-level-unroll-for-amx.mlir
@@ -2063,3 +2063,58 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// A blocked K loop cannot be tiled when either operand has a non-zero offset
+// in its vector K dimension.
+#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_blocked_lhs_k_offset(%A: memref<?x?x32x96xbf16>, %B: memref<?x?x64x32xbf16>, %C: memref<?x?x32x32xf32>,
+ %m: index, %n: index, %k_start: index, %k_end: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : 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, %c0], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?x32x32xf32>, vector<32x32xf32>
+ %res = scf.for %k = %k_start to %k_end step %c1 iter_args(%acc = %c) -> (vector<32x32xf32>) {
+ %a = vector.transfer_read %A[%m, %k, %c0, %c32], %c0_0 {in_bounds = [true, true]} : memref<?x?x32x96xbf16>, vector<32x64xbf16>
+ %b = vector.transfer_read %B[%n, %k, %c0, %c0], %c0_0 {in_bounds = [true, true]} : memref<?x?x64x32xbf16>, vector<64x32xbf16>
+ %d = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %a, %b, %acc : vector<32x64xbf16>, vector<64x32xbf16> into vector<32x32xf32>
+ scf.yield %d : vector<32x32xf32>
+ }
+ vector.transfer_write %res, %C[%m, %n, %c0, %c0] {in_bounds = [true, true]} : vector<32x32xf32>, memref<?x?x32x32xf32>
+ func.return
+}
+
+// CHECK-LABEL: func.func @negative_blocked_lhs_k_offset(
+// CHECK: vector.contract {{.*}} : vector<32x64xbf16>, vector<64x32xbf16> into vector<32x32xf32>
+
+// NANO-LABEL: func.func @negative_blocked_lhs_k_offset(
+// 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_for_amx
+ } : !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_for_amx
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
>From 4a69833f70495c425471d2225ab68134f52c7866 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Wed, 28 Oct 2026 05:48:57 -0700
Subject: [PATCH 4/4] Naming
---
.../Transforms/VectorContractMultiLevelUnroll.cpp | 13 ++++++-------
1 file changed, 6 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
index 2d70561b2b6b3a..c293a19f98d15f 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractMultiLevelUnroll.cpp
@@ -495,14 +495,13 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return rewriter.notifyMatchFailure(
contract, "LHS/RHS are not vector.transfer_read ops");
- auto matchOperandRead = [&](vector::TransferReadOp transferOp,
- unsigned kVecDim) {
- if (!isSimpleTransferOp(transferOp))
+ auto checkOperandRead = [&](vector::TransferReadOp read, unsigned kVecDim) {
+ if (!isSimpleTransferOp(read))
return rewriter.notifyMatchFailure(
contract, "operand read has non-identity permutation, mask, or "
"out-of-bounds check");
- auto indices = transferOp.getIndices();
+ auto indices = read.getIndices();
if (candidate.isInVnniLayout && !isZeroInteger(indices.back()))
return rewriter.notifyMatchFailure(
contract, "operand read VNNI dimension offset must be zero");
@@ -535,7 +534,7 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
// v v
// vector.transfer_read %A[%block_M, %IV, %c0, %offset_k] ...
// : memref<?x?x?x?xi8>, vector<64x64xi8>
- unsigned firstVecDimPos = transferOp.getLeadingShapedRank();
+ unsigned firstVecDimPos = read.getLeadingShapedRank();
unsigned ivPos = std::distance(indices.begin(), ivIt);
if (ivPos < firstVecDimPos &&
candidate.contractK > candidate.contractNativeShape[2] &&
@@ -547,8 +546,8 @@ static LogicalResult matchAccumulationLoop(MLUCandidate &candidate,
return success();
};
- if (failed(matchOperandRead(lhsRead, /*[m, k] -> kVecDim=*/1)) ||
- failed(matchOperandRead(rhsRead, /*[k, n] -> kVecDim=*/0)))
+ if (failed(checkOperandRead(lhsRead, /*[m, k] -> kVecDim=*/1)) ||
+ failed(checkOperandRead(rhsRead, /*[k, n] -> kVecDim=*/0)))
return failure();
candidate.accLoop = accLoop;
More information about the Mlir-commits
mailing list