[Mlir-commits] [mlir] [mlir][x86] Decouple Accumulator from ADD-Based Vector Contractions (over loops) (PR #204327)
Arun Thangamani
llvmlistbot at llvm.org
Fri Jul 31 07:18:07 PDT 2026
https://github.com/arun-thmn updated https://github.com/llvm/llvm-project/pull/204327
>From 23c66ee4340af485894825ac68c082ea2c28b07e Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Wed, 17 Jun 2026 02:48:59 -0700
Subject: [PATCH 1/7] does vector.contract(A, B, Acc) into vector.contract(A,
B, 0) + Acc
---
.../X86/TransformOps/X86TransformOps.td | 10 ++
mlir/include/mlir/Dialect/X86/Transforms.h | 5 +
.../include/mlir/Dialect/X86/Utils/X86Utils.h | 10 +-
.../X86/TransformOps/X86TransformOps.cpp | 5 +
.../lib/Dialect/X86/Transforms/CMakeLists.txt | 1 +
.../MoveAccumulatorForContractLoop.cpp | 114 +++++++++++++++
.../Transforms/VectorContractBF16ToFMA.cpp | 28 ++--
.../VectorContractToPackedTypeDotProduct.cpp | 30 ++--
mlir/lib/Dialect/X86/Utils/X86Utils.cpp | 63 +++++---
.../X86/move-acc-for-contract-loop.mlir | 136 ++++++++++++++++++
10 files changed, 351 insertions(+), 51 deletions(-)
create mode 100644 mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
create mode 100644 mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
diff --git a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
index c474cfb47d003..48aabf90914e2 100644
--- a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
+++ b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
@@ -82,5 +82,15 @@ def ApplyVectorContractToAMXDotProductPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplyMoveAccumulatorForContractLoopPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.x86.move_accumulator_for_contract_loop",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Rewrites an additive vector contraction by zeroing the accumulator input
+ and materializing the accumulation as a separate add after the contraction.
+ }];
+
+ let assemblyFormat = "attr-dict";
+}
#endif // X86_TRANSFORM_OPS
diff --git a/mlir/include/mlir/Dialect/X86/Transforms.h b/mlir/include/mlir/Dialect/X86/Transforms.h
index 6ebba5e94ec7c..8bcaf9a67feb0 100644
--- a/mlir/include/mlir/Dialect/X86/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86/Transforms.h
@@ -110,6 +110,11 @@ void populateShuffleVectorFMAOpsPatterns(RewritePatternSet &patterns);
// Int8).
void populateVectorContractToAMXDotProductPatterns(RewritePatternSet &patterns);
+// Rewrites an additive vector contraction by zeroing the accumulator input
+// and materializing the accumulation as a separate add after the contraction.
+void populateMoveAccumulatorForContractLoopPatterns(
+ RewritePatternSet &patterns);
+
//===----------------------------------------------------------------------===//
/// Helpers extracted from:
/// - clang/lib/Headers/avxintrin.h
diff --git a/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h b/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
index 2ff8f4f4283a2..f067ce3bb9b65 100644
--- a/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
+++ b/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
@@ -29,6 +29,10 @@ namespace x86 {
bool isInVnniLayout(Operation *op, llvm::ArrayRef<AffineMap> indexingMaps,
std::optional<unsigned> blockingFactor = std::nullopt);
+// Recursively follows single-use values through scf.yield operations
+// and returns the first non-yield user result in the contraction chain.
+Value contractionUsersAfterYield(Value v);
+
// Returns true if two contraction ops form a valid pair for VNNI packing.
// It verifies that both contractions share the appropriate operand, read from
// the same source buffer, and use constant indices that differ by 8 or 16.
@@ -59,9 +63,9 @@ LogicalResult shuffleAfterReadLikeOp(PatternRewriter &rewriter, Operation *opA,
// Shuffles vectors produced by vector.contraction ops into a flat layout
// before they are written to memory.
-LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter,
- Operation *opA, Operation *opB,
- int64_t nonUnitDimAcc, VectorType accTy);
+LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value opA,
+ Value opB, int64_t nonUnitDimAcc,
+ VectorType accTy);
} // namespace x86
} // namespace mlir
diff --git a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
index 390b21e12b0ed..a52f08a5cd061 100644
--- a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
+++ b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
@@ -52,6 +52,11 @@ void mlir::transform::ApplyVectorContractToAMXDotProductPatternsOp::
x86::populateVectorContractToAMXDotProductPatterns(patterns);
}
+void mlir::transform::ApplyMoveAccumulatorForContractLoopPatternsOp::
+ populatePatterns(RewritePatternSet &patterns) {
+ x86::populateMoveAccumulatorForContractLoopPatterns(patterns);
+}
+
//===----------------------------------------------------------------------===//
// Transform op registration
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt b/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
index 9c3695536cda9..5118563c03f16 100644
--- a/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
@@ -7,6 +7,7 @@ add_mlir_dialect_library(MLIRX86Transforms
SinkVectorProducerOps.cpp
ShuffleVectorFMAOps.cpp
VectorContractToAMXDotProduct.cpp
+ MoveAccumulatorForContractLoop.cpp
LINK_LIBS PUBLIC
MLIRArithDialect
diff --git a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
new file mode 100644
index 0000000000000..6f87ebfa441dc
--- /dev/null
+++ b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
@@ -0,0 +1,114 @@
+//===- MoveAccumulatorForContractLoop.cpp ---------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Dialect/X86/Transforms.h"
+#include "mlir/Dialect/X86/Utils/X86Utils.h"
+#include "mlir/Dialect/X86/X86Dialect.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Casting.h"
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+using namespace mlir::x86;
+
+namespace {
+// Transforms vector.contract(A, B, Acc) into vector.contract(A, B, 0) + Acc
+// to decouple the contraction computation from the accumulator update.
+struct MoveAccumulatorForContractLoop
+ : public OpRewritePattern<vector::ContractionOp> {
+ using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
+ PatternRewriter &rewriter) const override {
+
+ if (contractOp.getKind() != vector::CombiningKind::ADD)
+ return rewriter.notifyMatchFailure(contractOp,
+ "Expects add combining kind.");
+
+ Operation *accReadOp =
+ traceToVectorReadLikeParentOperation(contractOp.getAcc());
+
+ Value contractValue = contractionUsersAfterYield(contractOp.getResult());
+
+ if (!contractValue)
+ return rewriter.notifyMatchFailure(
+ contractOp, "Final acc write might have multiple users.");
+
+ Operation *resultWriteOp = *contractValue.getUsers().begin();
+
+ if (!accReadOp || !resultWriteOp)
+ return rewriter.notifyMatchFailure(
+ contractOp, "Read/write from/to acc matrix is not by "
+ "transfer_read/load/transfer_write/store ops.");
+
+ if (dyn_cast<arith::ConstantOp>(accReadOp))
+ return rewriter.notifyMatchFailure(
+ contractOp,
+ "The input acc to contract is already a constant zero vector.");
+ ;
+
+ if ((accReadOp->getBlock() == contractOp->getBlock()) ||
+ (resultWriteOp->getBlock() == contractOp->getBlock()))
+ return rewriter.notifyMatchFailure(
+ contractOp, "Acc read/write should be in a separate block.");
+
+ // Replace acc of a contraction operation with vector constant.
+ rewriter.setInsertionPointAfter(accReadOp);
+ Value accValue = accReadOp->getResult(0);
+ auto vecTy = llvm::dyn_cast<VectorType>(accValue.getType());
+ if (!vecTy)
+ return rewriter.notifyMatchFailure(contractOp, "Excepts vector type.");
+
+ Location loc = accReadOp->getLoc();
+ Type elemTy = vecTy.getElementType();
+
+ Value zeroScalar = arith::ConstantOp::create(rewriter, loc, elemTy,
+ rewriter.getZeroAttr(elemTy));
+
+ Value zeroVec =
+ vector::BroadcastOp::create(rewriter, loc, vecTy, zeroScalar);
+
+ accValue.replaceAllUsesWith(zeroVec);
+
+ // Adds the initial acc value with acontract results before storing to acc
+ // matrix.
+ rewriter.setInsertionPoint(resultWriteOp);
+ Location locUser = resultWriteOp->getLoc();
+
+ Value addition;
+
+ if (llvm::isa<FloatType>(elemTy)) {
+ addition =
+ arith::AddFOp::create(rewriter, locUser, contractValue, accValue);
+ }
+
+ if (llvm::isa<IntegerType>(elemTy)) {
+ addition =
+ arith::AddIOp::create(rewriter, locUser, contractValue, accValue);
+ }
+
+ resultWriteOp->replaceUsesOfWith(contractValue, addition);
+ return success();
+ }
+};
+
+} // namespace
+
+void x86::populateMoveAccumulatorForContractLoopPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<MoveAccumulatorForContractLoop>(patterns.getContext());
+}
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp
index 287892c3a660f..c70ad221c772d 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp
@@ -415,25 +415,21 @@ struct VectorContractBF16ToFMA
Operation *accReadOp1 =
traceToVectorReadLikeParentOperation(pairContractOp.getAcc());
- // Iterate down to find the users of contact operations until it is store
- // or transfer_write.
- Operation *resultWriteOp0 =
- traceToVectorWriteLikeUserOperation(contractOp.getResult());
- Operation *resultWriteOp1 =
- traceToVectorWriteLikeUserOperation(pairContractOp.getResult());
-
- // Shuffle the accumulators of the contract operations.
- LogicalResult readShuffle =
- shuffleAfterReadLikeOp(rewriter, accReadOp0, accReadOp1, contractOp,
- pairContractOp, nonUnitDim, accTy);
-
- if (failed(readShuffle))
- return rewriter.notifyMatchFailure(
- contractOp, "Accumulator read is not by transfer_read or load");
+ if (!(dyn_cast<arith::ConstantOp>(accReadOp0))) {
+ // Shuffle the accumulators of the contract operations.
+ LogicalResult readShuffle =
+ shuffleAfterReadLikeOp(rewriter, accReadOp0, accReadOp1, contractOp,
+ pairContractOp, nonUnitDim, accTy);
+
+ if (failed(readShuffle))
+ return rewriter.notifyMatchFailure(
+ contractOp, "Accumulator read is not by transfer_read or load");
+ }
// Shuffle the output of contract operations before its use.
LogicalResult writeShuffle = shuffleBeforeWriteLikeOp(
- rewriter, resultWriteOp0, resultWriteOp1, nonUnitDim, accTy);
+ rewriter, contractOp.getResult(), pairContractOp.getResult(),
+ nonUnitDim, accTy);
if (failed(writeShuffle))
return rewriter.notifyMatchFailure(
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
index d3487369355f0..394c950189ead 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
@@ -325,17 +325,17 @@ struct VectorContractToPackedTypeDotProduct
}
// If the accumulators are shuffled we get nullptr else the
- // transfer_read or load operations.
- Operation *accRead =
- traceToVectorReadLikeParentOperation(contractOp.getAcc());
+ // transfer_load or store operations.
+ Operation *accWrite =
+ traceToVectorWriteLikeUserOperation(contractOp.getResult());
if (!pairContractOp &&
- (!isNonUnitDimOperandShuffled(nonUnitDimOperand) || accRead))
+ (!isNonUnitDimOperandShuffled(nonUnitDimOperand) || accWrite))
return rewriter.notifyMatchFailure(contractOp,
"Could not find a contract pair");
// Validate and shuffle the accumulator
- if (accRead) {
+ if (accWrite) {
// Trace back to the load or transfer_read operations of the contract
// accumulators.
Operation *accReadOp0 =
@@ -373,18 +373,22 @@ struct VectorContractToPackedTypeDotProduct
return rewriter.notifyMatchFailure(
contractOp, "The store/write operation of contract operation is "
"before the pair contract operation");
- // Shuffle the accumulators of the contract operations.
- LogicalResult readShuffle =
- shuffleAfterReadLikeOp(rewriter, accReadOp0, accReadOp1, contractOp,
- pairContractOp, nonUnitDimValue, accTy);
- if (failed(readShuffle))
- return rewriter.notifyMatchFailure(
- contractOp, "Accumulator read is not by transfer_read or load");
+ if (!(dyn_cast<arith::ConstantOp>(accReadOp0))) {
+ // Shuffle the accumulators of the contract operations.
+ LogicalResult readShuffle = shuffleAfterReadLikeOp(
+ rewriter, accReadOp0, accReadOp1, contractOp, pairContractOp,
+ nonUnitDimValue, accTy);
+
+ if (failed(readShuffle))
+ return rewriter.notifyMatchFailure(
+ contractOp, "Accumulator read is not by transfer_read or load");
+ }
// Shuffle the output of contract operations before it's use.
LogicalResult writeShuffle = shuffleBeforeWriteLikeOp(
- rewriter, resultWriteOp0, resultWriteOp1, nonUnitDimValue, accTy);
+ rewriter, contractOp.getResult(), pairContractOp.getResult(),
+ nonUnitDimValue, accTy);
if (failed(writeShuffle))
return rewriter.notifyMatchFailure(
diff --git a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
index a04a026f35ae6..232a912a84661 100644
--- a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
+++ b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
@@ -144,6 +144,25 @@ inline ShuffleMasks getShuffleMasks(int64_t nonUnitDimAcc, bool isInt8Avx2) {
return {maskLo8, maskHi8};
}
+// Recursively follows single-use values through scf.yield operations
+// and returns the first non-yield user result in the contraction chain.
+Value contractionUsersAfterYield(Value v) {
+ if (v.getNumUses() != 1)
+ return nullptr;
+
+ OpOperand &use = *v.use_begin();
+ Operation *user = use.getOwner();
+
+ if (!isa<scf::YieldOp>(user))
+ return v;
+
+ auto yield = cast<scf::YieldOp>(user);
+ Operation *parent = yield->getParentOp();
+ unsigned idx = use.getOperandNumber();
+
+ return contractionUsersAfterYield(parent->getResult(idx));
+}
+
// This function walks backward from a value to locate its originating
// vector read-like operation (`vector.transfer_read` or `vector.load`).
// It follows simple forwarding through unary ops and across `scf.for`
@@ -155,9 +174,21 @@ Operation *traceToVectorReadLikeParentOperation(Value v) {
while (true) {
// Case 1: Value defined by an operation
if (Operation *defOp = v.getDefiningOp()) {
- if (isa<vector::TransferReadOp, vector::LoadOp>(defOp))
+ if (isa<vector::TransferReadOp, vector::LoadOp, arith::ConstantOp>(defOp))
return defOp;
+ if (isa<arith::ConstantOp>(defOp)) {
+ Attribute value = (dyn_cast<arith::ConstantOp>(defOp)).getValue();
+
+ if ((dyn_cast<IntegerAttr>(value)).getValue().isZero())
+ return defOp;
+
+ if ((dyn_cast<FloatAttr>(value)).getValue().isZero())
+ return defOp;
+
+ return nullptr;
+ }
+
return nullptr;
}
@@ -289,27 +320,23 @@ LogicalResult shuffleAfterReadLikeOp(PatternRewriter &rewriter, Operation *opA,
// This function shuffles the vectors written by vector.contract operation
// as a flat layout structure before they are stored.
-LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter,
- Operation *opA, Operation *opB,
- int64_t nonUnitDimAcc,
+LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value opA,
+ Value opB, int64_t nonUnitDimAcc,
VectorType accTy) {
- // Helper to extract vector operand from write-like ops
- auto getWrittenVector = [](Operation *op) -> Value {
- if (auto write = dyn_cast<vector::TransferWriteOp>(op))
- return write.getVector();
- if (auto store = dyn_cast<vector::StoreOp>(op))
- return store.getValueToStore();
- return nullptr;
- };
- Value vecA = getWrittenVector(opA);
- Value vecB = getWrittenVector(opB);
+ Value vecA = contractionUsersAfterYield(opA);
+ Value vecB = contractionUsersAfterYield(opB);
if (!vecA || !vecB)
return failure();
+ Operation *resultWriteOpA = *vecA.getUsers().begin();
+ Operation *resultWriteOpB = *vecB.getUsers().begin();
+
// Decide insertion point and location
- Operation *insertBefore = opA->isBeforeInBlock(opB) ? opA : opB;
+ Operation *insertBefore = resultWriteOpA->isBeforeInBlock(resultWriteOpB)
+ ? resultWriteOpA
+ : resultWriteOpB;
rewriter.setInsertionPoint(insertBefore);
Location loc = insertBefore->getLoc();
@@ -335,10 +362,8 @@ LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter,
auto newVecB = vector::ShapeCastOp::create(rewriter, loc, accTy, shuffledHi);
// Update write operands in place via the rewriter to notify it of changes.
- rewriter.modifyOpInPlace(opA,
- [&]() { opA->setOperand(0, newVecA.getResult()); });
- rewriter.modifyOpInPlace(opB,
- [&]() { opB->setOperand(0, newVecB.getResult()); });
+ resultWriteOpA->replaceUsesOfWith(vecA, newVecA);
+ resultWriteOpB->replaceUsesOfWith(vecB, newVecB);
return success();
}
diff --git a/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir b/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
new file mode 100644
index 0000000000000..8ddec33b086df
--- /dev/null
+++ b/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
@@ -0,0 +1,136 @@
+// RUN: mlir-opt %s -transform-interpreter -cse -split-input-file | FileCheck %s
+
+!vecA = vector<1x1x2xbf16>
+!vecB = vector<1x2x16xbf16>
+!vecC = vector<1x16xf32>
+!memrefA = memref<1x1x2xbf16, strided<[2048, 32, 1], offset: ?>>
+!memrefB = memref<1x2x32xbf16, strided<[2048, 64, 1], offset: ?>>
+!memrefC = memref<1x32xf32, strided<[64, 1], offset: ?>>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @brmatmul_acc_mv(%arg0: memref<16x64x32xbf16>, %arg1: memref<16x32x64xbf16>,
+ %arg2: memref<64x64xf32>) -> memref<64x64xf32> {
+ %0 = ub.poison : f32
+ %1 = ub.poison : bf16
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ scf.for %arg3 = %c0 to %c64 step %c1 {
+ scf.for %arg4 = %c0 to %c64 step %c32 {
+ %subview = memref.subview %arg2[%arg3, %arg4] [1, 32] [1, 1]
+ : memref<64x64xf32> 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 %c16 step %c1 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
+ %5:2 = scf.for %arg8 = %c0 to %c32 step %c2 iter_args(%arg9 = %arg6, %arg10 = %arg7) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg5, %arg3, %arg8] [1, 1, 2] [1, 1, 1]
+ : memref<16x64x32xbf16> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg8, %arg4] [1, 2, 32] [1, 1, 1]
+ : memref<16x32x64xbf16> to !memrefB
+
+ %6 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]}
+ : !memrefA, !vecA
+ %7 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]}
+ : !memrefB, !vecB
+ %8 = vector.transfer_read %subview_1[%c0, %c0, %c16], %1 {in_bounds = [true, true, true]}
+ : !memrefB, !vecB
+
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %6, %7, %arg9
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+ %10 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %6, %8, %arg10
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+
+ scf.yield %9, %10 : !vecC, !vecC
+ }
+ scf.yield %5#0, %5#1 : !vecC, !vecC
+ }
+
+ vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]}
+ : !vecC, !memrefC
+ vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]}
+ : !vecC, !memrefC
+ }
+ }
+
+ return %arg2 : memref<64x64xf32>
+}
+
+// CHECK-LABEL: @brmatmul_acc_mv
+// CHECK: arith.constant dense<0.000000e+00> : vector<1x16xf32>
+// CHECK: arith.addf
+// CHECK-NEXT: vector.transfer_write
+// CHECK: arith.addf
+// CHECK-NEXT: vector.transfer_write
+
+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.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<1x16x16x2xbf16>
+!vecB = vector<1x16x16x2xbf16>
+!vecC = vector<16x16xf32>
+!memrefA = memref<1x32x16x2xbf16>
+!memrefB = memref<1x16x32x2xbf16>
+!memrefC = memref<32x32xf32>
+#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
+#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
+#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d1, d2)>
+func.func @negative_no_loop(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+
+ %1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
+ !memrefB, !vecB
+
+ %3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
+
+ %4 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["reduction", "reduction", "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
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_no_loop
+// CHECK-NOT: arith.constant dense<0.000000e+00> {{.*}}
+// CHECK-NOT: arith.addf
+// CHECK-NOT: arith.addf
+
+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.any_op
+ transform.yield
+ }
+}
+
>From 63ea2855e6bb37d179f13598c973cd5e83375d71 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Thu, 23 Jul 2026 00:13:56 -0700
Subject: [PATCH 2/7] fix typos, expand test-case
---
.../include/mlir/Dialect/X86/Utils/X86Utils.h | 4 ++--
.../MoveAccumulatorForContractLoop.cpp | 11 ++++++-----
mlir/lib/Dialect/X86/Utils/X86Utils.cpp | 10 +++++-----
.../X86/move-acc-for-contract-loop.mlir | 19 ++++++++++++++-----
4 files changed, 27 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h b/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
index f067ce3bb9b65..edc17e08cde7c 100644
--- a/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
+++ b/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
@@ -63,8 +63,8 @@ LogicalResult shuffleAfterReadLikeOp(PatternRewriter &rewriter, Operation *opA,
// Shuffles vectors produced by vector.contraction ops into a flat layout
// before they are written to memory.
-LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value opA,
- Value opB, int64_t nonUnitDimAcc,
+LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value valA,
+ Value valB, int64_t nonUnitDimAcc,
VectorType accTy);
} // namespace x86
diff --git a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
index 6f87ebfa441dc..3c15c680e65e6 100644
--- a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
@@ -58,8 +58,7 @@ struct MoveAccumulatorForContractLoop
if (dyn_cast<arith::ConstantOp>(accReadOp))
return rewriter.notifyMatchFailure(
contractOp,
- "The input acc to contract is already a constant zero vector.");
- ;
+ "The input acc to contract is already a constant vector.");
if ((accReadOp->getBlock() == contractOp->getBlock()) ||
(resultWriteOp->getBlock() == contractOp->getBlock()))
@@ -67,11 +66,13 @@ struct MoveAccumulatorForContractLoop
contractOp, "Acc read/write should be in a separate block.");
// Replace acc of a contraction operation with vector constant.
- rewriter.setInsertionPointAfter(accReadOp);
Value accValue = accReadOp->getResult(0);
+ Operation *firstUser = *accValue.getUsers().begin();
+ rewriter.setInsertionPoint(firstUser);
+
auto vecTy = llvm::dyn_cast<VectorType>(accValue.getType());
if (!vecTy)
- return rewriter.notifyMatchFailure(contractOp, "Excepts vector type.");
+ return rewriter.notifyMatchFailure(contractOp, "Expects vector type.");
Location loc = accReadOp->getLoc();
Type elemTy = vecTy.getElementType();
@@ -84,7 +85,7 @@ struct MoveAccumulatorForContractLoop
accValue.replaceAllUsesWith(zeroVec);
- // Adds the initial acc value with acontract results before storing to acc
+ // Adds the initial acc value with contract results before storing to acc
// matrix.
rewriter.setInsertionPoint(resultWriteOp);
Location locUser = resultWriteOp->getLoc();
diff --git a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
index 232a912a84661..bd3f0d5244f2e 100644
--- a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
+++ b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
@@ -147,7 +147,7 @@ inline ShuffleMasks getShuffleMasks(int64_t nonUnitDimAcc, bool isInt8Avx2) {
// Recursively follows single-use values through scf.yield operations
// and returns the first non-yield user result in the contraction chain.
Value contractionUsersAfterYield(Value v) {
- if (v.getNumUses() != 1)
+ if (!v || v.getNumUses() != 1)
return nullptr;
OpOperand &use = *v.use_begin();
@@ -320,12 +320,12 @@ LogicalResult shuffleAfterReadLikeOp(PatternRewriter &rewriter, Operation *opA,
// This function shuffles the vectors written by vector.contract operation
// as a flat layout structure before they are stored.
-LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value opA,
- Value opB, int64_t nonUnitDimAcc,
+LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value valA,
+ Value valB, int64_t nonUnitDimAcc,
VectorType accTy) {
- Value vecA = contractionUsersAfterYield(opA);
- Value vecB = contractionUsersAfterYield(opB);
+ Value vecA = contractionUsersAfterYield(valA);
+ Value vecB = contractionUsersAfterYield(valB);
if (!vecA || !vecB)
return failure();
diff --git a/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir b/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
index 8ddec33b086df..547b6a591c1d6 100644
--- a/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
+++ b/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
@@ -67,11 +67,20 @@ func.func @brmatmul_acc_mv(%arg0: memref<16x64x32xbf16>, %arg1: memref<16x32x64x
}
// CHECK-LABEL: @brmatmul_acc_mv
-// CHECK: arith.constant dense<0.000000e+00> : vector<1x16xf32>
-// CHECK: arith.addf
-// CHECK-NEXT: vector.transfer_write
-// CHECK: arith.addf
-// CHECK-NEXT: vector.transfer_write
+// CHECK: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<1x16xf32>
+// CHECK: scf.for
+// CHECK-NEXT: scf.for
+// CHECK: memref.subview
+// CHECK: %[[TRANSFER_READ_0:.*]] = vector.transfer_read
+// CHECK: %[[TRANSFER_READ_1:.*]] = vector.transfer_read
+// CHECK: scf.for {{.*}} iter_args(%[[VAL_3:.*]] = %[[CST]], %[[VAL_4:.*]] = %[[CST]]) -> (vector<1x16xf32>, vector<1x16xf32>) {
+// CHECK-NEXT: scf.for
+// CHECK: scf.yield
+// CHECK: scf.yield
+// CHECK: %[[ADDF_0:.*]] = arith.addf {{.*}}, %[[TRANSFER_READ_1]] : vector<1x16xf32>
+// CHECK-NEXT: vector.transfer_write %[[ADDF_0:.*]], %subview[%c0, %c16] {{.*}}
+// CHECK: %[[ADDF_1:.*]] = arith.addf {{.*}}, %[[TRANSFER_READ_0]] : vector<1x16xf32>
+// CHECK-NEXT: vector.transfer_write %[[ADDF_1:.*]], %subview[%c0, %c0] {{.*}}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
>From 3f60705d506d5134f124d549972c4de1d38c6f0d Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Wed, 29 Jul 2026 06:28:58 -0700
Subject: [PATCH 3/7] clean ups, new test-case, and code refactor
---
.../include/mlir/Dialect/X86/Utils/X86Utils.h | 6 +-
.../MoveAccumulatorForContractLoop.cpp | 19 ++--
.../Transforms/VectorContractBF16ToFMA.cpp | 2 +-
.../VectorContractToPackedTypeDotProduct.cpp | 2 +-
mlir/lib/Dialect/X86/Utils/X86Utils.cpp | 29 +++---
.../X86/move-acc-for-contract-loop.mlir | 89 +++++++++++++++++++
6 files changed, 119 insertions(+), 28 deletions(-)
diff --git a/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h b/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
index edc17e08cde7c..3e8b33f712b35 100644
--- a/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
+++ b/mlir/include/mlir/Dialect/X86/Utils/X86Utils.h
@@ -63,9 +63,9 @@ LogicalResult shuffleAfterReadLikeOp(PatternRewriter &rewriter, Operation *opA,
// Shuffles vectors produced by vector.contraction ops into a flat layout
// before they are written to memory.
-LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value valA,
- Value valB, int64_t nonUnitDimAcc,
- VectorType accTy);
+LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter,
+ Value contractARes, Value contractBRes,
+ int64_t nonUnitDimAcc, VectorType accTy);
} // namespace x86
} // namespace mlir
diff --git a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
index 3c15c680e65e6..72bded7732430 100644
--- a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
@@ -52,8 +52,9 @@ struct MoveAccumulatorForContractLoop
if (!accReadOp || !resultWriteOp)
return rewriter.notifyMatchFailure(
- contractOp, "Read/write from/to acc matrix is not by "
- "transfer_read/load/transfer_write/store ops.");
+ contractOp, "Read from acc matrix is not by "
+ "transfer_read/load/constant_zero or multiple users of "
+ "contract operation.");
if (dyn_cast<arith::ConstantOp>(accReadOp))
return rewriter.notifyMatchFailure(
@@ -77,11 +78,9 @@ struct MoveAccumulatorForContractLoop
Location loc = accReadOp->getLoc();
Type elemTy = vecTy.getElementType();
- Value zeroScalar = arith::ConstantOp::create(rewriter, loc, elemTy,
- rewriter.getZeroAttr(elemTy));
-
- Value zeroVec =
- vector::BroadcastOp::create(rewriter, loc, vecTy, zeroScalar);
+ Value zeroVec = arith::ConstantOp::create(
+ rewriter, loc,
+ DenseElementsAttr::get(vecTy, rewriter.getZeroAttr(elemTy)));
accValue.replaceAllUsesWith(zeroVec);
@@ -95,11 +94,11 @@ struct MoveAccumulatorForContractLoop
if (llvm::isa<FloatType>(elemTy)) {
addition =
arith::AddFOp::create(rewriter, locUser, contractValue, accValue);
- }
-
- if (llvm::isa<IntegerType>(elemTy)) {
+ } else if (llvm::isa<IntegerType>(elemTy)) {
addition =
arith::AddIOp::create(rewriter, locUser, contractValue, accValue);
+ } else {
+ llvm_unreachable("expected floating-point or integer element type");
}
resultWriteOp->replaceUsesOfWith(contractValue, addition);
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp
index c70ad221c772d..289a5c1ae8ef2 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp
@@ -415,7 +415,7 @@ struct VectorContractBF16ToFMA
Operation *accReadOp1 =
traceToVectorReadLikeParentOperation(pairContractOp.getAcc());
- if (!(dyn_cast<arith::ConstantOp>(accReadOp0))) {
+ if (!isa<arith::ConstantOp>(accReadOp0)) {
// Shuffle the accumulators of the contract operations.
LogicalResult readShuffle =
shuffleAfterReadLikeOp(rewriter, accReadOp0, accReadOp1, contractOp,
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
index 394c950189ead..bd55a5bb6ab8f 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
@@ -374,7 +374,7 @@ struct VectorContractToPackedTypeDotProduct
contractOp, "The store/write operation of contract operation is "
"before the pair contract operation");
- if (!(dyn_cast<arith::ConstantOp>(accReadOp0))) {
+ if (!isa<arith::ConstantOp>(accReadOp0)) {
// Shuffle the accumulators of the contract operations.
LogicalResult readShuffle = shuffleAfterReadLikeOp(
rewriter, accReadOp0, accReadOp1, contractOp, pairContractOp,
diff --git a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
index bd3f0d5244f2e..36b84c685edc5 100644
--- a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
+++ b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
@@ -174,19 +174,21 @@ Operation *traceToVectorReadLikeParentOperation(Value v) {
while (true) {
// Case 1: Value defined by an operation
if (Operation *defOp = v.getDefiningOp()) {
- if (isa<vector::TransferReadOp, vector::LoadOp, arith::ConstantOp>(defOp))
- return defOp;
+ if (isa<vector::TransferReadOp, vector::LoadOp, arith::ConstantOp>(
+ defOp)) {
- if (isa<arith::ConstantOp>(defOp)) {
- Attribute value = (dyn_cast<arith::ConstantOp>(defOp)).getValue();
+ if (auto constOp = dyn_cast<arith::ConstantOp>(defOp)) {
+ Attribute value = constOp.getValue();
- if ((dyn_cast<IntegerAttr>(value)).getValue().isZero())
- return defOp;
+ if (auto intAttr = dyn_cast<IntegerAttr>(value))
+ return intAttr.getValue().isZero() ? defOp : nullptr;
- if ((dyn_cast<FloatAttr>(value)).getValue().isZero())
- return defOp;
+ if (auto floatAttr = dyn_cast<FloatAttr>(value))
+ return floatAttr.getValue().isZero() ? defOp : nullptr;
- return nullptr;
+ return nullptr;
+ }
+ return defOp;
}
return nullptr;
@@ -320,12 +322,13 @@ LogicalResult shuffleAfterReadLikeOp(PatternRewriter &rewriter, Operation *opA,
// This function shuffles the vectors written by vector.contract operation
// as a flat layout structure before they are stored.
-LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Value valA,
- Value valB, int64_t nonUnitDimAcc,
+LogicalResult shuffleBeforeWriteLikeOp(PatternRewriter &rewriter,
+ Value contractARes, Value contractBRes,
+ int64_t nonUnitDimAcc,
VectorType accTy) {
- Value vecA = contractionUsersAfterYield(valA);
- Value vecB = contractionUsersAfterYield(valB);
+ Value vecA = contractionUsersAfterYield(contractARes);
+ Value vecB = contractionUsersAfterYield(contractBRes);
if (!vecA || !vecB)
return failure();
diff --git a/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir b/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
index 547b6a591c1d6..b732f97d0fdb5 100644
--- a/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
+++ b/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
@@ -143,3 +143,92 @@ module attributes {transform.with_named_sequence} {
}
}
+// -----
+
+!vecA = vector<1x1x2xbf16>
+!vecB = vector<1x2x16xbf16>
+!vecC = vector<1x16xf32>
+!memrefA = memref<1x1x2xbf16, strided<[2048, 32, 1], offset: ?>>
+!memrefB = memref<1x2x32xbf16, strided<[2048, 64, 1], offset: ?>>
+!memrefC = memref<1x32xf32, strided<[64, 1], offset: ?>>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @negative_non_zero_cst(%arg0: memref<16x64x32xbf16>, %arg1: memref<16x32x64xbf16>,
+ %arg2: memref<64x64xf32>) -> memref<64x64xf32> {
+ %cst2 = arith.constant dense<2.000000e+00> : vector<1x16xf32>
+ %cst3 = arith.constant dense<3.000000e+00> : vector<1x16xf32>
+ %0 = ub.poison : f32
+ %1 = ub.poison : bf16
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ scf.for %arg3 = %c0 to %c64 step %c1 {
+ scf.for %arg4 = %c0 to %c64 step %c32 {
+ %subview = memref.subview %arg2[%arg3, %arg4] [1, 32] [1, 1]
+ : memref<64x64xf32> to !memrefC
+
+ %4:2 = scf.for %arg5 = %c0 to %c16 step %c1 iter_args(%arg6 = %cst2, %arg7 = %cst3) -> (!vecC, !vecC) {
+ %5:2 = scf.for %arg8 = %c0 to %c32 step %c2 iter_args(%arg9 = %arg6, %arg10 = %arg7) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg5, %arg3, %arg8] [1, 1, 2] [1, 1, 1]
+ : memref<16x64x32xbf16> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg8, %arg4] [1, 2, 32] [1, 1, 1]
+ : memref<16x32x64xbf16> to !memrefB
+
+ %6 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]}
+ : !memrefA, !vecA
+ %7 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]}
+ : !memrefB, !vecB
+ %8 = vector.transfer_read %subview_1[%c0, %c0, %c16], %1 {in_bounds = [true, true, true]}
+ : !memrefB, !vecB
+
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %6, %7, %arg9
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+ %10 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %6, %8, %arg10
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+
+ scf.yield %9, %10 : !vecC, !vecC
+ }
+ scf.yield %5#0, %5#1 : !vecC, !vecC
+ }
+
+ vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]}
+ : !vecC, !memrefC
+ vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]}
+ : !vecC, !memrefC
+ }
+ }
+
+ return %arg2 : memref<64x64xf32>
+}
+
+// CHECK-LABEL: @negative_non_zero_cst
+// CHECK-NOT: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<1x16xf32>
+// CHECK: scf.for
+// CHECK-NEXT: scf.for
+// CHECK: memref.subview
+// CHECK: scf.for
+// CHECK-NEXT: scf.for
+// CHECK: scf.yield
+// CHECK: scf.yield
+// CHECK-NOT: arith.addf
+
+
+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.any_op
+ transform.yield
+ }
+}
+
+
>From 398275548b52bfcbca433f2bffb5dd94749db3d6 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Wed, 29 Jul 2026 11:11:49 -0700
Subject: [PATCH 4/7] test-case modification
---
.../VectorContractToPackedTypeDotProduct.cpp | 2 +-
mlir/lib/Dialect/X86/Utils/X86Utils.cpp | 16 +++++++++++-----
.../Dialect/X86/vector-contract-bf16-to-fma.mlir | 3 +--
...ector-contract-to-packed-type-dotproduct.mlir | 3 +--
4 files changed, 14 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
index bd55a5bb6ab8f..a4f0875d0c6fe 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
@@ -353,7 +353,7 @@ struct VectorContractToPackedTypeDotProduct
if (!accReadOp0 || !accReadOp1)
return rewriter.notifyMatchFailure(
contractOp,
- "Operands doesn't have load or transfer_read as it's parent op");
+ "Operands doesn't have load or transfer_read or dense constant attribute as it's parent op");
if (!resultWriteOp0 || !resultWriteOp1)
return rewriter.notifyMatchFailure(
diff --git a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
index 36b84c685edc5..74139cd219b18 100644
--- a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
+++ b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
@@ -178,13 +178,19 @@ Operation *traceToVectorReadLikeParentOperation(Value v) {
defOp)) {
if (auto constOp = dyn_cast<arith::ConstantOp>(defOp)) {
- Attribute value = constOp.getValue();
+ if (auto denseAttr =
+ dyn_cast<DenseElementsAttr>(constOp.getValue())) {
+ if (!denseAttr.isSplat())
+ return nullptr;
- if (auto intAttr = dyn_cast<IntegerAttr>(value))
- return intAttr.getValue().isZero() ? defOp : nullptr;
+ Attribute splat = denseAttr.getSplatValue<Attribute>();
- if (auto floatAttr = dyn_cast<FloatAttr>(value))
- return floatAttr.getValue().isZero() ? defOp : nullptr;
+ if (auto floatAttr = dyn_cast<FloatAttr>(splat))
+ return floatAttr.getValue().isZero() ? defOp : nullptr;
+
+ if (auto intAttr = dyn_cast<IntegerAttr>(splat))
+ return intAttr.getValue().isZero() ? defOp : nullptr;
+ }
return nullptr;
}
diff --git a/mlir/test/Dialect/X86/vector-contract-bf16-to-fma.mlir b/mlir/test/Dialect/X86/vector-contract-bf16-to-fma.mlir
index 4f0e5c5f3c907..94abf95051854 100644
--- a/mlir/test/Dialect/X86/vector-contract-bf16-to-fma.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-bf16-to-fma.mlir
@@ -444,8 +444,6 @@ func.func @matmul_to_fma_flat_layout_loop(%arg0: memref<16x64x32xbf16>, %arg1: m
}
// CHECK-LABEL: @matmul_to_fma_flat_layout_loop
-// CHECK: vector.shuffle{{.*}}[0, 8, 1, 9, 2, 10, 3, 11] : vector<8xf32>, vector<8xf32>
-// CHECK-NEXT: vector.shuffle{{.*}}[4, 12, 5, 13, 6, 14, 7, 15] : vector<8xf32>, vector<8xf32>
// CHECK: scf.for
// CHECK: scf.for
// CHECK: x86.avx.bcst_to_f32.packed
@@ -461,6 +459,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_bf16_to_fma
} : !transform.any_op
transform.yield
diff --git a/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir b/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
index 5d32570b21cbe..181a789a9adc0 100644
--- a/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
@@ -687,8 +687,6 @@ func.func @brmatmul_bf16dp_flat_layout_loop(%arg0: memref<16x64x32xbf16>, %arg1:
}
// CHECK-LABEL: @brmatmul_bf16dp_flat_layout_loop
-// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xf32>, vector<16xf32>
-// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
// CHECK: scf.for
// CHECK: scf.for
// CHECK: vector.transfer_read {{.*}} vector<32xbf16>
@@ -706,6 +704,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_packed_type_dot_product
} : !transform.any_op
transform.yield
>From eff9171389caf1f83b73a5ebf6b13dc9beb30c18 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Wed, 29 Jul 2026 19:13:37 -0700
Subject: [PATCH 5/7] clang-fix
---
.../X86/Transforms/VectorContractToPackedTypeDotProduct.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
index a4f0875d0c6fe..fd2320b3af63d 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
@@ -352,8 +352,8 @@ struct VectorContractToPackedTypeDotProduct
if (!accReadOp0 || !accReadOp1)
return rewriter.notifyMatchFailure(
- contractOp,
- "Operands doesn't have load or transfer_read or dense constant attribute as it's parent op");
+ contractOp, "Operands doesn't have load or transfer_read or "
+ "dense constant attribute as it's parent op");
if (!resultWriteOp0 || !resultWriteOp1)
return rewriter.notifyMatchFailure(
>From 8447b57e5b5a46a25fa78578371d4e0a012b75a3 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Thu, 30 Jul 2026 03:29:06 -0700
Subject: [PATCH 6/7] variable name update
---
.../Transforms/MoveAccumulatorForContractLoop.cpp | 12 ++++++------
1 file changed, 6 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
index 72bded7732430..40b8412aa5efd 100644
--- a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
@@ -48,9 +48,9 @@ struct MoveAccumulatorForContractLoop
return rewriter.notifyMatchFailure(
contractOp, "Final acc write might have multiple users.");
- Operation *resultWriteOp = *contractValue.getUsers().begin();
+ Operation *resultUserOp = *contractValue.getUsers().begin();
- if (!accReadOp || !resultWriteOp)
+ if (!accReadOp || !resultUserOp)
return rewriter.notifyMatchFailure(
contractOp, "Read from acc matrix is not by "
"transfer_read/load/constant_zero or multiple users of "
@@ -62,7 +62,7 @@ struct MoveAccumulatorForContractLoop
"The input acc to contract is already a constant vector.");
if ((accReadOp->getBlock() == contractOp->getBlock()) ||
- (resultWriteOp->getBlock() == contractOp->getBlock()))
+ (resultUserOp->getBlock() == contractOp->getBlock()))
return rewriter.notifyMatchFailure(
contractOp, "Acc read/write should be in a separate block.");
@@ -86,8 +86,8 @@ struct MoveAccumulatorForContractLoop
// Adds the initial acc value with contract results before storing to acc
// matrix.
- rewriter.setInsertionPoint(resultWriteOp);
- Location locUser = resultWriteOp->getLoc();
+ rewriter.setInsertionPoint(resultUserOp);
+ Location locUser = resultUserOp->getLoc();
Value addition;
@@ -101,7 +101,7 @@ struct MoveAccumulatorForContractLoop
llvm_unreachable("expected floating-point or integer element type");
}
- resultWriteOp->replaceUsesOfWith(contractValue, addition);
+ resultUserOp->replaceUsesOfWith(contractValue, addition);
return success();
}
};
>From 0b2fef8451e16bb718498f62a4f64d38c4a73e80 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Fri, 31 Jul 2026 07:17:37 -0700
Subject: [PATCH 7/7] added extra check for input acc has many users +
test-case
---
.../MoveAccumulatorForContractLoop.cpp | 4 +
.../X86/move-acc-for-contract-loop.mlir | 89 +++++++++++++++++++
2 files changed, 93 insertions(+)
diff --git a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
index 40b8412aa5efd..325509727d799 100644
--- a/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp
@@ -68,6 +68,10 @@ struct MoveAccumulatorForContractLoop
// Replace acc of a contraction operation with vector constant.
Value accValue = accReadOp->getResult(0);
+ if (!accValue.hasOneUse())
+ return rewriter.notifyMatchFailure(
+ contractOp, "The input accumulator has multiple users.");
+
Operation *firstUser = *accValue.getUsers().begin();
rewriter.setInsertionPoint(firstUser);
diff --git a/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir b/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
index b732f97d0fdb5..8007f5cd8fc8d 100644
--- a/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
+++ b/mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
@@ -231,4 +231,93 @@ module attributes {transform.with_named_sequence} {
}
}
+// -----
+
+!vecA = vector<1x1x2xbf16>
+!vecB = vector<1x2x16xbf16>
+!vecC = vector<1x16xf32>
+!memrefA = memref<1x1x2xbf16, strided<[2048, 32, 1], offset: ?>>
+!memrefB = memref<1x2x32xbf16, strided<[2048, 64, 1], offset: ?>>
+!memrefC = memref<1x32xf32, strided<[64, 1], offset: ?>>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @negative_input_acc_many_users(%arg0: memref<16x64x32xbf16>, %arg1: memref<16x32x64xbf16>,
+ %arg2: memref<64x64xf32>) -> memref<64x64xf32> {
+ %0 = ub.poison : f32
+ %1 = ub.poison : bf16
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ scf.for %arg3 = %c0 to %c64 step %c1 {
+ scf.for %arg4 = %c0 to %c64 step %c32 {
+ %subview = memref.subview %arg2[%arg3, %arg4] [1, 32] [1, 1]
+ : memref<64x64xf32> 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 %c16 step %c1 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
+ %5:2 = scf.for %arg8 = %c0 to %c32 step %c2 iter_args(%arg9 = %arg6, %arg10 = %arg7) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg5, %arg3, %arg8] [1, 1, 2] [1, 1, 1]
+ : memref<16x64x32xbf16> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg8, %arg4] [1, 2, 32] [1, 1, 1]
+ : memref<16x32x64xbf16> to !memrefB
+
+ %6 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]}
+ : !memrefA, !vecA
+ %7 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]}
+ : !memrefB, !vecB
+ %8 = vector.transfer_read %subview_1[%c0, %c0, %c16], %1 {in_bounds = [true, true, true]}
+ : !memrefB, !vecB
+
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %6, %7, %arg9
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+ %10 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %6, %8, %arg10
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+
+ scf.yield %9, %10 : !vecC, !vecC
+ }
+ scf.yield %5#0, %5#1 : !vecC, !vecC
+ }
+
+ %7 = arith.addf %4#1, %2 : vector<1x16xf32>
+ vector.transfer_write %7, %subview[%c0, %c16] {in_bounds = [true, true]}
+ : !vecC, !memrefC
+ %8 = arith.addf %4#0, %3 : vector<1x16xf32>
+ vector.transfer_write %8, %subview[%c0, %c0] {in_bounds = [true, true]}
+ : !vecC, !memrefC
+ }
+ }
+
+ return %arg2 : memref<64x64xf32>
+}
+// CHECK-LABEL: @negative_input_acc_many_users
+// CHECK-NOT: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<1x16xf32>
+// CHECK: scf.for
+// CHECK-NEXT: scf.for
+// CHECK: memref.subview
+// CHECK: scf.for
+// CHECK-NEXT: scf.for
+// CHECK: scf.yield
+// CHECK: scf.yield
+
+
+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.any_op
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list