[Mlir-commits] [mlir] [mlir][x86] Make AMX nanokernel independent of result write (PR #236699)
Julian Oppermann
llvmlistbot at llvm.org
Thu Oct 29 03:54:24 PDT 2026
https://github.com/jopperm updated https://github.com/llvm/llvm-project/pull/236699
>From 8774781520e830269a4345e97a47ddfccf9bd672 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Thu, 29 Oct 2026 02:03:02 -0700
Subject: [PATCH 1/2] [mlir][x86] Make AMX nanokernel independent of result
write
---
.../VectorContractToAMXDotProduct.cpp | 115 +++++++++---------
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 111 +++++++++++++++--
2 files changed, 162 insertions(+), 64 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 9714a2e36a893d..bece3235802635 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -591,34 +591,45 @@ getAccTileAccess(Operation *op) {
// Returns the accumulator read and result write for each contraction if the
// accumulator can be updated in place through direct AMX tile accesses.
-static std::optional<SmallVector<std::pair<Operation *, Operation *>>>
-getInPlaceAccAccesses(ArrayRef<vector::ContractionOp> ops) {
- SmallVector<std::pair<Operation *, Operation *>> accesses;
+// If writeOnly is true, only the result writes are considered, and the read
+// accesses are ignored.
+static std::pair<std::optional<SmallVector<Operation *>>,
+ std::optional<SmallVector<Operation *>>>
+getInPlaceAccAccesses(ArrayRef<vector::ContractionOp> ops, bool writeOnly) {
+ SmallVector<Operation *> reads, writes;
for (vector::ContractionOp op : ops) {
- Operation *read = traceToVectorReadLikeParentOperation(op.getAcc());
Value result = contractionUsersAfterYield(op.getResult());
- if (!read || !result || !result.hasOneUse())
- return std::nullopt;
+ if (!result || !result.hasOneUse())
+ return {std::nullopt, std::nullopt};
Operation *write = *result.user_begin();
if (!isa<vector::TransferWriteOp, vector::StoreOp>(write) ||
- write->getOperand(0) != result || !getAccTileAccess(read) ||
- !getAccTileAccess(write))
- return std::nullopt;
- accesses.emplace_back(read, write);
+ write->getOperand(0) != result || !getAccTileAccess(write))
+ return {std::nullopt, std::nullopt};
+ writes.push_back(write);
+
+ if (writeOnly)
+ continue;
+
+ Operation *read = traceToVectorReadLikeParentOperation(op.getAcc());
+ if (!read || !getAccTileAccess(read))
+ return {std::nullopt, std::nullopt};
+ reads.push_back(read);
}
- return accesses;
+
+ if (writeOnly)
+ return {std::nullopt, writes};
+ return std::make_pair(reads, writes);
}
// Loads the initial accumulator tiles before the given loop.
-static SmallVector<Value>
-createAccTileLoads(OpBuilder &rewriter, Location loc, Type opType,
- scf::ForOp loop,
- ArrayRef<std::pair<Operation *, Operation *>> accesses) {
+static SmallVector<Value> createAccTileLoads(OpBuilder &rewriter, Location loc,
+ Type opType, scf::ForOp loop,
+ ArrayRef<Operation *> reads) {
OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPoint(loop);
auto tileType = amx::TileType::get({16, 16}, opType);
SmallVector<Value> tiles;
- for (auto [read, write] : accesses) {
+ for (Operation *read : reads) {
auto [base, indices] = *getAccTileAccess(read);
tiles.push_back(
amx::TileLoadOp::create(rewriter, loc, tileType, base, indices));
@@ -936,9 +947,6 @@ struct VectorContractToAMXDotProduct
Operation *accReadOp =
traceToVectorReadLikeParentOperation(contractOp.getAcc());
- 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.
@@ -1229,10 +1237,6 @@ struct VectorContractToAMXDotProduct
return rewriter.notifyMatchFailure(
contractOp, "Rewrite is supported until reduction loop depth of 2.");
- if (isAccZeroVectorConstant && !accWrite)
- return rewriter.notifyMatchFailure(
- contractOp, "The result of the contract is not written to memory.");
-
auto srcIndxLhs = getSrcIndxValue(rewriter, contractOp.getLoc(),
contractOp.getLhs(), false);
if (failed(srcIndxLhs))
@@ -1340,13 +1344,14 @@ struct VectorContractToAMXDotProduct
// For VNNI layout, accumulator tiles stay in their original layout.
// Load and and store them back in place when possible.
// Otherwise, start from zero tile and accumulate only the final result.
- std::optional<SmallVector<std::pair<Operation *, Operation *>>> inPlaceAcc;
- if (isVnni && !isAccZeroVectorConstant)
- inPlaceAcc = getInPlaceAccAccesses(ops);
+ std::optional<SmallVector<Operation *>> inPlaceAccReads, inPlaceAccWrites;
+ if (isVnni)
+ std::tie(inPlaceAccReads, inPlaceAccWrites) =
+ getInPlaceAccAccesses(ops, /*writeOnly=*/isAccZeroVectorConstant);
auto createInitTiles = [&](scf::ForOp loop) {
- if (inPlaceAcc)
+ if (inPlaceAccReads)
return createAccTileLoads(rewriter, loop.getLoc(), opType, loop,
- *inPlaceAcc);
+ *inPlaceAccReads);
return createTileZeros(rewriter, loop.getLoc(), opType, loop, ops.size());
};
@@ -1619,11 +1624,11 @@ struct VectorContractToAMXDotProduct
// initial accumulation value, and store back to the C-Matrix
Location loc = outerLoop.getLoc();
- if (inPlaceAcc) {
- // The tiles already include the initial accumulator: store each one
- // where its contraction result was written.
- for (auto [i, access] : llvm::enumerate(*inPlaceAcc)) {
- Operation *write = access.second;
+ if (inPlaceAccWrites) {
+ // The tiles already include the initial accumulator or were
+ // zero-initialized: store each one where its contraction result was
+ // written.
+ for (auto [i, write] : llvm::enumerate(*inPlaceAccWrites)) {
auto [base, indices] = *getAccTileAccess(write);
rewriter.setInsertionPoint(write);
amx::TileStoreOp::create(rewriter, write->getLoc(), base, indices,
@@ -1633,23 +1638,18 @@ struct VectorContractToAMXDotProduct
return success();
}
+ // Carried-in accumulators are added to the results instead.
+ bool addAccFromMemory = !isAccZeroVectorConstant && carriedAccs.empty();
+
+ // Direct tile stores weren't possible, so we need a temporary buffer to
+ // move tile register contents through memory.
+ unsigned int M, N;
Value srcBuffAcc;
SmallVector<Value> indicesAcc;
- if (isAccZeroVectorConstant) {
- llvm::TypeSwitch<Operation *>(accWrite)
- .Case<vector::TransferWriteOp, vector::StoreOp>([&](auto writeOp) {
- srcBuffAcc = writeOp->getOperand(1);
-
- auto indices = writeOp.getIndices();
- indicesAcc.reserve(indices.size());
-
- llvm::transform(
- indices, std::back_inserter(indicesAcc), [&](OpFoldResult ofr) {
- return getValueOrCreateConstantIndexOp(rewriter, loc, ofr);
- });
- });
- } else {
+ if (addAccFromMemory) {
+ // We accumulate onto a register tile in memory, hence we let the
+ // temporary buffer match its shape.
llvm::TypeSwitch<Operation *>(accReadOp).Case<TransferReadOp, LoadOp>(
[&](auto readOp) {
srcBuffAcc = readOp.getOperand(0);
@@ -1663,12 +1663,18 @@ struct VectorContractToAMXDotProduct
rewriter, loc, ofr);
});
});
+ auto outputShapes = cast<MemRefType>(srcBuffAcc.getType()).getShape();
+ M = outputShapes[outputShapes.size() - 2];
+ N = outputShapes[outputShapes.size() - 1];
+ } else {
+ // Otherwise we can just pick a suitable shape to hold all result tiles.
+ // For the VNNI-packed case, we pass the result from each tile register to
+ // memory individually. For the online packing case, we ensure two tiles
+ // are neighboring in the N-dimension to facilitate the shuffling.
+ M = isVnni ? ops.size() * 16 : ops.size() / 2 * 16;
+ N = isVnni ? 16 : 2 * 16;
}
- auto outputShapes = cast<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);
@@ -1687,17 +1693,14 @@ struct VectorContractToAMXDotProduct
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 nBound = arith::ConstantIndexOp::create(rewriter, loc, N);
-
- // Carried-in accumulators are added to the results instead.
- bool addAccFromMemory = !isAccZeroVectorConstant && carriedAccs.empty();
+ auto mBound = arith::ConstantIndexOp::create(rewriter, loc, M);
if (!isVnni || addAccFromMemory) {
// 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{},
+ rewriter, loc, c0, mBound, one, ValueRange{},
[&](OpBuilder &nestedBuilder, Location loc, Value iv,
ValueRange iterArgs) {
auto row = vector::LoadOp::create(rewriter, 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 125ea77937d5e4..61b86a510f4d51 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
@@ -1271,6 +1271,92 @@ module attributes {transform.with_named_sequence} {
// -----
+// Zero-initialized VNNI contractions store a 1x3 tile grid directly to the
+// output, preserving the destination of each contraction rather than using
+// contraction order as the tile order in memory.
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x48x4xi8, strided<[192, 4, 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_zero_init_direct_store(%lhs: memref<16x64x4xi8>, %rhs: memref<64x48x4xi8>, %out: memref<16x48xi32>) {
+ %zero = arith.constant dense<0> : !vecC
+ %padding = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c64 = arith.constant 64 : index
+ %tiles:3 = scf.for %k = %c0 to %c64 step %c16 iter_args(%acc0 = %zero, %acc1 = %zero, %acc2 = %zero) -> (!vecC, !vecC, !vecC) {
+ %lhs_tile = memref.subview %lhs[0, %k, 0] [16, 16, 4] [1, 1, 1] :
+ memref<16x64x4xi8> to !memrefA
+ %rhs_tiles = memref.subview %rhs[%k, 0, 0] [16, 48, 4] [1, 1, 1] :
+ memref<64x48x4xi8> to !memrefB
+ %a = vector.transfer_read %lhs_tile[%c0, %c0, %c0], %padding {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %b0 = vector.transfer_read %rhs_tiles[%c0, %c16, %c0], %padding {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %b1 = vector.transfer_read %rhs_tiles[%c0, %c32, %c0], %padding {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %b2 = vector.transfer_read %rhs_tiles[%c0, %c0, %c0], %padding {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+
+ %res0 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %a, %b0, %acc0 : !vecAB, !vecAB into !vecC
+ %res1 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %a, %b1, %acc1 : !vecAB, !vecAB into !vecC
+ %res2 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %a, %b2, %acc2 : !vecAB, !vecAB into !vecC
+
+ scf.yield %res0, %res1, %res2 : !vecC, !vecC, !vecC
+ }
+
+ vector.transfer_write %tiles#0, %out[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, memref<16x48xi32>
+ vector.transfer_write %tiles#1, %out[%c0, %c32] {in_bounds = [true, true]} :
+ !vecC, memref<16x48xi32>
+ vector.transfer_write %tiles#2, %out[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, memref<16x48xi32>
+
+ return
+}
+
+// CHECK-LABEL: @matmul_int8_loop_zero_init_direct_store(
+// CHECK-SAME: %[[LHS:[^:]+]]: memref<16x64x4xi8>, %[[RHS:[^:]+]]: memref<64x48x4xi8>, %[[OUT:[^:]+]]: memref<16x48xi32>
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-NOT: memref.alloca
+// CHECK-COUNT-3: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK: %[[TILES:.+]]:3 = scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK-COUNT-3: x86.amx.tile_muli
+// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
+// CHECK-NEXT: }
+// CHECK-NEXT: x86.amx.tile_store %[[OUT]][%[[C0]], %[[C16]]], %[[TILES]]#0
+// CHECK-NEXT: x86.amx.tile_store %[[OUT]][%[[C0]], %[[C32]]], %[[TILES]]#1
+// CHECK-NEXT: x86.amx.tile_store %[[OUT]][%[[C0]], %[[C0]]], %[[TILES]]#2
+// CHECK-NEXT: return
+// 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
+ }
+}
+
+// -----
+
// The contraction results are carried through another loop before being
// written.
@@ -1443,7 +1529,7 @@ func.func @matmul_int8_loop_result_through_loop_extra_args(
// CHECK-LABEL: @matmul_int8_loop_result_through_loop_extra_args
// CHECK: %[[TILES:.+]]:2 = scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
-// CHECK: %[[BUF:.+]] = memref.alloca() : memref<16x32xi32>
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<32x16xi32>
// CHECK: x86.amx.tile_store %[[BUF]][%{{.+}}, %{{.+}}], %[[TILES]]#0
// CHECK: x86.amx.tile_store %[[BUF]][%{{.+}}, %{{.+}}], %[[TILES]]#1
// CHECK: %[[RES0:.+]] = vector.transfer_read %[[BUF]]
@@ -4049,8 +4135,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// The zero initialized accumulation results are not written to memory: not
-// supported.
+// The zero initialized accumulation results are not written to memory.
!vecAB = vector<16x16x4xi8>
!vecC = vector<16x16xi32>
@@ -4061,7 +4146,7 @@ module attributes {transform.with_named_sequence} {
#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
-func.func @negative_matmul_int8_loop_result_not_written(
+func.func @matmul_int8_loop_result_not_written(
%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>) -> (!vecC, !vecC) {
%cst = arith.constant dense<0> : !vecC
%1 = ub.poison : i8
@@ -4095,10 +4180,20 @@ func.func @negative_matmul_int8_loop_result_not_written(
return %4#0, %4#1 : !vecC, !vecC
}
-// CHECK-LABEL: @negative_matmul_int8_loop_result_not_written
-// CHECK-NOT: x86.amx
-// CHECK-COUNT-2: vector.contract
-// CHECK-NOT: x86.amx
+// CHECK-LABEL: @matmul_int8_loop_result_not_written
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-COUNT-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK: %[[TILES:.+]]:2 = scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK-COUNT-2: x86.amx.tile_muli
+// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<32x16xi32>
+// CHECK-NEXT: x86.amx.tile_store %[[BUF]][%[[C0]], %[[C0]]], %[[TILES]]#0
+// CHECK-NEXT: x86.amx.tile_store %[[BUF]][%[[C16]], %[[C0]]], %[[TILES]]#1
+// CHECK-NEXT: %[[RES0:.+]] = vector.transfer_read %[[BUF]][%[[C0]], %[[C0]]]
+// CHECK-NEXT: %[[RES1:.+]] = vector.transfer_read %[[BUF]][%[[C16]], %[[C0]]]
+// CHECK-NEXT: return %[[RES0]], %[[RES1]] : 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}) {
>From 7a6a869ef62135929f680603bc967550e8477fa2 Mon Sep 17 00:00:00 2001
From: Julian Oppermann <julian.oppermann at intel.com>
Date: Thu, 29 Oct 2026 03:53:51 -0700
Subject: [PATCH 2/2] Fix typos in vicinity
---
.../X86/Transforms/VectorContractToAMXDotProduct.cpp | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index bece3235802635..600686258d9886 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -1342,7 +1342,7 @@ struct VectorContractToAMXDotProduct
scf::ForOp outerLoop;
// For VNNI layout, accumulator tiles stay in their original layout.
- // Load and and store them back in place when possible.
+ // Load and store them back in place when possible.
// Otherwise, start from zero tile and accumulate only the final result.
std::optional<SmallVector<Operation *>> inPlaceAccReads, inPlaceAccWrites;
if (isVnni)
@@ -1696,8 +1696,8 @@ struct VectorContractToAMXDotProduct
auto mBound = arith::ConstantIndexOp::create(rewriter, loc, M);
if (!isVnni || addAccFromMemory) {
- // 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
+ // Create a loop that iterates over the MxN memref, retrieves two rows +
+ // shuffles them, adds up the C element values and stores them to the temp
// buffer.
scf::ForOp::create(
rewriter, loc, c0, mBound, one, ValueRange{},
More information about the Mlir-commits
mailing list