[Mlir-commits] [mlir] [mlir][x86] Fix - multiple issues with AMX dot-product lowering (PR #196984)
Arun Thangamani
llvmlistbot at llvm.org
Tue May 12 05:55:49 PDT 2026
https://github.com/arun-thmn updated https://github.com/llvm/llvm-project/pull/196984
>From 91394a7e8e01385a767407aa2340cf9e2486bc4c Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Mon, 11 May 2026 09:02:23 -0700
Subject: [PATCH 1/2] fixex issues with AMX dot-product lowering
---
.../VectorContractToAMXDotProduct.cpp | 391 +++++++++---------
1 file changed, 202 insertions(+), 189 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 94b94292e675f..551fccb47e114 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -27,6 +27,31 @@ using namespace mlir::x86;
namespace {
+static Value traceToVectorWriteLikeUserOperationForAMX(Value v) {
+ if (v.getNumUses() > 1)
+ return nullptr;
+
+ for (OpOperand &use : v.getUses()) {
+ Operation *user = use.getOwner();
+
+ if (!isa<scf::YieldOp>(user)) {
+ return v;
+ }
+
+ // --- SCF YIELD ---
+ if (auto yield = dyn_cast<scf::YieldOp>(user)) {
+ Operation *parent = yield->getParentOp();
+ unsigned idx = use.getOperandNumber();
+ if (auto res =
+ traceToVectorWriteLikeUserOperationForAMX(parent->getResult(idx)))
+ return res;
+ continue;
+ }
+ }
+
+ return nullptr;
+}
+
// Function to collapse the last two dimension (vnni and k) to help the
// amx.tile_load to correctly load the packed element type.
static Value collapseInnerDims(OpBuilder &builder, mlir::Location loc,
@@ -216,22 +241,30 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
Value c16 = arith::ConstantIndexOp::create(rewriter, loc, 16);
auto subview = matB.getDefiningOp<mlir::memref::SubViewOp>();
- SmallVector<Value> subviewOffset(subview.getOffsets().size(), c0);
+ SmallVector<Value> subviewOffset(subview.getMixedOffsets().size(), c0);
Value cStep = arith::ConstantIndexOp::create(rewriter, loc, offset);
Value cBound = arith::ConstantIndexOp::create(rewriter, loc, (16 * offset));
Value offsetIndx =
arith::ConstantIndexOp::create(rewriter, loc, (offset / 2));
+ // llvm::outs() << "check-a:" << matB << " subview:" << subview << "\n";
+ // llvm::outs() << "The size:" << subview.getMixedOffsets().size() << "\n";
+
scf::ForOp::create(
rewriter, loc, c0, cBound, cStep, ValueRange{},
[&](OpBuilder &nestedBuilder, Location loc, Value iv,
ValueRange iterArgs) {
+ // llvm::outs() << "check-a0" << subviewOffset.size() << "\n";
subviewOffset[subviewOffset.size() - 2] = iv;
+
+ // llvm::outs() << "check-a1" << "\n";
auto vec1 = vector::LoadOp::create(
rewriter, loc, VectorType::get((16 * offset), ipType), matB,
ValueRange(subviewOffset));
+ // llvm::outs() << "check-b" << "\n";
+
// Increment the iv by 1 or 2 based on the type to load the next 32/64
// elements
Value incIV = arith::AddIOp::create(rewriter, loc, offsetIndx, iv);
@@ -243,6 +276,8 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
vector::ShuffleOp shuffle1;
vector::ShuffleOp shuffle2;
+ // llvm::outs() << "check-c" << "\n";
+
if (ipType.isBF16()) {
shuffle1 = vector::ShuffleOp::create(
@@ -283,6 +318,8 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
30, 62, 94, 126, 31, 63, 95, 127});
}
+ // llvm::outs() << "check-d" << "\n";
+
// iv to store the shuffled elements
Value ivShuff1 = arith::DivUIOp::create(rewriter, loc, iv, cStep);
Value ivShuff2 = arith::AddIOp::create(rewriter, loc, ivShuff1, c16);
@@ -468,6 +505,8 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
Value c1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
Value c2 = arith::ConstantIndexOp::create(rewriter, loc, 2);
+ int64_t offset = step.getDefiningOp<arith::ConstantIndexOp>().value();
+
auto newLoop = scf::ForOp::create(
rewriter, loc, lowerBound, upperBound, step, loopItrArgs,
[&](OpBuilder &rewriterNewInnerLoop, Location locNewInnerLoop,
@@ -485,7 +524,6 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
Value indxToStoreInBuffer = c0;
Value indxToLoadFromBuffer = c0;
-
if (!isVnni) {
if (outerLoop) {
if (innerLoopIndex.value() == 0) {
@@ -509,7 +547,7 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
} else {
Value nLoadIndx = arith::ConstantIndexOp::create(
- rewriter, locNewInnerLoop, (16 * blockingFactor));
+ rewriter, locNewInnerLoop, offset);
ivNewInnerLoop = arith::AddIOp::create(rewriter, locNewInnerLoop,
nLoadIndx, ivNewInnerLoop);
indxToStoreInBuffer = getIndxToLoadStoreFromPckBuffer(
@@ -525,7 +563,7 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
} else {
if (pack) {
Value nLoadIndx = arith::ConstantIndexOp::create(
- rewriter, locNewInnerLoop, (16 * blockingFactor));
+ rewriter, locNewInnerLoop, offset);
ivNewInnerLoop = arith::AddIOp::create(rewriter, locNewInnerLoop,
nLoadIndx, ivNewInnerLoop);
Value quotient_K = arith::DivUIOp::create(
@@ -541,27 +579,49 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
}
}
}
-
IRMapping rhsMapping;
- if (outerLoop)
- rhsMapping.map(
- vectorOpRhs->getOperand(
- getIndexPosition(contractOp.getRhs(), outerLoop) + 1),
- ivOuterLoop);
- rhsMapping.map(
- vectorOpRhs->getOperand(
- getIndexPosition(contractOp.getRhs(), innerLoop) + 1),
- ivNewInnerLoop);
- auto rhsClone = rewriterNewInnerLoop.clone(*vectorOpRhs, rhsMapping);
+ Value matB;
+ Operation *rhsOp = vectorOpRhs;
- Value matB = rhsClone->getResult(0);
+ // Clone only if the op has operands.
+ if (rhsOp->getNumOperands() > 0) {
+ if (outerLoop) {
+ int64_t outerPos = getIndexPosition(contractOp.getRhs(), outerLoop);
+
+ if (outerPos >= 0) {
+ unsigned operandIdx = static_cast<unsigned>(outerPos + 1);
+
+ if (operandIdx < rhsOp->getNumOperands()) {
+ rhsMapping.map(rhsOp->getOperand(operandIdx), ivOuterLoop);
+ }
+ }
+ }
+
+ int64_t innerPos = getIndexPosition(contractOp.getRhs(), innerLoop);
+
+ if (innerPos >= 0) {
+ unsigned operandIdx = static_cast<unsigned>(innerPos + 1);
+
+ if (operandIdx < rhsOp->getNumOperands()) {
+ rhsMapping.map(rhsOp->getOperand(operandIdx), ivNewInnerLoop);
+ }
+ }
+
+ auto rhsClone = rewriterNewInnerLoop.clone(*rhsOp, rhsMapping);
+
+ matB = rhsClone->getResult(0);
+
+ } else {
+ // memref.get_global / constants
+ matB = rhsOp->getResult(0);
+ }
if (!isVnni) {
if (outerLoop) {
if (!pack) {
Value nLoadIndx = arith::ConstantIndexOp::create(
- rewriter, locNewInnerLoop, (16 * blockingFactor));
+ rewriter, locNewInnerLoop, offset);
matB = Value();
indxToLoadFromBuffer = c0;
indxToLoadFromBuffer = getIndxToLoadStoreFromPckBuffer(
@@ -572,7 +632,7 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
} else {
if (!pack) {
Value nLoadIndx = arith::ConstantIndexOp::create(
- rewriter, locNewInnerLoop, (16 * blockingFactor));
+ rewriter, locNewInnerLoop, offset);
matB = Value();
Value quotient_K = arith::DivUIOp::create(
rewriter, loc, ivNewInnerLoop, nLoadIndx);
@@ -581,7 +641,6 @@ createLoops(OpBuilder &rewriter, Location loc, Value lowerBound,
}
}
}
-
// compute tiled dot-product
SmallVector<Value> accumulators = createTiledDp(
rewriter, locNewInnerLoop, ops, lhsClone->getResult(0), matB,
@@ -860,7 +919,7 @@ struct VectorContractToAMXDotProduct
14, 30, 46, 62, 15, 31, 47, 63});
}
- auto rem = arith::RemUIOp::create(
+ auto rem = arith::DivUIOp::create(
rewriter, loc, rewriter.getIndexType(), iv, step);
vector::StoreOp::create(rewriter, loc, shuffle1, packedBuffer,
@@ -988,6 +1047,7 @@ struct VectorContractToAMXDotProduct
scf::ForOp newLoop;
// Case 2a: Reduction loop depth is 2.
if (loopLists.size() == 2) {
+
outerLoop = loopLists[1];
innerLoop = loopLists[0];
@@ -1120,8 +1180,8 @@ struct VectorContractToAMXDotProduct
// Case 2b: Reduction loop depth is 1.
if (loopLists.size() == 1) {
- innerLoop = loopLists[0];
+ innerLoop = loopLists[0];
SmallVector<Value> loopItrArgs = createTileZeros(
rewriter, innerLoop.getLoc(), opType, innerLoop, ops.size());
@@ -1135,6 +1195,7 @@ struct VectorContractToAMXDotProduct
nullptr, false, false);
} else {
+
bool isInnerLoopUBLarger = false;
bool isInnerLoopUBHasOddQuot = false;
@@ -1154,8 +1215,12 @@ struct VectorContractToAMXDotProduct
rewriter.setInsertionPoint(innerLoop);
auto c0 =
arith::ConstantIndexOp::create(rewriter, innerLoop.getLoc(), 0);
+
+ int64_t stepVal =
+ innerLoop.getStep().getDefiningOp<arith::ConstantIndexOp>().value();
+
auto spillLoopBound = arith::ConstantIndexOp::create(
- rewriter, innerLoop.getLoc(), 16 * blockingFactor);
+ rewriter, innerLoop.getLoc(), stepVal);
Value spillInnerLoop =
arith::SubIOp::create(rewriter, innerLoop.getLoc(),
@@ -1173,10 +1238,8 @@ struct VectorContractToAMXDotProduct
getIndexPosition(contractOp.getRhs(), innerLoop) + 1),
c0);
auto rhsClone = rewriter.clone(*vectorOpRhs, rhsMapping);
-
performShuffle(rewriter, innerLoop.getLoc(), rhsClone->getResult(0),
ipType, blockingFactor, packedBuffer, c0);
-
auto newLoopNonSpill = createLoops(
rewriter, innerLoop.getLoc(), innerLoop.getLowerBound(),
spillInnerLoop, innerLoop.getStep(), loopItrArgs, ipType, opType,
@@ -1200,194 +1263,144 @@ struct VectorContractToAMXDotProduct
// Copy the amx tile accumulation results to a MemRef buffer, add the
// initial accumulation value, and store back to the C-Matrix
+ Location loc = outerLoop.getLoc();
+ Value srcBuffAcc;
+ SmallVector<Value> indicesAcc;
+
+ llvm::TypeSwitch<Operation *>(accReadOp).Case<TransferReadOp, LoadOp>(
+ [&](auto readOp) {
+ srcBuffAcc = readOp.getOperand(0);
+
+ auto indices = readOp.getIndices();
+ indicesAcc.reserve(indices.size());
+
+ llvm::transform(indices, std::back_inserter(indicesAcc),
+ [&](OpFoldResult ofr) {
+ return mlir::getValueOrCreateConstantIndexOp(
+ rewriter, loc, ofr);
+ });
+ });
- if (!isVnni) {
- Location loc = outerLoop.getLoc();
- Operation *accReadOp =
- traceToVectorReadLikeParentOperation(contractOp.getAcc());
-
- Value srcBuffAcc;
- SmallVector<Value> indicesAcc;
-
- llvm::TypeSwitch<Operation *>(accReadOp).Case<TransferReadOp, LoadOp>(
- [&](auto readOp) {
- srcBuffAcc = readOp.getOperand(0);
-
- auto indices = readOp.getIndices();
- indicesAcc.reserve(indices.size());
-
- llvm::transform(indices, std::back_inserter(indicesAcc),
- [&](OpFoldResult ofr) {
- return mlir::getValueOrCreateConstantIndexOp(
- rewriter, loc, ofr);
- });
- });
-
- auto outputShapes =
- mlir::cast<mlir::MemRefType>(srcBuffAcc.getType()).getShape();
- unsigned int M = outputShapes[outputShapes.size() - 2];
- unsigned int N = outputShapes[outputShapes.size() - 1];
-
- SmallVector<Value> dps = newLoop.getResults();
- auto bufferType = MemRefType::get({M, N}, opType);
- auto resultBuffer = memref::AllocaOp::create(rewriter, loc, bufferType);
-
- // Store the amx tiled-dot product output into an MxN memref.
- for (unsigned int i = 0, k = 0; i < M; i = i + 16) {
- for (unsigned int j = 0; j < N; j = j + 16) {
- Value indexOp_i = arith::ConstantIndexOp::create(rewriter, loc, i);
- Value indexOp_j = arith::ConstantIndexOp::create(rewriter, loc, j);
- amx::TileStoreOp::create(rewriter, loc, resultBuffer,
- ValueRange{indexOp_i, indexOp_j}, dps[k]);
- k++;
- }
+ auto outputShapes =
+ mlir::cast<mlir::MemRefType>(srcBuffAcc.getType()).getShape();
+ unsigned int M = outputShapes[outputShapes.size() - 2];
+ unsigned int N = outputShapes[outputShapes.size() - 1];
+
+ SmallVector<Value> dps = newLoop.getResults();
+ auto bufferType = MemRefType::get({M, N}, opType);
+ auto resultBuffer = memref::AllocaOp::create(rewriter, loc, bufferType);
+
+ // Store the amx tiled-dot product output into an MxN memref.
+ for (unsigned int i = 0, k = 0; i < M; i = i + 16) {
+ for (unsigned int j = 0; j < N; j = j + 16) {
+ Value indexOp_i = arith::ConstantIndexOp::create(rewriter, loc, i);
+ Value indexOp_j = arith::ConstantIndexOp::create(rewriter, loc, j);
+ amx::TileStoreOp::create(rewriter, loc, resultBuffer,
+ ValueRange{indexOp_i, indexOp_j}, dps[k]);
+ k++;
}
- auto c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
- auto c16 = arith::ConstantIndexOp::create(rewriter, loc, 16);
- auto one = arith::ConstantIndexOp::create(rewriter, loc, 1);
- auto mBound = arith::ConstantIndexOp::create(rewriter, loc, N);
-
- // Create a loop that iterates over the MxN memerf, retrives two rows +
- // shuffle them, add up the C element values and stores them back.
- scf::ForOp::create(
- rewriter, loc, c0, mBound, one, ValueRange{},
- [&](OpBuilder &nestedBuilder, Location loc, Value iv,
- ValueRange iterArgs) {
- auto row = vector::LoadOp::create(rewriter, loc,
- VectorType::get(16, opType),
- resultBuffer, ValueRange{iv, c0});
-
- auto row2 = vector::LoadOp::create(
- rewriter, loc, VectorType::get(16, opType), resultBuffer,
- ValueRange{iv, c16});
-
- auto shuffle1 = vector::ShuffleOp::create(
+ }
+ auto c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ auto c16 = arith::ConstantIndexOp::create(rewriter, loc, 16);
+ auto one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ auto mBound = arith::ConstantIndexOp::create(rewriter, loc, N);
+
+ // Create a loop that iterates over the MxN memerf, retrives two rows +
+ // shuffle them, add up the C element values and stores them back.
+ scf::ForOp::create(
+ rewriter, loc, c0, mBound, one, ValueRange{},
+ [&](OpBuilder &nestedBuilder, Location loc, Value iv,
+ ValueRange iterArgs) {
+ auto row =
+ vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
+ resultBuffer, ValueRange{iv, c0});
+
+ auto row2 =
+ vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
+ resultBuffer, ValueRange{iv, c16});
+
+ Value shuffle1 = row;
+ Value shuffle2 = row2;
+
+ if (!isVnni) {
+ shuffle1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get(16, opType), row, row2,
ArrayRef<int64_t>{0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20,
21, 22, 23});
- auto shuffle2 = vector::ShuffleOp::create(
+ shuffle2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get(16, opType), row, row2,
ArrayRef<int64_t>{8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15,
28, 29, 30, 31});
+ }
+ indicesAcc[indicesAcc.size() - 2] = iv;
+ indicesAcc[indicesAcc.size() - 1] = c0;
- indicesAcc[indicesAcc.size() - 2] = iv;
- indicesAcc[indicesAcc.size() - 1] = c0;
+ Value valueCRow1 =
+ vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
+ srcBuffAcc, indicesAcc);
+ indicesAcc[indicesAcc.size() - 1] = c16;
- Value valueCRow1 = vector::LoadOp::create(
- rewriter, loc, VectorType::get(16, opType), srcBuffAcc,
- indicesAcc);
- indicesAcc[indicesAcc.size() - 1] = c16;
+ Value valueCRow2 =
+ vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
+ srcBuffAcc, indicesAcc);
- Value valueCRow2 = vector::LoadOp::create(
- rewriter, loc, VectorType::get(16, opType), srcBuffAcc,
- indicesAcc);
+ Value addOp;
+ Value addOp2;
- Value addOp;
- Value addOp2;
+ if (ipType.isBF16()) {
+ addOp = arith::AddFOp::create(rewriter, loc, shuffle1, valueCRow1);
- if (ipType.isBF16()) {
- addOp =
- arith::AddFOp::create(rewriter, loc, shuffle1, valueCRow1);
+ addOp2 = arith::AddFOp::create(rewriter, loc, shuffle2, valueCRow2);
+ }
- addOp2 =
- arith::AddFOp::create(rewriter, loc, shuffle2, valueCRow2);
- }
+ if (ipType.isSignlessInteger(8)) {
+ addOp = arith::AddIOp::create(rewriter, loc, shuffle1, valueCRow1);
- if (ipType.isSignlessInteger(8)) {
- addOp =
- arith::AddIOp::create(rewriter, loc, shuffle1, valueCRow1);
+ addOp2 = arith::AddIOp::create(rewriter, loc, shuffle2, valueCRow2);
+ }
- addOp2 =
- arith::AddIOp::create(rewriter, loc, shuffle2, valueCRow2);
- }
- indicesAcc[indicesAcc.size() - 1] = c0;
- vector::StoreOp::create(rewriter, loc, addOp, srcBuffAcc,
- indicesAcc);
- indicesAcc[indicesAcc.size() - 1] = c16;
- vector::StoreOp::create(rewriter, loc, addOp2, srcBuffAcc,
- indicesAcc);
-
- scf::YieldOp::create(nestedBuilder, loc);
- });
- }
+ vector::StoreOp::create(rewriter, loc, addOp, resultBuffer,
+ ValueRange{iv, c0});
+ vector::StoreOp::create(rewriter, loc, addOp2, resultBuffer,
+ ValueRange{iv, c16});
- auto bufferType = MemRefType::get({16, 16}, opType);
- auto resultBuffer =
- memref::AllocaOp::create(rewriter, outerLoop.getLoc(), bufferType);
- SmallVector<Value> dps = newLoop.getResults();
+ scf::YieldOp::create(nestedBuilder, loc);
+ });
+
+ SmallVector<Value> writeResults;
+ for (unsigned int i = 0; i < M; i = i + 16) {
+ for (unsigned int j = 0; j < N; j = j + 16) {
+ Value indexOp_i = arith::ConstantIndexOp::create(rewriter, loc, i);
+ Value indexOp_j = arith::ConstantIndexOp::create(rewriter, loc, j);
+
+ auto flatTy = mlir::VectorType::get({16, 16}, opType);
+
+ int64_t srcRank =
+ (dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
+ Value padding = ub::PoisonOp::create(rewriter, loc, opType);
+ auto map = AffineMap::getMinorIdentityMap(srcRank, flatTy.getRank(),
+ rewriter.getContext());
+ SmallVector<bool> inBounds(flatTy.getRank(), true);
+
+ auto vec1 = vector::TransferReadOp::create(
+ rewriter, loc, flatTy, resultBuffer,
+ ValueRange{indexOp_i, indexOp_j}, padding, map, inBounds);
+ writeResults.push_back(vec1);
+ }
+ }
for (size_t i = 0; i < ops.size(); i++) {
vector::ContractionOp contOp = ops[i];
- Operation *resultWriteOp =
- traceToVectorWriteLikeUserOperation(contOp.getResult());
- if (isVnni) {
- rewriter.setInsertionPoint(resultWriteOp);
-
- Value indexOp_0 =
- arith::ConstantIndexOp::create(rewriter, outerLoop.getLoc(), 0);
-
- amx::TileStoreOp::create(rewriter, outerLoop.getLoc(), resultBuffer,
- ValueRange{indexOp_0, indexOp_0}, dps[i]);
-
- auto c0 =
- arith::ConstantIndexOp::create(rewriter, outerLoop.getLoc(), 0);
- auto one =
- arith::ConstantIndexOp::create(rewriter, outerLoop.getLoc(), 1);
- auto mBound =
- arith::ConstantIndexOp::create(rewriter, outerLoop.getLoc(), 16);
-
- scf::ForOp::create(
- rewriter, outerLoop.getLoc(), c0, mBound, one, ValueRange{},
- [&](OpBuilder &builder, Location loc, Value iv,
- ValueRange iterArgs) {
- auto resultAcc = vector::LoadOp::create(
- rewriter, loc, VectorType::get(16, opType), resultBuffer,
- ValueRange{iv, c0});
-
- Operation *accReadOp =
- traceToVectorReadLikeParentOperation(ops[i].getAcc());
-
- Value srcBuffAcc;
- SmallVector<Value> indicesAcc;
+ Value vecRoc = writeResults[i];
- llvm::TypeSwitch<Operation *>(accReadOp)
- .Case<TransferReadOp, LoadOp>([&](auto readOp) {
- srcBuffAcc = readOp.getOperand(0);
-
- auto indices = readOp.getIndices();
- indicesAcc.reserve(indices.size());
-
- llvm::transform(
- indices, std::back_inserter(indicesAcc),
- [&](OpFoldResult ofr) {
- return mlir::getValueOrCreateConstantIndexOp(
- rewriter, loc, ofr);
- });
- });
-
- Value sum =
- arith::AddIOp::create(builder, loc, iv, indicesAcc[0]);
- indicesAcc[indicesAcc.size() - 2] = sum;
-
- auto acc = vector::LoadOp::create(rewriter, loc,
- VectorType::get(16, opType),
- srcBuffAcc, indicesAcc);
- Value addition;
- if (ipType.isBF16())
- addition = arith::AddFOp::create(rewriter, loc, resultAcc, acc);
-
- if (ipType.isSignlessInteger(8))
- addition = arith::AddIOp::create(rewriter, loc, resultAcc, acc);
-
- vector::StoreOp::create(builder, loc, addition, srcBuffAcc,
- indicesAcc);
-
- scf::YieldOp::create(builder, outerLoop.getLoc());
- });
+ Value resultWriteOp =
+ traceToVectorWriteLikeUserOperationForAMX(contOp.getResult());
+ if (auto vecType = llvm::dyn_cast<VectorType>(resultWriteOp.getType())) {
+ vecRoc = mlir::vector::ShapeCastOp::create(rewriter, loc, vecType,
+ writeResults[i]);
}
-
- rewriter.eraseOp(resultWriteOp);
+ resultWriteOp.replaceAllUsesWith(vecRoc);
}
return success();
>From 8e448d5f2a3e369dcc9514587e242e913d3e7f67 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 12 May 2026 05:55:29 -0700
Subject: [PATCH 2/2] counting offset on the subview result
---
.../VectorContractToAMXDotProduct.cpp | 47 ++++----
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 100 ++++++++++++++++--
2 files changed, 122 insertions(+), 25 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 551fccb47e114..64e5a6b56504b 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -239,32 +239,24 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
Value c16 = arith::ConstantIndexOp::create(rewriter, loc, 16);
-
- auto subview = matB.getDefiningOp<mlir::memref::SubViewOp>();
- SmallVector<Value> subviewOffset(subview.getMixedOffsets().size(), c0);
+ SmallVector<Value> subviewOffset(
+ llvm::cast<MemRefType>(matB.getType()).getRank(), c0);
Value cStep = arith::ConstantIndexOp::create(rewriter, loc, offset);
Value cBound = arith::ConstantIndexOp::create(rewriter, loc, (16 * offset));
Value offsetIndx =
arith::ConstantIndexOp::create(rewriter, loc, (offset / 2));
- // llvm::outs() << "check-a:" << matB << " subview:" << subview << "\n";
- // llvm::outs() << "The size:" << subview.getMixedOffsets().size() << "\n";
-
scf::ForOp::create(
rewriter, loc, c0, cBound, cStep, ValueRange{},
[&](OpBuilder &nestedBuilder, Location loc, Value iv,
ValueRange iterArgs) {
- // llvm::outs() << "check-a0" << subviewOffset.size() << "\n";
subviewOffset[subviewOffset.size() - 2] = iv;
- // llvm::outs() << "check-a1" << "\n";
auto vec1 = vector::LoadOp::create(
rewriter, loc, VectorType::get((16 * offset), ipType), matB,
ValueRange(subviewOffset));
- // llvm::outs() << "check-b" << "\n";
-
// Increment the iv by 1 or 2 based on the type to load the next 32/64
// elements
Value incIV = arith::AddIOp::create(rewriter, loc, offsetIndx, iv);
@@ -276,8 +268,6 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
vector::ShuffleOp shuffle1;
vector::ShuffleOp shuffle2;
- // llvm::outs() << "check-c" << "\n";
-
if (ipType.isBF16()) {
shuffle1 = vector::ShuffleOp::create(
@@ -318,8 +308,6 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
30, 62, 94, 126, 31, 63, 95, 127});
}
- // llvm::outs() << "check-d" << "\n";
-
// iv to store the shuffled elements
Value ivShuff1 = arith::DivUIOp::create(rewriter, loc, iv, cStep);
Value ivShuff2 = arith::AddIOp::create(rewriter, loc, ivShuff1, c16);
@@ -829,6 +817,8 @@ struct VectorContractToAMXDotProduct
"The ACC src is not a MemRef type.");
auto [srcBuffAcc, indicesAcc] = *srcIndxAcc;
+ Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
+
// amx.tile_loads
auto tileType = amx::TileType::get({16, (16 * blockingFactor)}, ipType);
auto loadLhs = amx::TileLoadOp::create(rewriter, loc, tileType,
@@ -856,7 +846,6 @@ struct VectorContractToAMXDotProduct
auto packedBuffer = memref::AllocaOp::create(rewriter, loc, bufferType);
// create a loop that does online packing.
- Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
Value step =
arith::ConstantIndexOp::create(rewriter, loc, blockingFactor);
Value uBound = arith::ConstantIndexOp::create(rewriter, loc,
@@ -951,9 +940,32 @@ struct VectorContractToAMXDotProduct
dp = amx::TileMulIOp::create(rewriter, loc, tileTypeAcc, loadLhs,
loadRhs, loadAcc);
- amx::TileStoreOp::create(rewriter, loc, srcBuffAcc, indicesAcc, dp);
+ auto bufferType = MemRefType::get({16, 16}, opType);
+ auto resultBuffer = memref::AllocaOp::create(rewriter, loc, bufferType);
+
+ amx::TileStoreOp::create(rewriter, loc, resultBuffer, ValueRange{c0, c0},
+ dp);
+
+ auto flatTy = mlir::VectorType::get({16, 16}, opType);
+ int64_t srcRank =
+ (dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
+ Value padding = ub::PoisonOp::create(rewriter, loc, opType);
+ auto map = AffineMap::getMinorIdentityMap(srcRank, flatTy.getRank(),
+ rewriter.getContext());
+ SmallVector<bool> inBounds(flatTy.getRank(), true);
+
+ Value vecRow = vector::TransferReadOp::create(
+ rewriter, loc, flatTy, resultBuffer, ValueRange{c0, c0}, padding, map,
+ inBounds);
+
+ Value resultOp =
+ traceToVectorWriteLikeUserOperationForAMX(contractOp.getResult());
+ if (auto vecType = llvm::dyn_cast<VectorType>(resultOp.getType())) {
+ vecRow =
+ mlir::vector::ShapeCastOp::create(rewriter, loc, vecType, vecRow);
+ }
- rewriter.eraseOp(resultWriteOp);
+ resultOp.replaceAllUsesWith(vecRow);
return success();
}
@@ -1186,7 +1198,6 @@ struct VectorContractToAMXDotProduct
rewriter, innerLoop.getLoc(), opType, innerLoop, ops.size());
if (isVnni) {
-
newLoop = createLoops(
rewriter, innerLoop.getLoc(), innerLoop.getLowerBound(),
innerLoop.getUpperBound(), innerLoop.getStep(), loopItrArgs, ipType,
diff --git a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
index 1a6deed31eceb..20d269fd6ff88 100644
--- a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
+++ b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
@@ -239,13 +239,17 @@ func.func @online_packing_int8(
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
+ %bias = arith.constant dense<13> : !vecC
+
%4 = vector.contract {
indexing_maps = [#map, #map1, #map2],
iterator_types = ["parallel", "parallel", "reduction"],
kind = #vector.kind<add>}
%1, %2, %3 : !vecA, !vecB into !vecC
- vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+ %5 = arith.addi %4, %bias : !vecC
+
+ vector.transfer_write %5, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
return %arg2 : !memrefC
}
@@ -259,10 +263,11 @@ func.func @online_packing_int8(
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xi32>
// CHECK: x86.amx.tile_muli
// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xi32>
+// CHECK: vector.transfer_read
+// CHECK: arith.addi
+// CHECK: vector.transfer_write
// CHECK-NOT: 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
@@ -695,7 +700,80 @@ func.func @online_packing_bf16_loop(%arg0: memref<16x64x96xbf16>, %arg1: memref<
// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
// CHECK-NOT: 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_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecAB = vector<1x16x16x2xbf16>
+!vecC = vector<16x16xf32>
+!memrefA = memref<1x32x16x2xbf16, strided<[1024, 32, 2, 1], offset: ?>>
+!memrefB = memref<1x16x32x2xbf16, strided<[1024, 64, 2, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4, d1)>
+#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3, d1)>
+#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>
+
+func.func @brgemm_bf16_with_cano(%arg0: memref<16x32x16x2xbf16>, %arg1: memref<16x16x32x2xbf16>, %arg2: memref<32x32xf32>) -> memref<32x32xf32> {
+ %0 = ub.poison : f32
+ %1 = ub.poison : bf16
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c1 = arith.constant 1 : index
+ %2 = vector.transfer_read %arg2[%c0, %c0], %0 {in_bounds = [true, true]} : memref<32x32xf32>, !vecC
+ %3 = vector.transfer_read %arg2[%c0, %c16], %0 {in_bounds = [true, true]} : memref<32x32xf32>, !vecC
+ %4 = vector.transfer_read %arg2[%c16, %c0], %0 {in_bounds = [true, true]} : memref<32x32xf32>, !vecC
+ %5 = vector.transfer_read %arg2[%c16, %c16], %0 {in_bounds = [true, true]} : memref<32x32xf32>, !vecC
+
+ %6:4 = scf.for %arg3 = %c0 to %c16 step %c1 iter_args(%arg4 = %2, %arg5 = %3, %arg6 = %4, %arg7 = %5) -> (!vecC, !vecC, !vecC, !vecC) {
+
+ %subview = memref.subview %arg0[%arg3, 0, 0, 0] [1, 32, 16, 2] [1, 1, 1, 1] : memref<16x32x16x2xbf16> to !memrefA
+ %subview_0 = memref.subview %arg1[%arg3, 0, 0, 0] [1, 16, 32, 2] [1, 1, 1, 1] : memref<16x16x32x2xbf16> to !memrefB
+
+ %7 = vector.transfer_read %subview[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} : !memrefA, !vecAB
+ %8 = vector.transfer_read %subview[%c0, %c16, %c0, %c0], %1 {in_bounds = [true, true, true, true]} : !memrefA, !vecAB
+ %9 = vector.transfer_read %subview_0[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} : !memrefB, !vecAB
+ %10 = vector.transfer_read %subview_0[%c0, %c0, %c16, %c0], %1 {in_bounds = [true, true, true, true]} : !memrefB, !vecAB
+
+ %11 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %7, %9, %arg4 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %7, %10, %arg5 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %8, %9, %arg6 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %8, %10, %arg7 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ scf.yield %11, %12, %13, %14 : !vecC, !vecC, !vecC, !vecC
+ }
+ vector.transfer_write %6#3, %arg2[%c16, %c16] {in_bounds = [true, true]} : !vecC, memref<32x32xf32>
+ vector.transfer_write %6#2, %arg2[%c16, %c0] {in_bounds = [true, true]} : !vecC, memref<32x32xf32>
+ vector.transfer_write %6#1, %arg2[%c0, %c16] {in_bounds = [true, true]} : !vecC, memref<32x32xf32>
+ vector.transfer_write %6#0, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, memref<32x32xf32>
+ %alloc = memref.alloc() : memref<32x32xf32>
+ memref.copy %arg2, %alloc : memref<32x32xf32> to memref<32x32xf32>
+ return %alloc : memref<32x32xf32>
+}
+
+// CHECK-LABEL: @brgemm_bf16_with_cano
+// CHECK-1: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
+// CHECK-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
+// CHECK-4: x86.amx.tile_load
+// CHECK-4: x86.amx.tile_mulf
+// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
+// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
+// CHECK-NOT: vector.contract
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -728,6 +806,7 @@ func.func @online_packing_int8_matmul_loop(%arg0: memref<64x256xi8>, %arg1: memr
%c128 = arith.constant 128 : index
%c256 = arith.constant 256 : index
%c32 = arith.constant 32 : index
+ %bias = arith.constant dense<13> : !vecC
scf.for %arg3 = %c0 to %c64 step %c32 {
scf.for %arg4 = %c0 to %c128 step %c32 {
%subview = memref.subview %arg2[%arg3, %arg4] [32, 32] [1, 1] : memref<64x128xi32> to !memrefC
@@ -756,10 +835,16 @@ func.func @online_packing_int8_matmul_loop(%arg0: memref<64x256xi8>, %arg1: memr
%8, %10, %arg9 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
scf.yield %11, %12, %13, %14 : !vecC, !vecC, !vecC, !vecC
}
- vector.transfer_write %6#3, %subview[%c16, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
- vector.transfer_write %6#2, %subview[%c16, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
- vector.transfer_write %6#1, %subview[%c0, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
- vector.transfer_write %6#0, %subview[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ %7 = arith.addi %6#3, %bias : !vecC
+ %8 = arith.addi %6#2, %bias : !vecC
+ %9 = arith.addi %6#1, %bias : !vecC
+ %10 = arith.addi %6#0, %bias : !vecC
+
+ vector.transfer_write %7, %subview[%c16, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
+ vector.transfer_write %8, %subview[%c16, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+ vector.transfer_write %9, %subview[%c0, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
+ vector.transfer_write %10, %subview[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
}
}
%alloc = memref.alloc() : memref<64x128xi32>
@@ -777,6 +862,7 @@ func.func @online_packing_int8_matmul_loop(%arg0: memref<64x256xi8>, %arg1: memr
// CHECK: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xi32>, vector<16xi32>
// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xi32>, vector<16xi32>
+// CHECK-COUNT-4: arith.addi
// CHECK-NOT: scf.for {{.*}} vector<16x16xi32>, vector<16x16xi32>, vector<16x16xi32>, vector<16x16xi32>
// CHECK-NOT: vector.contract
More information about the Mlir-commits
mailing list