[Mlir-commits] [mlir] 739a430 - [mlir][x86] Add extra patterns for AMX nano-kernel lowering (acc with zero constant) (#222513)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 30 06:19:47 PDT 2026
Author: Arun Thangamani
Date: 2026-09-30T18:49:39+05:30
New Revision: 739a4306427e1be1cd02058ecb108ce1b2dc2916
URL: https://github.com/llvm/llvm-project/commit/739a4306427e1be1cd02058ecb108ce1b2dc2916
DIFF: https://github.com/llvm/llvm-project/commit/739a4306427e1be1cd02058ecb108ce1b2dc2916.diff
LOG: [mlir][x86] Add extra patterns for AMX nano-kernel lowering (acc with zero constant) (#222513)
This patch extends AMX nano-kernel lowering for `zero-initialization
constant` of input accumulation
(https://github.com/llvm/llvm-project/pull/204327).
Added:
Modified:
mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 365a3fb8e523b..32ab0500a3921 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -27,6 +27,23 @@ using namespace mlir::x86;
namespace {
+// Return true if the operation is a constant dense element
+// attribute.
+static bool isZeroVectorConstant(Operation *op) {
+ auto constantOp = dyn_cast<arith::ConstantOp>(op);
+ if (!constantOp)
+ return false;
+
+ if (!isa<VectorType>(constantOp.getType()))
+ return false;
+
+ auto denseAttr = dyn_cast<DenseElementsAttr>(constantOp.getValue());
+ if (!denseAttr || !denseAttr.isSplat())
+ return false;
+
+ return isZeroIntegerOrFloat(denseAttr.getSplatValue<Attribute>());
+}
+
// Recursively follows single-use values through scf.yield operations
// and returns the first non-yield user result in the contraction chain.
static Value contractionUsersAfterYield(Value v) {
@@ -48,8 +65,7 @@ static Value contractionUsersAfterYield(Value v) {
// Function to collapse the last two dimension (vnni and k) to help the
// amx.tile_load to correctly load the packed element type.
-static Value collapseInnerDims(OpBuilder &builder, mlir::Location loc,
- Value input) {
+static Value collapseInnerDims(OpBuilder &builder, Location loc, Value input) {
ShapedType inputType = cast<ShapedType>(input.getType());
int64_t firstDimToCollapse = inputType.getRank() - 2;
@@ -174,7 +190,7 @@ static LogicalResult validateContractOps(OpBuilder &rewriter,
// The Accumulator dims should be 16 or 1. Like <1x16x16> or <16x16>.
ArrayRef<int64_t> accShape = accTy.getShape();
- llvm::SmallVector<int64_t> nonUnitDimAcc;
+ SmallVector<int64_t> nonUnitDimAcc;
llvm::copy_if(accShape, std::back_inserter(nonUnitDimAcc),
[](int64_t dim) { return (dim != 16 && dim != 1); });
@@ -185,7 +201,7 @@ static LogicalResult validateContractOps(OpBuilder &rewriter,
// <16x16x4>. The vnni dims should be 2 or 4.
VectorType lhsTy = contractOp.getLhsType();
ArrayRef<int64_t> lhsShape = lhsTy.getShape();
- llvm::SmallVector<int64_t> nonUnitDimLhs;
+ SmallVector<int64_t> nonUnitDimLhs;
llvm::copy_if(lhsShape, std::back_inserter(nonUnitDimLhs),
[](int64_t dim) { return (dim != 16 && dim != 1); });
@@ -309,19 +325,19 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
vec2 = vector::ShapeCastOp::create(
rewriter, loc, VectorType::get((16 * offset), ipType), vec2);
- vector::ShuffleOp shuffle1;
- vector::ShuffleOp shuffle2;
+ vector::ShuffleOp writeVal1;
+ vector::ShuffleOp writeVal2;
if (ipType.isBF16()) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({(16 * offset)}, ipType), vec1,
vec2,
ArrayRef<int64_t>{0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9,
41, 10, 42, 11, 43, 16, 48, 17, 49, 18, 50,
19, 51, 24, 56, 25, 57, 26, 58, 27, 59});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({(16 * offset)}, ipType), vec1,
vec2,
ArrayRef<int64_t>{4, 36, 5, 37, 6, 38, 7, 39, 12, 44, 13,
@@ -332,7 +348,7 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
if (ipType.isSignlessInteger(8) || ipType.isF8E5M2() ||
ipType.isF8E4M3FN()) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({(16 * offset)}, ipType), vec1,
vec2,
ArrayRef<int64_t>{
@@ -342,7 +358,7 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
113, 18, 50, 82, 114, 19, 51, 83, 115, 24, 56, 88, 120,
25, 57, 89, 121, 26, 58, 90, 122, 27, 59, 91, 123});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({(16 * offset)}, ipType), vec1,
vec2,
ArrayRef<int64_t>{
@@ -357,9 +373,9 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
Value ivShuff1 = arith::DivUIOp::create(rewriter, loc, iv, cStep);
Value ivShuff2 = arith::AddIOp::create(rewriter, loc, ivShuff1, c16);
- vector::StoreOp::create(rewriter, loc, shuffle1, packedBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal1, packedBuffer,
ValueRange{indxToStoreInBuffer, ivShuff1, c0});
- vector::StoreOp::create(rewriter, loc, shuffle2, packedBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal2, packedBuffer,
ValueRange{indxToStoreInBuffer, ivShuff2, c0});
scf::YieldOp::create(nestedBuilder, loc);
@@ -802,6 +818,11 @@ struct VectorContractToAMXDotProduct
Operation *accReadOp =
traceToVectorReadLikeParentOperation(contractOp.getAcc());
+ bool isAccZeroVectorConstant = isZeroVectorConstant(accReadOp);
+
+ Operation *accWrite =
+ traceToVectorWriteLikeUserOperation(contractOp.getResult());
+
// Only the contract result's first consumer is needed, not the final
// store. This keeps the lowering independent of the epilogue ops (truncf,
// bias add, ReLU, ...) that sit between the contraction and the write.
@@ -955,31 +976,31 @@ struct VectorContractToAMXDotProduct
VectorType::get(16 * (blockingFactor / 2), ipType), subview,
range2);
- vector::ShuffleOp shuffle1;
- vector::ShuffleOp shuffle2;
+ vector::ShuffleOp writeVal1;
+ vector::ShuffleOp writeVal2;
if (blockingFactor == 2) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({16}, ipType), vec1, vec2,
ArrayRef<int64_t>{0, 16, 1, 17, 2, 18, 3, 19, 4, 20, 5, 21,
6, 22, 7, 23});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({16}, ipType), vec1, vec2,
ArrayRef<int64_t>{8, 24, 9, 25, 10, 26, 11, 27, 12, 28, 13,
29, 14, 30, 15, 31});
}
if (blockingFactor == 4) {
- shuffle1 = vector::ShuffleOp::create(
+ writeVal1 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({32}, ipType), vec1, vec2,
ArrayRef<int64_t>{0, 16, 32, 48, 1, 17, 33, 49,
2, 18, 34, 50, 3, 19, 35, 51,
4, 20, 36, 52, 5, 21, 37, 53,
6, 22, 38, 54, 7, 23, 39, 55});
- shuffle2 = vector::ShuffleOp::create(
+ writeVal2 = vector::ShuffleOp::create(
rewriter, loc, VectorType::get({32}, ipType), vec1, vec2,
ArrayRef<int64_t>{8, 24, 40, 56, 9, 25, 41, 57,
10, 26, 42, 58, 11, 27, 43, 59,
@@ -990,9 +1011,9 @@ struct VectorContractToAMXDotProduct
auto rem = arith::DivUIOp::create(
rewriter, loc, rewriter.getIndexType(), iv, step);
- vector::StoreOp::create(rewriter, loc, shuffle1, packedBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal1, packedBuffer,
ValueRange{rem, c0});
- vector::StoreOp::create(rewriter, loc, shuffle2, packedBuffer,
+ vector::StoreOp::create(rewriter, loc, writeVal2, packedBuffer,
ValueRange{rem, nextStoreIndx});
scf::YieldOp::create(nestedBuilder, loc);
@@ -1025,7 +1046,7 @@ struct VectorContractToAMXDotProduct
amx::TileStoreOp::create(rewriter, loc, resultBuffer, ValueRange{c0, c0},
dp);
- auto vectorType = mlir::VectorType::get({16, 16}, opType);
+ auto vectorType = VectorType::get({16, 16}, opType);
int64_t srcRank =
(dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
Value padding = ub::PoisonOp::create(rewriter, loc, opType);
@@ -1038,7 +1059,7 @@ struct VectorContractToAMXDotProduct
map, inBounds);
Value resultOp = contractionUsersAfterYield(contractOp.getResult());
- if (auto vecType = llvm::dyn_cast<VectorType>(resultOp.getType()))
+ if (auto vecType = dyn_cast<VectorType>(resultOp.getType()))
vecRow = vector::ShapeCastOp::create(rewriter, loc, vecType, vecRow);
rewriter.replaceAllUsesWith(resultOp, vecRow);
@@ -1069,12 +1090,13 @@ struct VectorContractToAMXDotProduct
loopLists.push_back(dyn_cast<scf::ForOp>(parent));
- if (accReadOp->getBlock() == parent->getBlock()) {
+ if (resultBlock == parent->getBlock()) {
break;
}
current = parent;
}
+
if (loopLists.size() > 2 || loopLists.size() == 0)
return rewriter.notifyMatchFailure(
contractOp, "Rewrite is supported until reduction loop depth of 2.");
@@ -1110,10 +1132,9 @@ struct VectorContractToAMXDotProduct
// Retrive all the contaction operation within the loop.
SmallVector<vector::ContractionOp> ops;
- for (mlir::Operation &op : loopLists[0].getBody()->getOperations()) {
-
- if (auto contract = llvm::dyn_cast<mlir::vector::ContractionOp>(op)) {
+ for (Operation &op : loopLists[0].getBody()->getOperations()) {
+ if (auto contract = dyn_cast<vector::ContractionOp>(op)) {
LogicalResult validate = validateContractOps(
rewriter, contract, dimValue, srcBuffLhs, srcBuffRhs, true);
@@ -1124,6 +1145,14 @@ struct VectorContractToAMXDotProduct
"the re-write conditions either the dimensions are "
"wrong or MemRef source are
diff erent or many users.");
+ Operation *contractReadOp =
+ traceToVectorReadLikeParentOperation(contract.getAcc());
+
+ if (isAccZeroVectorConstant != isZeroVectorConstant(contractReadOp))
+ return rewriter.notifyMatchFailure(
+ contractOp, "The input acc should be a zero constant or "
+ "transfer_read for all contracts");
+
ops.push_back(contract);
}
}
@@ -1192,10 +1221,9 @@ struct VectorContractToAMXDotProduct
bool isInnerLoopUBHasOddQuot = false;
int64_t ubVal = 16 * blockingFactor;
- mlir::Value ub = innerLoop.getUpperBound();
- if (auto constOp = ub.getDefiningOp<mlir::arith::ConstantOp>()) {
- if (auto intAttr =
- llvm::dyn_cast<mlir::IntegerAttr>(constOp.getValue())) {
+ Value ub = innerLoop.getUpperBound();
+ if (auto constOp = ub.getDefiningOp<arith::ConstantOp>()) {
+ if (auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue())) {
ubVal = intAttr.getInt();
}
}
@@ -1340,10 +1368,9 @@ struct VectorContractToAMXDotProduct
bool isInnerLoopUBHasOddQuot = false;
int64_t ubVal = 16 * blockingFactor;
- mlir::Value ub = innerLoop.getUpperBound();
- if (auto constOp = ub.getDefiningOp<mlir::arith::ConstantOp>()) {
- if (auto intAttr =
- llvm::dyn_cast<mlir::IntegerAttr>(constOp.getValue())) {
+ Value ub = innerLoop.getUpperBound();
+ if (auto constOp = ub.getDefiningOp<arith::ConstantOp>()) {
+ if (auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue())) {
ubVal = intAttr.getInt();
}
}
@@ -1418,22 +1445,36 @@ struct VectorContractToAMXDotProduct
Value srcBuffAcc;
SmallVector<Value> indicesAcc;
- llvm::TypeSwitch<Operation *>(accReadOp).Case<TransferReadOp, LoadOp>(
- [&](auto readOp) {
- srcBuffAcc = readOp.getOperand(0);
+ if (isAccZeroVectorConstant) {
+ llvm::TypeSwitch<Operation *>(accWrite)
+ .Case<vector::TransferWriteOp, vector::StoreOp>([&](auto writeOp) {
+ srcBuffAcc = writeOp->getOperand(1);
- auto indices = readOp.getIndices();
- indicesAcc.reserve(indices.size());
+ auto indices = writeOp.getIndices();
+ indicesAcc.reserve(indices.size());
- llvm::transform(indices, std::back_inserter(indicesAcc),
- [&](OpFoldResult ofr) {
- return mlir::getValueOrCreateConstantIndexOp(
- rewriter, loc, ofr);
- });
- });
+ llvm::transform(
+ indices, std::back_inserter(indicesAcc), [&](OpFoldResult ofr) {
+ return getValueOrCreateConstantIndexOp(rewriter, loc, ofr);
+ });
+ });
+ } else {
+ llvm::TypeSwitch<Operation *>(accReadOp).Case<TransferReadOp, LoadOp>(
+ [&](auto readOp) {
+ srcBuffAcc = readOp.getOperand(0);
+
+ auto indices = readOp.getIndices();
+ indicesAcc.reserve(indices.size());
+
+ llvm::transform(indices, std::back_inserter(indicesAcc),
+ [&](OpFoldResult ofr) {
+ return mlir::getValueOrCreateConstantIndexOp(
+ rewriter, loc, ofr);
+ });
+ });
+ }
- auto outputShapes =
- mlir::cast<mlir::MemRefType>(srcBuffAcc.getType()).getShape();
+ auto outputShapes = cast<MemRefType>(srcBuffAcc.getType()).getShape();
unsigned int M = outputShapes[outputShapes.size() - 2];
unsigned int N = outputShapes[outputShapes.size() - 1];
@@ -1451,73 +1492,81 @@ struct VectorContractToAMXDotProduct
k++;
}
}
+
auto c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
auto c16 = arith::ConstantIndexOp::create(rewriter, loc, 16);
auto one = arith::ConstantIndexOp::create(rewriter, loc, 1);
auto nBound = arith::ConstantIndexOp::create(rewriter, loc, N);
- // Create a loop that iterates over the MxN memerf, retrives two rows +
- // shuffle them, add up the C element values and stores them to temp buffer.
- scf::ForOp::create(
- rewriter, loc, c0, nBound, one, ValueRange{},
- [&](OpBuilder &nestedBuilder, Location loc, Value iv,
- ValueRange iterArgs) {
- auto row =
- vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
- resultBuffer, ValueRange{iv, c0});
-
- auto row2 =
- vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
- resultBuffer, ValueRange{iv, c16});
-
- Value shuffle1 = row;
- Value shuffle2 = row2;
-
- if (!isVnni) {
- shuffle1 = vector::ShuffleOp::create(
- rewriter, loc, VectorType::get(16, opType), row, row2,
- ArrayRef<int64_t>{0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20,
- 21, 22, 23});
-
- shuffle2 = vector::ShuffleOp::create(
- rewriter, loc, VectorType::get(16, opType), row, row2,
- ArrayRef<int64_t>{8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15,
- 28, 29, 30, 31});
- }
- indicesAcc[indicesAcc.size() - 2] = iv;
- indicesAcc[indicesAcc.size() - 1] = c0;
+ if (!isVnni || !isAccZeroVectorConstant) {
+ // Create a loop that iterates over the MxN memerf, retrives two rows +
+ // shuffle them, add up the C element values and stores them to temp
+ // buffer.
+ scf::ForOp::create(
+ rewriter, loc, c0, nBound, one, ValueRange{},
+ [&](OpBuilder &nestedBuilder, Location loc, Value iv,
+ ValueRange iterArgs) {
+ auto row = vector::LoadOp::create(rewriter, loc,
+ VectorType::get(16, opType),
+ resultBuffer, ValueRange{iv, c0});
+
+ auto row2 = vector::LoadOp::create(
+ rewriter, loc, VectorType::get(16, opType), resultBuffer,
+ ValueRange{iv, c16});
+
+ Value writeVal1 = row;
+ Value writeVal2 = row2;
+
+ if (!isVnni) {
+ writeVal1 = vector::ShuffleOp::create(
+ rewriter, loc, VectorType::get(16, opType), row, row2,
+ ArrayRef<int64_t>{0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20,
+ 21, 22, 23});
+
+ writeVal2 = vector::ShuffleOp::create(
+ rewriter, loc, VectorType::get(16, opType), row, row2,
+ ArrayRef<int64_t>{8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14,
+ 15, 28, 29, 30, 31});
+ }
- Value valueCRow1 =
- vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
- srcBuffAcc, indicesAcc);
- indicesAcc[indicesAcc.size() - 1] = c16;
+ if (!isAccZeroVectorConstant) {
+ indicesAcc[indicesAcc.size() - 2] = iv;
+ indicesAcc[indicesAcc.size() - 1] = c0;
- Value valueCRow2 =
- vector::LoadOp::create(rewriter, loc, VectorType::get(16, opType),
- srcBuffAcc, indicesAcc);
+ Value valueCRow1 = vector::LoadOp::create(
+ rewriter, loc, VectorType::get(16, opType), srcBuffAcc,
+ indicesAcc);
+ indicesAcc[indicesAcc.size() - 1] = c16;
- Value addOp;
- Value addOp2;
+ Value valueCRow2 = vector::LoadOp::create(
+ rewriter, loc, VectorType::get(16, opType), srcBuffAcc,
+ indicesAcc);
- if (ipType.isBF16() || ipType.isF8E5M2() || ipType.isF8E4M3FN()) {
- addOp = arith::AddFOp::create(rewriter, loc, shuffle1, valueCRow1);
+ if (ipType.isBF16() || ipType.isF8E5M2() || ipType.isF8E4M3FN()) {
+ writeVal1 =
+ arith::AddFOp::create(rewriter, loc, writeVal1, valueCRow1);
- addOp2 = arith::AddFOp::create(rewriter, loc, shuffle2, valueCRow2);
- }
+ writeVal2 =
+ arith::AddFOp::create(rewriter, loc, writeVal2, valueCRow2);
+ }
- if (ipType.isSignlessInteger(8)) {
- addOp = arith::AddIOp::create(rewriter, loc, shuffle1, valueCRow1);
+ if (ipType.isSignlessInteger(8)) {
+ writeVal1 =
+ arith::AddIOp::create(rewriter, loc, writeVal1, valueCRow1);
- addOp2 = arith::AddIOp::create(rewriter, loc, shuffle2, valueCRow2);
- }
+ writeVal2 =
+ arith::AddIOp::create(rewriter, loc, writeVal2, valueCRow2);
+ }
+ }
- vector::StoreOp::create(rewriter, loc, addOp, resultBuffer,
- ValueRange{iv, c0});
- vector::StoreOp::create(rewriter, loc, addOp2, resultBuffer,
- ValueRange{iv, c16});
+ vector::StoreOp::create(rewriter, loc, writeVal1, resultBuffer,
+ ValueRange{iv, c0});
+ vector::StoreOp::create(rewriter, loc, writeVal2, resultBuffer,
+ ValueRange{iv, c16});
- scf::YieldOp::create(nestedBuilder, loc);
- });
+ scf::YieldOp::create(nestedBuilder, loc);
+ });
+ }
SmallVector<Value> writeResults;
for (unsigned int i = 0; i < M; i = i + 16) {
@@ -1525,7 +1574,7 @@ struct VectorContractToAMXDotProduct
Value indexOp_i = arith::ConstantIndexOp::create(rewriter, loc, i);
Value indexOp_j = arith::ConstantIndexOp::create(rewriter, loc, j);
- auto vectorType = mlir::VectorType::get({16, 16}, opType);
+ auto vectorType = VectorType::get({16, 16}, opType);
int64_t srcRank =
(dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
@@ -1547,9 +1596,9 @@ struct VectorContractToAMXDotProduct
Value vecRow = writeResults[i];
Value resultWriteOp = contractionUsersAfterYield(contOp.getResult());
- if (auto vecType = llvm::dyn_cast<VectorType>(resultWriteOp.getType()))
- vecRow = mlir::vector::ShapeCastOp::create(rewriter, loc, vecType,
- writeResults[i]);
+ if (auto vecType = dyn_cast<VectorType>(resultWriteOp.getType()))
+ vecRow = vector::ShapeCastOp::create(rewriter, loc, vecType,
+ writeResults[i]);
rewriter.replaceAllUsesWith(resultWriteOp, vecRow);
}
diff --git a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
index 093e97af0520e..8e2791ed9153c 100644
--- a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
+++ b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
@@ -693,6 +693,200 @@ module attributes {transform.with_named_sequence} {
// -----
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @matmul_int8_loop_cst_acc(%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>) {
+ %cst = arith.constant dense<0> : vector<16x16xi32>
+ %0 = ub.poison : i32
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xi32> 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 %c64 step %c16 iter_args(%arg6 = %cst, %arg7 = %cst) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
+ memref<64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+
+ %5 = arith.addi %4#1, %3 : vector<16x16xi32>
+ vector.transfer_write %5, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ %6 = arith.addi %4#0, %2 : vector<16x16xi32>
+ vector.transfer_write %6, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @matmul_int8_loop_cst_acc
+// CHECK-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK-3: x86.amx.tile_load
+// CHECK-2: x86.amx.tile_muli
+// CHECK: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
+// CHECK-NOT: scf.for {{.*}} vector<16x16xi32>, vector<16x16xi32>
+// CHECK-NOT: vector.contract
+
+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.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<1x16x32xbf16>
+!vecB = vector<1x32x16xbf16>
+!vecC = vector<16x16xf32>
+!memrefA = memref<1x32x32xbf16, strided<[6144, 96, 1], offset: ?>>
+!memrefB = memref<1x32x32xbf16, strided<[12288, 128, 1], offset: ?>>
+!memrefC = memref<32x32xf32, strided<[128, 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 @online_packing_bf16_loop_acc_after_vc(%arg0: memref<16x64x96xbf16>, %arg1: memref<16x96x128xbf16>, %arg2: memref<64x128xf32>) -> memref<64x128xf32> {
+ %cst = arith.constant dense<0.000000e+00> : vector<16x16xf32>
+ %0 = ub.poison : f32
+ %1 = ub.poison : bf16
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c96 = arith.constant 96 : index
+ %c32 = arith.constant 32 : index
+ %c1 = arith.constant 1 : index
+ scf.for %arg3 = %c0 to %c64 step %c32 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %6:4 = scf.for %arg5 = %c0 to %c16 step %c1 iter_args(%arg6 = %cst, %arg7 = %cst, %arg8 = %cst, %arg9 = %cst) -> (!vecC, !vecC, !vecC, !vecC) {
+ %7:4 = scf.for %arg10 = %c0 to %c96 step %c32 iter_args(%arg11 = %arg6, %arg12 = %arg7, %arg13 = %arg8, %arg14 = %arg9) -> (!vecC, !vecC, !vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg5, %arg3, %arg10] [1, 32, 32] [1, 1, 1] :
+ memref<16x64x96xbf16> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg10, %arg4] [1, 32, 32] [1, 1, 1] :
+ memref<16x96x128xbf16> to !memrefB
+ %8 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecA
+ %9 = vector.transfer_read %subview_0[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecA
+ %10 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecB
+ %11 = vector.transfer_read %subview_1[%c0, %c0, %c16], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecB
+
+ %12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %8, %10, %arg11 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
+ %13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %8, %11, %arg12 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
+ %14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %9, %10, %arg13 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
+ %15 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %9, %11, %arg14 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
+
+ scf.yield %12, %13, %14, %15 : !vecC, !vecC, !vecC, !vecC
+ }
+ scf.yield %7#0, %7#1, %7#2, %7#3 : !vecC, !vecC, !vecC, !vecC
+ }
+
+ %cst_1 = arith.constant dense<4.000000e+00> : vector<16x16xf32>
+ %d = arith.addf %6#3, %cst_1 : vector<16x16xf32>
+ %c = arith.addf %6#2, %cst_1 : vector<16x16xf32>
+ %b = arith.addf %6#1, %cst_1 : vector<16x16xf32>
+ %a = arith.addf %6#0, %cst_1 : vector<16x16xf32>
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [32, 32] [1, 1] :
+ memref<64x128xf32> to !memrefC
+ vector.transfer_write %d, %subview[%c16, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %c, %subview[%c16, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %b, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %a, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+ %alloc = memref.alloc() : memref<64x128xf32>
+ memref.copy %arg2, %alloc : memref<64x128xf32> to memref<64x128xf32>
+ return %alloc : memref<64x128xf32>
+}
+
+// CHECK-LABEL: @online_packing_bf16_loop_acc_after_vc
+// CHECK-COUNT-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
+// CHECK-COUNT-4: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
+// CHECK: vector.transfer_read {{.*}} vector<32xbf16>
+// CHECK: vector.transfer_read {{.*}} vector<32xbf16>
+// CHECK: vector.shuffle{{.*}}[0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9, 41, 10, 42, 11, 43, 16, 48, 17, 49, 18, 50, 19, 51, 24, 56, 25, 57, 26, 58, 27, 59] : vector<32xbf16>, vector<32xbf16>
+// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 5, 37, 6, 38, 7, 39, 12, 44, 13, 45, 14, 46, 15, 47, 20, 52, 21, 53, 22, 54, 23, 55, 28, 60, 29, 61, 30, 62, 31, 63] : vector<32xbf16>, vector<32xbf16>
+// CHECK: x86.amx.tile_load
+// CHECK: x86.amx.tile_mulf
+// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
+// 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-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
+// CHECK-NOT: vector.contract
+
+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.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
!vecAB = vector<16x16x4xf8E4M3FN>
!vecC = vector<16x16xf32>
!memrefA = memref<16x16x4xf8E4M3FN, strided<[256, 4, 1], offset: ?>>
@@ -2329,6 +2523,88 @@ func.func @negative_acc_tensor(
// CHECK-NOT: x86.amx
// CHECK: vector.contract
+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.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @negative_mixed_acc(%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>) {
+ %cst_1 = arith.constant dense<0> : vector<16x16xi32>
+ %0 = ub.poison : i32
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xi32> 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 %c64 step %c16 iter_args(%arg6 = %2, %arg7 = %cst_1) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
+ memref<64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+ %add = arith.addi %4#1, %3 : vector<16x16xi32>
+ vector.transfer_write %add, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @negative_mixed_acc
+// CHECK-NOT: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK-NOT: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK-NOT: x86.amx.tile_load
+// CHECK-NOT: x86.amx.tile_muli
+// CHECK-NOT: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
+
+
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
More information about the Mlir-commits
mailing list