[Mlir-commits] [mlir] [mlir][x86] Add extra patterns for AMX nano-kernel lowering (acc with zero constant) (PR #222513)
Arun Thangamani
llvmlistbot at llvm.org
Tue Sep 22 07:00:44 PDT 2026
https://github.com/arun-thmn updated https://github.com/llvm/llvm-project/pull/222513
>From 8eda54c35b63138dfea81b24aa435e26740a90fc Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Wed, 9 Sep 2026 22:21:16 -0700
Subject: [PATCH 1/6] support zero constant acc for AMX nano-kernel lowering
---
.../VectorContractToAMXDotProduct.cpp | 152 ++++++++++++------
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 3 +
2 files changed, 109 insertions(+), 46 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 365a3fb8e523b..3ff7554b8fc25 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -27,6 +27,47 @@ using namespace mlir::x86;
namespace {
+static bool isZeroVectorConstant(mlir::Operation *op) {
+ auto constantOp = llvm::dyn_cast<mlir::arith::ConstantOp>(op);
+ if (!constantOp)
+ return false;
+
+ auto vectorType = llvm::dyn_cast<mlir::VectorType>(constantOp.getType());
+ if (!vectorType)
+ return false;
+
+ auto denseAttr =
+ llvm::dyn_cast<mlir::DenseElementsAttr>(constantOp.getValue());
+ if (!denseAttr || !denseAttr.isSplat())
+ return false;
+
+ mlir::Type elementType = vectorType.getElementType();
+
+ if (llvm::isa<mlir::FloatType>(elementType)) {
+ return denseAttr.getSplatValue<llvm::APFloat>().isZero();
+ }
+
+ if (llvm::isa<mlir::IntegerType>(elementType)) {
+ return denseAttr.getSplatValue<llvm::APInt>().isZero();
+ }
+
+ return false;
+}
+
+/*static bool isZeroVectorConstant(mlir::Operation *op) {
+ auto constantOp = llvm::dyn_cast<mlir::arith::ConstantOp>(op);
+ if (!constantOp)
+ return false;
+
+ auto denseAttr =
+ llvm::dyn_cast<mlir::DenseElementsAttr>(constantOp.getValue());
+ if (!denseAttr)
+ return false;
+
+ return denseAttr.isSplat() &&
+ denseAttr.getSplatValue<llvm::APFloat>().isZero();
+}*/
+
// Recursively follows single-use values through scf.yield operations
// and returns the first non-yield user result in the contraction chain.
static Value contractionUsersAfterYield(Value v) {
@@ -309,19 +350,19 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
vec2 = vector::ShapeCastOp::create(
rewriter, loc, VectorType::get((16 * offset), ipType), vec2);
- vector::ShuffleOp shuffle1;
- vector::ShuffleOp shuffle2;
+ vector::ShuffleOp writeVal1;
+ vector::ShuffleOp writeVal2;
if (ipType.isBF16()) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({(16 * offset)}, ipType), vec1,
vec2,
ArrayRef<int64_t>{0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9,
41, 10, 42, 11, 43, 16, 48, 17, 49, 18, 50,
19, 51, 24, 56, 25, 57, 26, 58, 27, 59});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({(16 * offset)}, ipType), vec1,
vec2,
ArrayRef<int64_t>{4, 36, 5, 37, 6, 38, 7, 39, 12, 44, 13,
@@ -332,7 +373,7 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
if (ipType.isSignlessInteger(8) || ipType.isF8E5M2() ||
ipType.isF8E4M3FN()) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({(16 * offset)}, ipType), vec1,
vec2,
ArrayRef<int64_t>{
@@ -342,7 +383,7 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
113, 18, 50, 82, 114, 19, 51, 83, 115, 24, 56, 88, 120,
25, 57, 89, 121, 26, 58, 90, 122, 27, 59, 91, 123});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({(16 * offset)}, ipType), vec1,
vec2,
ArrayRef<int64_t>{
@@ -357,9 +398,9 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
Value ivShuff1 = arith::DivUIOp::create(rewriter, loc, iv, cStep);
Value ivShuff2 = arith::AddIOp::create(rewriter, loc, ivShuff1, c16);
- vector::StoreOp::create(rewriter, loc, shuffle1, packedBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal1, packedBuffer,
ValueRange{indxToStoreInBuffer, ivShuff1, c0});
- vector::StoreOp::create(rewriter, loc, shuffle2, packedBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal2, packedBuffer,
ValueRange{indxToStoreInBuffer, ivShuff2, c0});
scf::YieldOp::create(nestedBuilder, loc);
@@ -802,6 +843,8 @@ struct VectorContractToAMXDotProduct
Operation *accReadOp =
traceToVectorReadLikeParentOperation(contractOp.getAcc());
+ bool isAccZeroVectorConstant = isZeroVectorConstant(accReadOp);
+
// Only the contract result's first consumer is needed, not the final
// store. This keeps the lowering independent of the epilogue ops (truncf,
// bias add, ReLU, ...) that sit between the contraction and the write.
@@ -955,31 +998,31 @@ struct VectorContractToAMXDotProduct
VectorType::get(16 * (blockingFactor / 2), ipType), subview,
range2);
- vector::ShuffleOp shuffle1;
- vector::ShuffleOp shuffle2;
+ vector::ShuffleOp writeVal1;
+ vector::ShuffleOp writeVal2;
if (blockingFactor == 2) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({16}, ipType), vec1, vec2,
ArrayRef<int64_t>{0, 16, 1, 17, 2, 18, 3, 19, 4, 20, 5, 21,
6, 22, 7, 23});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({16}, ipType), vec1, vec2,
ArrayRef<int64_t>{8, 24, 9, 25, 10, 26, 11, 27, 12, 28, 13,
29, 14, 30, 15, 31});
}
if (blockingFactor == 4) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({32}, ipType), vec1, vec2,
ArrayRef<int64_t>{0, 16, 32, 48, 1, 17, 33, 49,
2, 18, 34, 50, 3, 19, 35, 51,
4, 20, 36, 52, 5, 21, 37, 53,
6, 22, 38, 54, 7, 23, 39, 55});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({32}, ipType), vec1, vec2,
ArrayRef<int64_t>{8, 24, 40, 56, 9, 25, 41, 57,
10, 26, 42, 58, 11, 27, 43, 59,
@@ -990,9 +1033,9 @@ struct VectorContractToAMXDotProduct
auto rem = arith::DivUIOp::create(
rewriter, loc, rewriter.getIndexType(), iv, step);
- vector::StoreOp::create(rewriter, loc, shuffle1, packedBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal1, packedBuffer,
ValueRange{rem, c0});
- vector::StoreOp::create(rewriter, loc, shuffle2, packedBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal2, packedBuffer,
ValueRange{rem, nextStoreIndx});
scf::YieldOp::create(nestedBuilder, loc);
@@ -1069,12 +1112,13 @@ struct VectorContractToAMXDotProduct
loopLists.push_back(dyn_cast<scf::ForOp>(parent));
- if (accReadOp->getBlock() == parent->getBlock()) {
+ if (resultBlock == parent->getBlock()) {
break;
}
current = parent;
}
+
if (loopLists.size() > 2 || loopLists.size() == 0)
return rewriter.notifyMatchFailure(
contractOp, "Rewrite is supported until reduction loop depth of 2.");
@@ -1113,7 +1157,6 @@ struct VectorContractToAMXDotProduct
for (mlir::Operation &op : loopLists[0].getBody()->getOperations()) {
if (auto contract = llvm::dyn_cast<mlir::vector::ContractionOp>(op)) {
-
LogicalResult validate = validateContractOps(
rewriter, contract, dimValue, srcBuffLhs, srcBuffRhs, true);
@@ -1124,6 +1167,16 @@ struct VectorContractToAMXDotProduct
"the re-write conditions either the dimensions are "
"wrong or MemRef source are different or many users.");
+ Operation *contractReadOp =
+ traceToVectorReadLikeParentOperation(contract.getAcc());
+
+ if ((isAccZeroVectorConstant &&
+ !isZeroVectorConstant(contractReadOp)) ||
+ (!isAccZeroVectorConstant && isZeroVectorConstant(contractReadOp)))
+ return rewriter.notifyMatchFailure(
+ contractOp, "The input acc should be a zero constant or "
+ "transfer_read for all contracts");
+
ops.push_back(contract);
}
}
@@ -1418,11 +1471,14 @@ struct VectorContractToAMXDotProduct
Value srcBuffAcc;
SmallVector<Value> indicesAcc;
- llvm::TypeSwitch<Operation *>(accReadOp).Case<TransferReadOp, LoadOp>(
- [&](auto readOp) {
- srcBuffAcc = readOp.getOperand(0);
+ Operation *accWrite =
+ traceToVectorWriteLikeUserOperation(contractOp.getResult());
+
+ llvm::TypeSwitch<Operation *>(accWrite)
+ .Case<vector::TransferWriteOp, vector::StoreOp>([&](auto writeOp) {
+ srcBuffAcc = writeOp->getOperand(1);
- auto indices = readOp.getIndices();
+ auto indices = writeOp.getIndices();
indicesAcc.reserve(indices.size());
llvm::transform(indices, std::back_inserter(indicesAcc),
@@ -1470,50 +1526,54 @@ struct VectorContractToAMXDotProduct
vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
resultBuffer, ValueRange{iv, c16});
- Value shuffle1 = row;
- Value shuffle2 = row2;
+ Value writeVal1 = row;
+ Value writeVal2 = row2;
if (!isVnni) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = 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});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = 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;
- Value valueCRow1 =
- vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
- srcBuffAcc, indicesAcc);
- indicesAcc[indicesAcc.size() - 1] = c16;
+ if (!isAccZeroVectorConstant) {
+ indicesAcc[indicesAcc.size() - 2] = iv;
+ indicesAcc[indicesAcc.size() - 1] = c0;
- Value valueCRow2 =
- vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
- srcBuffAcc, indicesAcc);
+ Value valueCRow1 = vector::LoadOp::create(
+ rewriter, loc, VectorType::get(16, opType), srcBuffAcc,
+ indicesAcc);
+ indicesAcc[indicesAcc.size() - 1] = c16;
- Value addOp;
- Value addOp2;
+ Value valueCRow2 = vector::LoadOp::create(
+ rewriter, loc, VectorType::get(16, opType), srcBuffAcc,
+ indicesAcc);
- if (ipType.isBF16() || ipType.isF8E5M2() || ipType.isF8E4M3FN()) {
- addOp = arith::AddFOp::create(rewriter, loc, shuffle1, valueCRow1);
+ if (ipType.isBF16() || ipType.isF8E5M2() || ipType.isF8E4M3FN()) {
+ writeVal1 =
+ arith::AddFOp::create(rewriter, loc, writeVal1, valueCRow1);
- addOp2 = arith::AddFOp::create(rewriter, loc, shuffle2, valueCRow2);
- }
+ writeVal2 =
+ arith::AddFOp::create(rewriter, loc, writeVal2, valueCRow2);
+ }
- if (ipType.isSignlessInteger(8)) {
- addOp = arith::AddIOp::create(rewriter, loc, shuffle1, valueCRow1);
+ if (ipType.isSignlessInteger(8)) {
+ writeVal1 =
+ arith::AddIOp::create(rewriter, loc, writeVal1, valueCRow1);
- addOp2 = arith::AddIOp::create(rewriter, loc, shuffle2, valueCRow2);
+ writeVal2 =
+ arith::AddIOp::create(rewriter, loc, writeVal2, valueCRow2);
+ }
}
- vector::StoreOp::create(rewriter, loc, addOp, resultBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal1, resultBuffer,
ValueRange{iv, c0});
- vector::StoreOp::create(rewriter, loc, addOp2, resultBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal2, resultBuffer,
ValueRange{iv, c16});
scf::YieldOp::create(nestedBuilder, loc);
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 093e97af0520e..af2993e7f313c 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
@@ -547,6 +547,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
+ transform.apply_patterns.x86.move_accumulator_for_contract_loop
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -685,6 +686,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
+ transform.apply_patterns.x86.move_accumulator_for_contract_loop
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
@@ -958,6 +960,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
+ transform.apply_patterns.x86.move_accumulator_for_contract_loop
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
>From 12a255aa4985f8c5965d8534943530587143fc03 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Wed, 9 Sep 2026 23:02:31 -0700
Subject: [PATCH 2/6] clean-up
---
.../VectorContractToAMXDotProduct.cpp | 29 +++++--------------
1 file changed, 8 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 3ff7554b8fc25..057acd2b9a545 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -27,47 +27,34 @@ using namespace mlir::x86;
namespace {
-static bool isZeroVectorConstant(mlir::Operation *op) {
- auto constantOp = llvm::dyn_cast<mlir::arith::ConstantOp>(op);
+// Return true if the operation is a constant dense element
+// attribute.
+static bool isZeroVectorConstant(Operation *op) {
+ auto constantOp = llvm::dyn_cast<arith::ConstantOp>(op);
if (!constantOp)
return false;
- auto vectorType = llvm::dyn_cast<mlir::VectorType>(constantOp.getType());
+ auto vectorType = llvm::dyn_cast<VectorType>(constantOp.getType());
if (!vectorType)
return false;
- auto denseAttr =
- llvm::dyn_cast<mlir::DenseElementsAttr>(constantOp.getValue());
+ auto denseAttr = llvm::dyn_cast<DenseElementsAttr>(constantOp.getValue());
if (!denseAttr || !denseAttr.isSplat())
return false;
mlir::Type elementType = vectorType.getElementType();
- if (llvm::isa<mlir::FloatType>(elementType)) {
+ if (llvm::isa<FloatType>(elementType)) {
return denseAttr.getSplatValue<llvm::APFloat>().isZero();
}
- if (llvm::isa<mlir::IntegerType>(elementType)) {
+ if (llvm::isa<IntegerType>(elementType)) {
return denseAttr.getSplatValue<llvm::APInt>().isZero();
}
return false;
}
-/*static bool isZeroVectorConstant(mlir::Operation *op) {
- auto constantOp = llvm::dyn_cast<mlir::arith::ConstantOp>(op);
- if (!constantOp)
- return false;
-
- auto denseAttr =
- llvm::dyn_cast<mlir::DenseElementsAttr>(constantOp.getValue());
- if (!denseAttr)
- return false;
-
- return denseAttr.isSplat() &&
- denseAttr.getSplatValue<llvm::APFloat>().isZero();
-}*/
-
// Recursively follows single-use values through scf.yield operations
// and returns the first non-yield user result in the contraction chain.
static Value contractionUsersAfterYield(Value v) {
>From e2e93e85568b7359fd5073e968310265d0f06715 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Fri, 11 Sep 2026 00:01:41 -0700
Subject: [PATCH 3/6] code simplification + test-case
---
.../VectorContractToAMXDotProduct.cpp | 192 ++++++++----------
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 82 ++++++++
2 files changed, 171 insertions(+), 103 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 057acd2b9a545..7dadc87f97eaa 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -30,29 +30,19 @@ namespace {
// Return true if the operation is a constant dense element
// attribute.
static bool isZeroVectorConstant(Operation *op) {
- auto constantOp = llvm::dyn_cast<arith::ConstantOp>(op);
+ auto constantOp = dyn_cast<arith::ConstantOp>(op);
if (!constantOp)
return false;
- auto vectorType = llvm::dyn_cast<VectorType>(constantOp.getType());
+ auto vectorType = dyn_cast<VectorType>(constantOp.getType());
if (!vectorType)
return false;
- auto denseAttr = llvm::dyn_cast<DenseElementsAttr>(constantOp.getValue());
+ auto denseAttr = dyn_cast<DenseElementsAttr>(constantOp.getValue());
if (!denseAttr || !denseAttr.isSplat())
return false;
- mlir::Type elementType = vectorType.getElementType();
-
- if (llvm::isa<FloatType>(elementType)) {
- return denseAttr.getSplatValue<llvm::APFloat>().isZero();
- }
-
- if (llvm::isa<IntegerType>(elementType)) {
- return denseAttr.getSplatValue<llvm::APInt>().isZero();
- }
-
- return false;
+ return isZeroIntegerOrFloat(denseAttr.getSplatValue<Attribute>());
}
// Recursively follows single-use values through scf.yield operations
@@ -76,8 +66,7 @@ static Value contractionUsersAfterYield(Value v) {
// 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,
- Value input) {
+static Value collapseInnerDims(OpBuilder &builder, Location loc, Value input) {
ShapedType inputType = cast<ShapedType>(input.getType());
int64_t firstDimToCollapse = inputType.getRank() - 2;
@@ -202,7 +191,7 @@ static LogicalResult validateContractOps(OpBuilder &rewriter,
// The Accumulator dims should be 16 or 1. Like <1x16x16> or <16x16>.
ArrayRef<int64_t> accShape = accTy.getShape();
- llvm::SmallVector<int64_t> nonUnitDimAcc;
+ SmallVector<int64_t> nonUnitDimAcc;
llvm::copy_if(accShape, std::back_inserter(nonUnitDimAcc),
[](int64_t dim) { return (dim != 16 && dim != 1); });
@@ -213,7 +202,7 @@ static LogicalResult validateContractOps(OpBuilder &rewriter,
// <16x16x4>. The vnni dims should be 2 or 4.
VectorType lhsTy = contractOp.getLhsType();
ArrayRef<int64_t> lhsShape = lhsTy.getShape();
- llvm::SmallVector<int64_t> nonUnitDimLhs;
+ SmallVector<int64_t> nonUnitDimLhs;
llvm::copy_if(lhsShape, std::back_inserter(nonUnitDimLhs),
[](int64_t dim) { return (dim != 16 && dim != 1); });
@@ -832,6 +821,9 @@ struct VectorContractToAMXDotProduct
bool isAccZeroVectorConstant = isZeroVectorConstant(accReadOp);
+ Operation *accWrite =
+ traceToVectorWriteLikeUserOperation(contractOp.getResult());
+
// Only the contract result's first consumer is needed, not the final
// store. This keeps the lowering independent of the epilogue ops (truncf,
// bias add, ReLU, ...) that sit between the contraction and the write.
@@ -1055,7 +1047,7 @@ struct VectorContractToAMXDotProduct
amx::TileStoreOp::create(rewriter, loc, resultBuffer, ValueRange{c0, c0},
dp);
- auto vectorType = mlir::VectorType::get({16, 16}, opType);
+ auto vectorType = VectorType::get({16, 16}, opType);
int64_t srcRank =
(dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
Value padding = ub::PoisonOp::create(rewriter, loc, opType);
@@ -1068,7 +1060,7 @@ struct VectorContractToAMXDotProduct
map, inBounds);
Value resultOp = contractionUsersAfterYield(contractOp.getResult());
- if (auto vecType = llvm::dyn_cast<VectorType>(resultOp.getType()))
+ if (auto vecType = dyn_cast<VectorType>(resultOp.getType()))
vecRow = vector::ShapeCastOp::create(rewriter, loc, vecType, vecRow);
rewriter.replaceAllUsesWith(resultOp, vecRow);
@@ -1141,9 +1133,9 @@ struct VectorContractToAMXDotProduct
// Retrive all the contaction operation within the loop.
SmallVector<vector::ContractionOp> ops;
- for (mlir::Operation &op : loopLists[0].getBody()->getOperations()) {
+ for (Operation &op : loopLists[0].getBody()->getOperations()) {
- if (auto contract = llvm::dyn_cast<mlir::vector::ContractionOp>(op)) {
+ if (auto contract = dyn_cast<vector::ContractionOp>(op)) {
LogicalResult validate = validateContractOps(
rewriter, contract, dimValue, srcBuffLhs, srcBuffRhs, true);
@@ -1157,9 +1149,7 @@ struct VectorContractToAMXDotProduct
Operation *contractReadOp =
traceToVectorReadLikeParentOperation(contract.getAcc());
- if ((isAccZeroVectorConstant &&
- !isZeroVectorConstant(contractReadOp)) ||
- (!isAccZeroVectorConstant && isZeroVectorConstant(contractReadOp)))
+ if (isAccZeroVectorConstant != isZeroVectorConstant(contractReadOp))
return rewriter.notifyMatchFailure(
contractOp, "The input acc should be a zero constant or "
"transfer_read for all contracts");
@@ -1232,10 +1222,9 @@ struct VectorContractToAMXDotProduct
bool isInnerLoopUBHasOddQuot = false;
int64_t ubVal = 16 * blockingFactor;
- mlir::Value ub = innerLoop.getUpperBound();
- if (auto constOp = ub.getDefiningOp<mlir::arith::ConstantOp>()) {
- if (auto intAttr =
- llvm::dyn_cast<mlir::IntegerAttr>(constOp.getValue())) {
+ Value ub = innerLoop.getUpperBound();
+ if (auto constOp = ub.getDefiningOp<arith::ConstantOp>()) {
+ if (auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue())) {
ubVal = intAttr.getInt();
}
}
@@ -1380,10 +1369,9 @@ struct VectorContractToAMXDotProduct
bool isInnerLoopUBHasOddQuot = false;
int64_t ubVal = 16 * blockingFactor;
- mlir::Value ub = innerLoop.getUpperBound();
- if (auto constOp = ub.getDefiningOp<mlir::arith::ConstantOp>()) {
- if (auto intAttr =
- llvm::dyn_cast<mlir::IntegerAttr>(constOp.getValue())) {
+ Value ub = innerLoop.getUpperBound();
+ if (auto constOp = ub.getDefiningOp<arith::ConstantOp>()) {
+ if (auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue())) {
ubVal = intAttr.getInt();
}
}
@@ -1458,9 +1446,6 @@ struct VectorContractToAMXDotProduct
Value srcBuffAcc;
SmallVector<Value> indicesAcc;
- Operation *accWrite =
- traceToVectorWriteLikeUserOperation(contractOp.getResult());
-
llvm::TypeSwitch<Operation *>(accWrite)
.Case<vector::TransferWriteOp, vector::StoreOp>([&](auto writeOp) {
srcBuffAcc = writeOp->getOperand(1);
@@ -1468,15 +1453,13 @@ struct VectorContractToAMXDotProduct
auto indices = writeOp.getIndices();
indicesAcc.reserve(indices.size());
- llvm::transform(indices, std::back_inserter(indicesAcc),
- [&](OpFoldResult ofr) {
- return mlir::getValueOrCreateConstantIndexOp(
- rewriter, loc, ofr);
- });
+ llvm::transform(
+ indices, std::back_inserter(indicesAcc), [&](OpFoldResult ofr) {
+ return getValueOrCreateConstantIndexOp(rewriter, loc, ofr);
+ });
});
- auto outputShapes =
- mlir::cast<mlir::MemRefType>(srcBuffAcc.getType()).getShape();
+ auto outputShapes = cast<MemRefType>(srcBuffAcc.getType()).getShape();
unsigned int M = outputShapes[outputShapes.size() - 2];
unsigned int N = outputShapes[outputShapes.size() - 1];
@@ -1499,72 +1482,75 @@ struct VectorContractToAMXDotProduct
auto one = arith::ConstantIndexOp::create(rewriter, loc, 1);
auto nBound = 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 to temp buffer.
- scf::ForOp::create(
- rewriter, loc, c0, nBound, 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 writeVal1 = row;
- Value writeVal2 = row2;
-
- if (!isVnni) {
- writeVal1 = 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});
-
- writeVal2 = 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});
- }
+ if (!isVnni || !isAccZeroVectorConstant) {
+ // Create a loop that iterates over the MxN memerf, retrives two rows +
+ // shuffle them, add up the C element values and stores them to temp
+ // buffer.
+ scf::ForOp::create(
+ rewriter, loc, c0, nBound, 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 writeVal1 = row;
+ Value writeVal2 = row2;
+
+ if (!isVnni) {
+ writeVal1 = 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});
+
+ writeVal2 = 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});
+ }
- if (!isAccZeroVectorConstant) {
- indicesAcc[indicesAcc.size() - 2] = iv;
- indicesAcc[indicesAcc.size() - 1] = c0;
+ if (!isAccZeroVectorConstant) {
+ 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);
- if (ipType.isBF16() || ipType.isF8E5M2() || ipType.isF8E4M3FN()) {
- writeVal1 =
- arith::AddFOp::create(rewriter, loc, writeVal1, valueCRow1);
+ if (ipType.isBF16() || ipType.isF8E5M2() || ipType.isF8E4M3FN()) {
+ writeVal1 =
+ arith::AddFOp::create(rewriter, loc, writeVal1, valueCRow1);
- writeVal2 =
- arith::AddFOp::create(rewriter, loc, writeVal2, valueCRow2);
- }
+ writeVal2 =
+ arith::AddFOp::create(rewriter, loc, writeVal2, valueCRow2);
+ }
- if (ipType.isSignlessInteger(8)) {
- writeVal1 =
- arith::AddIOp::create(rewriter, loc, writeVal1, valueCRow1);
+ if (ipType.isSignlessInteger(8)) {
+ writeVal1 =
+ arith::AddIOp::create(rewriter, loc, writeVal1, valueCRow1);
- writeVal2 =
- arith::AddIOp::create(rewriter, loc, writeVal2, valueCRow2);
+ writeVal2 =
+ arith::AddIOp::create(rewriter, loc, writeVal2, valueCRow2);
+ }
}
- }
- vector::StoreOp::create(rewriter, loc, writeVal1, resultBuffer,
- ValueRange{iv, c0});
- vector::StoreOp::create(rewriter, loc, writeVal2, resultBuffer,
- ValueRange{iv, c16});
+ vector::StoreOp::create(rewriter, loc, writeVal1, resultBuffer,
+ ValueRange{iv, c0});
+ vector::StoreOp::create(rewriter, loc, writeVal2, resultBuffer,
+ ValueRange{iv, c16});
- scf::YieldOp::create(nestedBuilder, loc);
- });
+ scf::YieldOp::create(nestedBuilder, loc);
+ });
+ }
SmallVector<Value> writeResults;
for (unsigned int i = 0; i < M; i = i + 16) {
@@ -1572,7 +1558,7 @@ struct VectorContractToAMXDotProduct
Value indexOp_i = arith::ConstantIndexOp::create(rewriter, loc, i);
Value indexOp_j = arith::ConstantIndexOp::create(rewriter, loc, j);
- auto vectorType = mlir::VectorType::get({16, 16}, opType);
+ auto vectorType = VectorType::get({16, 16}, opType);
int64_t srcRank =
(dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
@@ -1594,9 +1580,9 @@ struct VectorContractToAMXDotProduct
Value vecRow = writeResults[i];
Value resultWriteOp = contractionUsersAfterYield(contOp.getResult());
- if (auto vecType = llvm::dyn_cast<VectorType>(resultWriteOp.getType()))
- vecRow = mlir::vector::ShapeCastOp::create(rewriter, loc, vecType,
- writeResults[i]);
+ if (auto vecType = dyn_cast<VectorType>(resultWriteOp.getType()))
+ vecRow = vector::ShapeCastOp::create(rewriter, loc, vecType,
+ writeResults[i]);
rewriter.replaceAllUsesWith(resultWriteOp, vecRow);
}
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 af2993e7f313c..6154e66a72101 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
@@ -2332,6 +2332,88 @@ func.func @negative_acc_tensor(
// CHECK-NOT: x86.amx
// CHECK: 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<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @negative_mixed_acc(%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>) {
+ %cst_1 = arith.constant dense<0> : vector<16x16xi32>
+ %0 = ub.poison : i32
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xi32> to !memrefC
+ %2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %4:2 = scf.for %arg5 = %c0 to %c64 step %c16 iter_args(%arg6 = %2, %arg7 = %cst_1) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
+ memref<64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+ %add = arith.addi %4#1, %3 : vector<16x16xi32>
+ vector.transfer_write %add, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @negative_mixed_acc
+// CHECK-NOT: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK-NOT: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK-NOT: x86.amx.tile_load
+// CHECK-NOT: x86.amx.tile_muli
+// CHECK-NOT: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
+
+
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
>From 4825733e88f46163e53a978afac98a52c3524346 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 22 Sep 2026 06:01:30 -0700
Subject: [PATCH 4/6] code simplification and extend test-cases
---
.../VectorContractToAMXDotProduct.cpp | 3 +-
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 87 ++++++++++++++++++-
2 files changed, 86 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 7dadc87f97eaa..95fbe621a8b25 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -34,8 +34,7 @@ static bool isZeroVectorConstant(Operation *op) {
if (!constantOp)
return false;
- auto vectorType = dyn_cast<VectorType>(constantOp.getType());
- if (!vectorType)
+ if (!isa<VectorType>(constantOp.getType()))
return false;
auto denseAttr = dyn_cast<DenseElementsAttr>(constantOp.getValue());
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 6154e66a72101..40bd8037a100e 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
@@ -686,7 +686,91 @@ 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.move_accumulator_for_contract_loop
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @matmul_int8_loop_cst_acc(%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>) {
+ %cst = arith.constant dense<0> : vector<16x16xi32>
+ %0 = ub.poison : i32
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xi32> to !memrefC
+ %2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %4:2 = scf.for %arg5 = %c0 to %c64 step %c16 iter_args(%arg6 = %cst, %arg7 = %cst) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
+ memref<64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+
+ %5 = arith.addi %4#1, %3 : vector<16x16xi32>
+ vector.transfer_write %5, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ %6 = arith.addi %4#0, %2 : vector<16x16xi32>
+ vector.transfer_write %6, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @matmul_int8_loop_cst_acc
+// CHECK-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK-3: x86.amx.tile_load
+// CHECK-2: x86.amx.tile_muli
+// CHECK: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
+// CHECK-NOT: scf.for {{.*}} vector<16x16xi32>, vector<16x16xi32>
+// 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
@@ -960,7 +1044,6 @@ 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.move_accumulator_for_contract_loop
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
>From 0de69b7962aaabf240966303d5b732be28a538ba Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 22 Sep 2026 06:19:49 -0700
Subject: [PATCH 5/6] simplify test-cases
---
mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir | 1 -
1 file changed, 1 deletion(-)
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 40bd8037a100e..81599b69ba545 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
@@ -547,7 +547,6 @@ 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.move_accumulator_for_contract_loop
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
} : !transform.any_op
transform.yield
>From 371f4364574158b0e5eb15788b6632f636e58c08 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 22 Sep 2026 07:00:26 -0700
Subject: [PATCH 6/6] update insertion point if acc memref after contract loop
---
.../X86/Transforms/VectorContractToAMXDotProduct.cpp | 9 +++++++++
1 file changed, 9 insertions(+)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 95fbe621a8b25..f6912d7119fa2 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -1476,6 +1476,15 @@ struct VectorContractToAMXDotProduct
k++;
}
}
+
+ // Update insertion point, if acc memref after contraction loop.
+ Operation *memrefOp = srcBuffAcc.getDefiningOp();
+ if (memrefOp && outerLoop &&
+ !(memrefOp->isBeforeInBlock(outerLoop.getOperation()))) {
+ rewriter.setInsertionPoint(memrefOp);
+ loc = memrefOp->getLoc();
+ }
+
auto c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
auto c16 = arith::ConstantIndexOp::create(rewriter, loc, 16);
auto one = arith::ConstantIndexOp::create(rewriter, loc, 1);
More information about the Mlir-commits
mailing list