[mlir][x86] Decouple Accumulator from ADD-Based Vector Contractions (over loops) (#204327)

This path transforms `vector.contract(A, B, Acc)` into
`vector.contract(A, B, 0) + Acc` to decouple the contraction computation
from the accumulator update for contractions over the loop.

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

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

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.
@@ -60,7 +64,7 @@ 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,
+                                       Value contractARes, Value contractBRes,
                                        int64_t nonUnitDimAcc, VectorType accTy);
 
 } // namespace x86

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
 //===----------------------------------------------------------------------===//

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

mlir/lib/Dialect/X86/Transforms/MoveAccumulatorForContractLoop.cpp

@@ -0,0 +1,118 @@
+//===- 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 *resultUserOp = *contractValue.getUsers().begin();
+
+    if (!accReadOp || !resultUserOp)
+      return rewriter.notifyMatchFailure(
+          contractOp, "Read from acc matrix is not by "
+                      "transfer_read/load/constant_zero or multiple users of "
+                      "contract operation.");
+
+    if (isa<arith::ConstantOp>(accReadOp))
+      return rewriter.notifyMatchFailure(
+          contractOp,
+          "The input acc to contract is already a constant vector.");
+
+    if ((accReadOp->getBlock() == contractOp->getBlock()) ||
+        (resultUserOp->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.
+    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);
+
+    auto vecTy = dyn_cast<VectorType>(accValue.getType());
+    if (!vecTy)
+      return rewriter.notifyMatchFailure(contractOp, "Expects vector type.");
+
+    Location loc = accReadOp->getLoc();
+    Type elemTy = vecTy.getElementType();
+
+    Value zeroVec = arith::ConstantOp::create(
+        rewriter, loc,
+        DenseElementsAttr::get(vecTy, rewriter.getZeroAttr(elemTy)));
+
+    accValue.replaceAllUsesWith(zeroVec);
+
+    // Adds the initial acc value with contract results before storing to acc
+    // matrix.
+    rewriter.setInsertionPoint(resultUserOp);
+    Location locUser = resultUserOp->getLoc();
+
+    Value addition;
+
+    if (isa<FloatType>(elemTy)) {
+      addition =
+          arith::AddFOp::create(rewriter, locUser, contractValue, accValue);
+    } else if (isa<IntegerType>(elemTy)) {
+      addition =
+          arith::AddIOp::create(rewriter, locUser, contractValue, accValue);
+    } else {
+      llvm_unreachable("expected floating-point or integer element type");
+    }
+
+    resultUserOp->replaceUsesOfWith(contractValue, addition);
+    return success();
+  }
+};
+
+} // namespace
+
+void x86::populateMoveAccumulatorForContractLoopPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<MoveAccumulatorForContractLoop>(patterns.getContext());
+}

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
+      if (!isa<arith::ConstantOp>(accReadOp0)) {
-      // or transfer_write.
+        // Shuffle the accumulators of the contract operations.
-      Operation *resultWriteOp0 =
+        LogicalResult readShuffle =
-          traceToVectorWriteLikeUserOperation(contractOp.getResult());
+            shuffleAfterReadLikeOp(rewriter, accReadOp0, accReadOp1, contractOp,
-      Operation *resultWriteOp1 =
+                                   pairContractOp, nonUnitDim, accTy);
-          traceToVectorWriteLikeUserOperation(pairContractOp.getResult());
 
-      // Shuffle the accumulators of the contract operations.
+        if (failed(readShuffle))
-      LogicalResult readShuffle =
+          return rewriter.notifyMatchFailure(
-          shuffleAfterReadLikeOp(rewriter, accReadOp0, accReadOp1, contractOp,
+              contractOp, "Accumulator read is not by transfer_read or load");
-                                 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(

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.
+      // transfer_load or store operations.
-      Operation *accRead =
+      Operation *accWrite =
-          traceToVectorReadLikeParentOperation(contractOp.getAcc());
+          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 =
@@ -352,8 +352,8 @@ struct VectorContractToPackedTypeDotProduct
 
         if (!accReadOp0 || !accReadOp1)
           return rewriter.notifyMatchFailure(
-              contractOp,
+              contractOp, "Operands doesn't have load or transfer_read or "
-              "Operands doesn't have load or transfer_read as it's parent op");
+                          "dense constant attribute as its parent op");
 
         if (!resultWriteOp0 || !resultWriteOp1)
           return rewriter.notifyMatchFailure(
@@ -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))
+        if (!isa<arith::ConstantOp>(accReadOp0)) {
-          return rewriter.notifyMatchFailure(
+          // Shuffle the accumulators of the contract operations.
-              contractOp, "Accumulator read is not by transfer_read or load");
+          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(

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 || 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,8 +174,28 @@ 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)) {
+
+        if (auto constOp = dyn_cast<arith::ConstantOp>(defOp)) {
+          if (auto denseAttr =
+                  dyn_cast<DenseElementsAttr>(constOp.getValue())) {
+            if (!denseAttr.isSplat())
+              return nullptr;
+
+            Attribute splat = denseAttr.getSplatValue<Attribute>();
+
+            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;
+        }
         return defOp;
+      }
 
       return nullptr;
     }
@@ -290,26 +329,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,
+                                       Value contractARes, Value contractBRes,
                                        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 vecA = contractionUsersAfterYield(contractARes);
-  Value vecB = getWrittenVector(opB);
+  Value vecB = contractionUsersAfterYield(contractBRes);
 
   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 +371,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,
+  resultWriteOpA->replaceUsesOfWith(vecA, newVecA);
-                           [&]() { opA->setOperand(0, newVecA.getResult()); });
+  resultWriteOpB->replaceUsesOfWith(vecB, newVecB);
-  rewriter.modifyOpInPlace(opB,
-                           [&]() { opB->setOperand(0, newVecB.getResult()); });
 
   return success();
 }

mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir

@@ -0,0 +1,323 @@
+// 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: %[[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}) {
+    %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
+  }
+}
+
+// -----
+
+!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
+  }
+}
+
+// -----
+
+!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
+  }
+}
+

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

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