[Mlir-commits] [mlir] [mlir][x86] Accumulate AMX VNNI tiles in place (PR #233388)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 19 02:47:43 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Adam Siemieniuk (adam-smnk)
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/233388.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp (+109-4)
- (modified) mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir (+8-5)
``````````diff
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} {
``````````
</details>
https://github.com/llvm/llvm-project/pull/233388
More information about the Mlir-commits
mailing list