[Mlir-commits] [mlir] [mlir][x86] Accumulate AMX VNNI tiles in place (PR #233388)
Adam Siemieniuk
llvmlistbot at llvm.org
Mon Oct 19 02:47:09 PDT 2026
https://github.com/adam-smnk created https://github.com/llvm/llvm-project/pull/233388
Accumulate AMX VNNI contraction results in place instead of through temporary storage when the accumulator and result support direct tile accesses.
This reduces memory traffic and improves performance for contractions in VNNI layout with a non-zero accumulator.
Assisted-by: Claude
>From 07c7f05411b626f4b9a9157b0d0782bb32535ef6 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Wed, 7 Oct 2026 11:38:33 -0700
Subject: [PATCH] [mlir][x86] Accumulate AMX VNNI tiles in place
Accumulate AMX VNNI contraction results in place instead of through
temporary storage when the accumulator and result support direct tile
accesses.
This reduces memory traffic and improves performance for contractions
in VNNI layout with a non-zero accumulator.
Assisted-by: Claude
---
.../VectorContractToAMXDotProduct.cpp | 113 +++++++++++++++++-
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 13 +-
2 files changed, 117 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 3dfdf1f09a07c..8e99b52dad5eb 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -536,6 +536,7 @@ createTiledDp(OpBuilder &rewriter, Location loc,
static SmallVector<Value> createTileZeros(OpBuilder &rewriter, Location loc,
Type opType, scf::ForOp outerLoop,
int64_t size) {
+ OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPoint(outerLoop);
SmallVector<Value> loopItrArgs;
@@ -548,6 +549,82 @@ static SmallVector<Value> createTileZeros(OpBuilder &rewriter, Location loc,
return loopItrArgs;
}
+// Find memref source and indices of the accumulator tile for an AMX
+// contraction. Returns nullopt if the access is not suitable for AMX tile
+// load/store.
+static std::optional<std::pair<Value, SmallVector<Value>>>
+getAccTileAccess(Operation *op) {
+ if (!op)
+ return std::nullopt;
+
+ Value base;
+ VectorType vecTy;
+ ValueRange indices;
+ if (auto transfer = dyn_cast<VectorTransferOpInterface>(op)) {
+ if (transfer.getMask() || !transfer.getPermutationMap().isMinorIdentity())
+ return std::nullopt;
+ base = transfer.getBase();
+ vecTy = transfer.getVectorType();
+ indices = transfer.getIndices();
+ } else if (auto load = dyn_cast<vector::LoadOp>(op)) {
+ base = load.getBase();
+ vecTy = load.getVectorType();
+ indices = load.getIndices();
+ } else if (auto store = dyn_cast<vector::StoreOp>(op)) {
+ base = store.getBase();
+ vecTy = store.getVectorType();
+ indices = store.getIndices();
+ } else {
+ return std::nullopt;
+ }
+
+ auto memrefTy = dyn_cast<MemRefType>(base.getType());
+ ArrayRef<int64_t> shape = vecTy.getShape();
+ if (!memrefTy || memrefTy.getRank() < 2 || !memrefTy.isLastDimUnitStride() ||
+ shape.size() < 2 || shape.end()[-1] != 16 || shape.end()[-2] != 16 ||
+ llvm::any_of(shape.drop_back(2), [](int64_t d) { return d != 1; }))
+ return std::nullopt;
+
+ return std::make_pair(base, SmallVector<Value>(indices));
+}
+
+// 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;
+ for (vector::ContractionOp op : ops) {
+ Operation *read = traceToVectorReadLikeParentOperation(op.getAcc());
+ Value result = contractionUsersAfterYield(op.getResult());
+ if (!read || !result || !result.hasOneUse())
+ return 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);
+ }
+ return accesses;
+}
+
+// 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) {
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(loop);
+ auto tileType = amx::TileType::get({16, 16}, opType);
+ SmallVector<Value> tiles;
+ for (auto [read, write] : accesses) {
+ auto [base, indices] = *getAccTileAccess(read);
+ tiles.push_back(
+ amx::TileLoadOp::create(rewriter, loc, tileType, base, indices));
+ }
+ return tiles;
+}
+
static Value getIndxToLoadStoreFromPckBuffer(OpBuilder &rewriter, Location loc,
Value ivInnerLoop,
Value ivOuterLoop,
@@ -1259,6 +1336,19 @@ struct VectorContractToAMXDotProduct
scf::ForOp innerLoop;
scf::ForOp outerLoop;
+ // 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);
+ auto createInitTiles = [&](scf::ForOp loop) {
+ if (inPlaceAcc)
+ return createAccTileLoads(rewriter, loop.getLoc(), opType, loop,
+ *inPlaceAcc);
+ return createTileZeros(rewriter, loop.getLoc(), opType, loop, ops.size());
+ };
+
scf::ForOp newLoop;
// Case 2a: Reduction loop depth is 2.
if (loopLists.size() == 2) {
@@ -1280,8 +1370,8 @@ struct VectorContractToAMXDotProduct
contractOp, "Invalid loop step. The step should be 32 for BF16 and "
"64 for Int8/F8.");
- SmallVector<Value> loopItrArgs = createTileZeros(
- rewriter, outerLoop.getLoc(), opType, outerLoop, ops.size());
+ SmallVector<Value> loopItrArgs = createInitTiles(outerLoop);
+ rewriter.setInsertionPoint(outerLoop);
if (isVnni) {
newLoop = scf::ForOp::create(
@@ -1436,8 +1526,8 @@ struct VectorContractToAMXDotProduct
"Invalid loop step. The step should be 32 for BF16 and "
"64 for Int8/F8 or 1 if it is rduction loop other than K.");
- SmallVector<Value> loopItrArgs = createTileZeros(
- rewriter, innerLoop.getLoc(), opType, innerLoop, ops.size());
+ SmallVector<Value> loopItrArgs = createInitTiles(innerLoop);
+ rewriter.setInsertionPoint(innerLoop);
if (isVnni) {
newLoop = createLoops(
@@ -1527,6 +1617,21 @@ 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();
+
+ 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;
+ auto [base, indices] = *getAccTileAccess(write);
+ rewriter.setInsertionPoint(write);
+ amx::TileStoreOp::create(rewriter, write->getLoc(), base, indices,
+ newLoop.getResult(i));
+ rewriter.eraseOp(write);
+ }
+ return success();
+ }
+
Value srcBuffAcc;
SmallVector<Value> indicesAcc;
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 b71ad093a9ce1..125ea77937d5e 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
@@ -772,12 +772,15 @@ func.func @brgemm_bf16_loop(%arg0: memref<16x64x64x2xbf16>, %arg1: memref<16x64x
}
// CHECK-LABEL: @brgemm_bf16_loop
-// CHECK-2: 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-COUNT-2: scf.for
+// CHECK-COUNT-4: x86.amx.tile_load {{.*}} : memref<32x32xf32, strided<[128, 1], offset: ?>> into !x86.amx.tile<16x16xf32>
+// CHECK-COUNT-2: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
+// CHECK-DAG: x86.amx.tile_load
+// CHECK-DAG: 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: x86.amx.tile_zero
+// CHECK-COUNT-4: x86.amx.tile_store {{.*}} : memref<32x32xf32, strided<[128, 1], offset: ?>>, !x86.amx.tile<16x16xf32>
+// CHECK-NOT: vector.transfer_write
// CHECK-NOT: vector.contract
module attributes {transform.with_named_sequence} {
More information about the Mlir-commits
mailing list