From b32d4f0c451f2c6c19e1508441e78033c6f95e9b Mon Sep 17 00:00:00 2001 From: Arun Thangamani Date: Sun, 2 Aug 2026 07:33:00 +0530 Subject: [PATCH] [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. --- .../X86/TransformOps/X86TransformOps.td | 10 + mlir/include/mlir/Dialect/X86/Transforms.h | 5 + .../include/mlir/Dialect/X86/Utils/X86Utils.h | 6 +- .../X86/TransformOps/X86TransformOps.cpp | 5 + .../lib/Dialect/X86/Transforms/CMakeLists.txt | 1 + .../MoveAccumulatorForContractLoop.cpp | 118 +++++++ .../Transforms/VectorContractBF16ToFMA.cpp | 26 +- .../VectorContractToPackedTypeDotProduct.cpp | 34 +- mlir/lib/Dialect/X86/Utils/X86Utils.cpp | 68 +++- .../X86/move-acc-for-contract-loop.mlir | 323 ++++++++++++++++++ .../X86/vector-contract-bf16-to-fma.mlir | 3 +- ...or-contract-to-packed-type-dotproduct.mlir | 3 +- 12 files changed, 550 insertions(+), 52 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 c474cfb47d00..48aabf90914e 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]> { + 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 6ebba5e94ec7..8bcaf9a67feb 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 2ff8f4f4283a..3e8b33f712b3 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 indexingMaps, std::optional 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 diff --git a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp index 390b21e12b0e..a52f08a5cd06 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 9c3695536cda..5118563c03f1 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 000000000000..b580ef15923f --- /dev/null +++ b/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 { + using OpRewritePattern::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(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(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(elemTy)) { + addition = + arith::AddFOp::create(rewriter, locUser, contractValue, accValue); + } else if (isa(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(patterns.getContext()); +} diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractBF16ToFMA.cpp index 287892c3a660..289a5c1ae8ef 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()); + if (!isa(accReadOp0)) { + // Shuffle the accumulators of the contract operations. + LogicalResult readShuffle = + shuffleAfterReadLikeOp(rewriter, accReadOp0, accReadOp1, contractOp, + pairContractOp, nonUnitDim, accTy); - // 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 (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 d3487369355f..9ec9bb04570c 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 = @@ -352,8 +352,8 @@ struct VectorContractToPackedTypeDotProduct if (!accReadOp0 || !accReadOp1) return rewriter.notifyMatchFailure( - contractOp, - "Operands doesn't have load or transfer_read as it's parent op"); + contractOp, "Operands doesn't have load or transfer_read or " + "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)) - return rewriter.notifyMatchFailure( - contractOp, "Accumulator read is not by transfer_read or load"); + if (!isa(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 a04a026f35ae..74139cd219b1 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 || v.getNumUses() != 1) + return nullptr; + + OpOperand &use = *v.use_begin(); + Operation *user = use.getOwner(); + + if (!isa(user)) + return v; + + auto yield = cast(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(defOp)) + if (isa( + defOp)) { + + if (auto constOp = dyn_cast(defOp)) { + if (auto denseAttr = + dyn_cast(constOp.getValue())) { + if (!denseAttr.isSplat()) + return nullptr; + + Attribute splat = denseAttr.getSplatValue(); + + if (auto floatAttr = dyn_cast(splat)) + return floatAttr.getValue().isZero() ? defOp : nullptr; + + if (auto intAttr = dyn_cast(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(op)) - return write.getVector(); - if (auto store = dyn_cast(op)) - return store.getValueToStore(); - return nullptr; - }; - Value vecA = getWrittenVector(opA); - Value vecB = getWrittenVector(opB); + Value vecA = contractionUsersAfterYield(contractARes); + 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, - [&]() { 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 000000000000..8007f5cd8fc8 --- /dev/null +++ b/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} %6, %7, %arg9 + {unroll_shape = array} : !vecA, !vecB into !vecC + %10 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = + ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind} %6, %8, %arg10 + {unroll_shape = array} : !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} + %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} %6, %7, %arg9 + {unroll_shape = array} : !vecA, !vecB into !vecC + %10 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = + ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind} %6, %8, %arg10 + {unroll_shape = array} : !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} %6, %7, %arg9 + {unroll_shape = array} : !vecA, !vecB into !vecC + %10 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = + ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind} %6, %8, %arg10 + {unroll_shape = array} : !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 + } +} + 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 4f0e5c5f3c90..94abf9505185 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 5d32570b21cb..181a789a9adc 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