[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.
This commit is contained in:
Arun Thangamani
GitHub
parent
676d48fb49
commit
b32d4f0c45
12 changed files with 550 additions and 52 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -52,6 +52,11 @@ void mlir::transform::ApplyVectorContractToAMXDotProductPatternsOp::
|
|||
x86::populateVectorContractToAMXDotProductPatterns(patterns);
|
||||
}
|
||||
|
||||
void mlir::transform::ApplyMoveAccumulatorForContractLoopPatternsOp::
|
||||
populatePatterns(RewritePatternSet &patterns) {
|
||||
x86::populateMoveAccumulatorForContractLoopPatterns(patterns);
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// Transform op registration
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ add_mlir_dialect_library(MLIRX86Transforms
|
|||
SinkVectorProducerOps.cpp
|
||||
ShuffleVectorFMAOps.cpp
|
||||
VectorContractToAMXDotProduct.cpp
|
||||
MoveAccumulatorForContractLoop.cpp
|
||||
|
||||
LINK_LIBS PUBLIC
|
||||
MLIRArithDialect
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
|
|
@ -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<arith::ConstantOp>(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(
|
||||
|
|
|
|||
|
|
@ -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<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(
|
||||
|
|
|
|||
|
|
@ -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 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();
|
||||
}
|
||||
|
|
|
|||
323
mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
Normal file
323
mlir/test/Dialect/X86/move-acc-for-contract-loop.mlir
Normal file
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue