[Mlir-commits] [mlir] [wip][mlir][x86vector] Shuffle BF16 vector.contract output for Flat layout. (PR #174590)
Arun Thangamani
llvmlistbot at llvm.org
Thu Jan 8 07:25:49 PST 2026
https://github.com/arun-thmn updated https://github.com/llvm/llvm-project/pull/174590
>From 5ddf7e96ddeb975534668b0bd0c991a4c483a171 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 6 Jan 2026 05:44:34 -0800
Subject: [PATCH 1/4] initial commit for shuffle VC output
---
.../TransformOps/X86VectorTransformOps.td | 11 +
.../mlir/Dialect/X86Vector/Transforms.h | 3 +
.../Dialect/X86Vector/Utils/X86VectorUtils.h | 26 ++
.../TransformOps/X86VectorTransformOps.cpp | 5 +
.../X86Vector/Transforms/CMakeLists.txt | 1 +
.../ShuffleBF16VectorContractResult.cpp | 150 ++++++++
.../X86Vector/Utils/X86VectorUtils.cpp | 325 ++++++++++++++++++
.../shuffle-bf16-vector-contract-result.mlir | 66 ++++
8 files changed, 587 insertions(+)
create mode 100644 mlir/lib/Dialect/X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp
create mode 100644 mlir/test/Dialect/shuffle-bf16-vector-contract-result.mlir
diff --git a/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td b/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
index 3c73eadf82167..00c611a9f3a7a 100644
--- a/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
@@ -60,6 +60,17 @@ def ApplySinkVectorProducerOpsPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplyShuffleBF16VectorContractResultPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.x86vector.shuffle_bf16_vector_contract_result",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Collect patterns to shuffle results of flat layout BF16 type
+ vector.contract operations.
+ }];
+
+ let assemblyFormat = "attr-dict";
+}
+
#endif // X86VECTOR_TRANSFORM_OPS
diff --git a/mlir/include/mlir/Dialect/X86Vector/Transforms.h b/mlir/include/mlir/Dialect/X86Vector/Transforms.h
index c25cdaf2d9428..538ffaac79998 100644
--- a/mlir/include/mlir/Dialect/X86Vector/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86Vector/Transforms.h
@@ -100,6 +100,9 @@ void populateVectorContractBF16ToFMAPatterns(RewritePatternSet &patterns);
// range by placing them at their earliest legal use site.
void populateSinkVectorProducerOpsPatterns(RewritePatternSet &patterns);
+void populateShuffleBF16VectorContractResultPatterns(
+ RewritePatternSet &patterns);
+
//===----------------------------------------------------------------------===//
/// Helpers extracted from:
/// - clang/lib/Headers/avxintrin.h
diff --git a/mlir/include/mlir/Dialect/X86Vector/Utils/X86VectorUtils.h b/mlir/include/mlir/Dialect/X86Vector/Utils/X86VectorUtils.h
index 2de9a3122cbd9..7267c5c2032d9 100644
--- a/mlir/include/mlir/Dialect/X86Vector/Utils/X86VectorUtils.h
+++ b/mlir/include/mlir/Dialect/X86Vector/Utils/X86VectorUtils.h
@@ -9,6 +9,9 @@
#ifndef MLIR_DIALECT_X86VECTOR_UTILS_X86VECTORUTILS_H_
#define MLIR_DIALECT_X86VECTOR_UTILS_X86VECTORUTILS_H_
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
#include "mlir/Support/LogicalResult.h"
#include "llvm/ADT/ArrayRef.h"
#include <cstdint>
@@ -26,6 +29,29 @@ namespace x86vector {
bool isInVnniLayout(Operation *op, llvm::ArrayRef<AffineMap> indexingMaps,
std::optional<unsigned> blockingFactor = std::nullopt);
+bool validatePairVectorContract(vector::ContractionOp contractOp,
+ vector::ContractionOp pairContOp,
+ bool rhsHasMultipleNonUnitDims,
+ int64_t nonUnitDimValue);
+
+Operation *traceToVectorReadLikeParentOperation(mlir::Value v);
+
+Operation *traceToVectorWriteLikeUserOperation(mlir::Value v);
+
+void shuffleBeforeWriteLikeOp(PatternRewriter &rewriter, Operation *op,
+ Operation *op1, int64_t nonUnitDimAcc,
+ VectorType accTy);
+
+void shuffleAfterReadLikeOp(PatternRewriter &rewriter, Operation *op,
+ Operation *op1, vector::ContractionOp contractOp,
+ vector::ContractionOp pairContractOp,
+ int64_t nonUnitDimAcc, VectorType accTy);
+
+void shuffleNonUnitDimOperand(PatternRewriter &rewriter, Operation *op,
+ Operation *op1, vector::ContractionOp contractOp,
+ vector::ContractionOp pairContractOp,
+ int64_t nonUnitDimAcc, VectorType Ty);
+
} // namespace x86vector
} // namespace mlir
diff --git a/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp b/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
index e77d30c9c5ffb..e40ddd3a4b1c0 100644
--- a/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
@@ -42,6 +42,11 @@ void mlir::transform::ApplySinkVectorProducerOpsPatternsOp::populatePatterns(
x86vector::populateSinkVectorProducerOpsPatterns(patterns);
}
+void mlir::transform::ApplyShuffleBF16VectorContractResultPatternsOp::
+ populatePatterns(RewritePatternSet &patterns) {
+ x86vector::populateShuffleBF16VectorContractResultPatterns(patterns);
+}
+
//===----------------------------------------------------------------------===//
// Transform op registration
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt b/mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt
index bbd9be880eb0a..acbc7fcfb635e 100644
--- a/mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt
@@ -5,6 +5,7 @@ add_mlir_dialect_library(MLIRX86VectorTransforms
VectorContractToPackedTypeDotProduct.cpp
VectorContractBF16ToFMA.cpp
SinkVectorProducerOps.cpp
+ ShuffleBF16VectorContractResult.cpp
LINK_LIBS PUBLIC
MLIRArithDialect
diff --git a/mlir/lib/Dialect/X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp b/mlir/lib/Dialect/X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp
new file mode 100644
index 0000000000000..2081de9f2be91
--- /dev/null
+++ b/mlir/lib/Dialect/X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp
@@ -0,0 +1,150 @@
+//===-
+//ShuffleBF16VectorContractResult.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/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Dialect/X86Vector/Transforms.h"
+#include "mlir/Dialect/X86Vector/Utils/X86VectorUtils.h"
+#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/PatternMatch.h"
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/Support/Casting.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+using namespace mlir::x86vector;
+
+struct ShuffleBF16VectorContractResult
+ : 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.");
+
+ // TODO: Move this validation to a common utility folder. Planned to
+ // do once (code refactoring), all architecture specific nanokernel
+ // passes are merged into the repo.
+ VectorType lhsTy = contractOp.getLhsType();
+ if (!lhsTy.getElementType().isBF16())
+ return rewriter.notifyMatchFailure(contractOp,
+ "Only BF16 lowering is supported.");
+
+ if (isInVnniLayout(contractOp.getOperation(),
+ contractOp.getIndexingMapsArray(),
+ /*blockingFactor=*/2))
+ return rewriter.notifyMatchFailure(contractOp,
+ "Input matrices in VNNI format.");
+
+ VectorType accTy = dyn_cast<VectorType>(contractOp.getAccType());
+ if (!accTy)
+ return rewriter.notifyMatchFailure(contractOp, "Wrong accmulator type.");
+
+ if (!accTy.getElementType().isF32())
+ return rewriter.notifyMatchFailure(
+ contractOp, "Only F32 acumulation supported for BF16 type.");
+
+ ArrayRef<int64_t> accShape = accTy.getShape();
+ llvm::SmallVector<int64_t> nonUnitDimAcc;
+ llvm::copy_if(accShape, std::back_inserter(nonUnitDimAcc),
+ [](int64_t dim) { return dim != 1; });
+
+ if (nonUnitDimAcc.size() != 1)
+ return rewriter.notifyMatchFailure(
+ contractOp, "A or B should be a non-unit dim in acc.");
+
+ int64_t nonUnitDimValue = nonUnitDimAcc.front();
+
+ if (nonUnitDimValue != 8 && nonUnitDimValue != 16)
+ return rewriter.notifyMatchFailure(
+ contractOp, "The accumulator dimension should be 8 or 16");
+
+ ArrayRef<int64_t> lhsShape = lhsTy.getShape();
+ llvm::SmallVector<int64_t> nonUnitDimLhs;
+ llvm::copy_if(lhsShape, std::back_inserter(nonUnitDimLhs),
+ [](int64_t dim) { return dim != 1; });
+
+ VectorType rhsTy = contractOp.getRhsType();
+ ArrayRef<int64_t> rhsShape = rhsTy.getShape();
+ llvm::SmallVector<int64_t> nonUnitDimRhs;
+ llvm::copy_if(rhsShape, std::back_inserter(nonUnitDimRhs),
+ [](int64_t dim) { return dim != 1; });
+
+ if ((nonUnitDimValue == 16) && (nonUnitDimLhs.size() - 1) > 0 &&
+ (nonUnitDimRhs.size() - 1) > 0)
+ return rewriter.notifyMatchFailure(contractOp,
+ "Excepts unit dimensions for either "
+ "LHS or RHS shape.");
+ if (nonUnitDimValue == 8 && nonUnitDimLhs.size() > 0 &&
+ nonUnitDimRhs.size() > 0)
+ return rewriter.notifyMatchFailure(contractOp,
+ "Excepts unit dimensions for either "
+ "LHS or RHS shape.");
+
+ vector::ContractionOp pairContractOp;
+ bool rhsHasMultipleNonUnitDims = nonUnitDimValue == 16
+ ? (nonUnitDimRhs.size() - 1) > 0
+ : nonUnitDimRhs.size() > 0;
+
+ Operation *nextOp = contractOp;
+ while ((nextOp = nextOp->getNextNode())) {
+ auto contOp = dyn_cast<vector::ContractionOp>(nextOp);
+
+ if (!contOp)
+ continue;
+
+ if (validatePairVectorContract(
+ contractOp, contOp, rhsHasMultipleNonUnitDims, nonUnitDimValue)) {
+ pairContractOp = contOp;
+ break;
+ }
+ }
+
+ if (!pairContractOp)
+ return failure();
+
+ Operation *accReadOp0 =
+ traceToVectorReadLikeParentOperation(contractOp.getAcc());
+ Operation *accReadOp1 =
+ traceToVectorReadLikeParentOperation(pairContractOp.getAcc());
+
+ Operation *resultWriteOp0 =
+ traceToVectorWriteLikeUserOperation(contractOp.getResult());
+ Operation *resultWriteOp1 =
+ traceToVectorWriteLikeUserOperation(pairContractOp.getResult());
+
+ if (!accReadOp0 || !accReadOp1)
+ return failure();
+
+ if (!resultWriteOp0 || !resultWriteOp1)
+ return failure();
+
+ shuffleAfterReadLikeOp(rewriter, accReadOp0, accReadOp1, contractOp,
+ pairContractOp, nonUnitDimValue, accTy);
+ shuffleBeforeWriteLikeOp(rewriter, resultWriteOp0, resultWriteOp1,
+ nonUnitDimValue, accTy);
+
+ return success();
+ }
+};
+
+void x86vector::populateShuffleBF16VectorContractResultPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<ShuffleBF16VectorContractResult>(patterns.getContext());
+}
diff --git a/mlir/lib/Dialect/X86Vector/Utils/X86VectorUtils.cpp b/mlir/lib/Dialect/X86Vector/Utils/X86VectorUtils.cpp
index ccb2e92fdd9e2..65cd0773160ef 100644
--- a/mlir/lib/Dialect/X86Vector/Utils/X86VectorUtils.cpp
+++ b/mlir/lib/Dialect/X86Vector/Utils/X86VectorUtils.cpp
@@ -10,11 +10,18 @@
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Types.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Casting.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include <cassert>
+
namespace mlir {
namespace x86vector {
@@ -104,5 +111,323 @@ bool isInVnniLayout(Operation *op, ArrayRef<AffineMap> indexingMaps,
return true;
}
+struct ShuffleMasks {
+ llvm::ArrayRef<int64_t> maskLo;
+ llvm::ArrayRef<int64_t> maskHi;
+};
+
+inline ShuffleMasks getShuffleMasks(int64_t nonUnitDimAcc) {
+ // We only support these two layouts for now.
+ assert((nonUnitDimAcc == 8 || nonUnitDimAcc == 16) &&
+ "Unsupported nonUnitDimAcc value");
+
+ static constexpr int64_t maskLo8[] = {0, 8, 1, 9, 2, 10, 3, 11};
+ static constexpr int64_t maskHi8[] = {4, 12, 5, 13, 6, 14, 7, 15};
+
+ static constexpr int64_t maskLo16[] = {0, 1, 2, 3, 16, 17, 18, 19,
+ 4, 5, 6, 7, 20, 21, 22, 23};
+ static constexpr int64_t maskHi16[] = {8, 9, 10, 11, 24, 25, 26, 27,
+ 12, 13, 14, 15, 28, 29, 30, 31};
+
+ if (nonUnitDimAcc == 16)
+ return {maskLo16, maskHi16};
+
+ // nonUnitDimAcc == 8
+ return {maskLo8, maskHi8};
+}
+
+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)) {
+ return defOp;
+ }
+
+ if (isa<vector::ShapeCastOp, vector::ShuffleOp>(defOp)) {
+ return nullptr;
+ }
+
+ // Continue tracing (accumulators usually forward the value)
+ if (defOp->getNumOperands() == 1) {
+ v = defOp->getOperand(0);
+ continue;
+ }
+
+ return nullptr;
+ }
+
+ // Case 2: BlockArgument (scf.for iter_arg)
+ if (auto barg = dyn_cast<BlockArgument>(v)) {
+ auto *parentOp = barg.getOwner()->getParentOp();
+
+ if (auto forOp = dyn_cast<scf::ForOp>(parentOp)) {
+ unsigned argNum = barg.getArgNumber();
+
+ // arg0 = induction variable (not an iter_arg)
+ if (argNum == 0)
+ return nullptr;
+
+ unsigned iterIdx = argNum - 1;
+ v = forOp.getInitArgs()[iterIdx];
+ continue;
+ }
+
+ return nullptr;
+ }
+
+ return nullptr;
+ }
+}
+
+Operation *traceToVectorWriteLikeUserOperation(Value v) {
+ for (OpOperand &use : v.getUses()) {
+ Operation *user = use.getOwner();
+
+ // --- TERMINAL OPS ---
+ if (isa<vector::TransferWriteOp>(user) || isa<vector::StoreOp>(user)) {
+ return user;
+ }
+
+ if (isa<vector::ShapeCastOp, vector::ShuffleOp>(user)) {
+ return nullptr;
+ }
+
+ // --- SCF YIELD ---
+ if (auto yield = dyn_cast<scf::YieldOp>(user)) {
+ Operation *parent = yield->getParentOp();
+ unsigned idx = use.getOperandNumber();
+ if (auto *res =
+ traceToVectorWriteLikeUserOperation(parent->getResult(idx)))
+ return res;
+ continue;
+ }
+
+ // --- SCF FOR ---
+ if (auto forOp = dyn_cast<scf::ForOp>(user)) {
+ unsigned idx = use.getOperandNumber();
+ if (auto *res = traceToVectorWriteLikeUserOperation(forOp.getResult(idx)))
+ return res;
+ continue;
+ }
+
+ // --- GENERIC CASE ---
+ for (Value res : user->getResults()) {
+ if (auto *found = traceToVectorWriteLikeUserOperation(res))
+ return found;
+ }
+ }
+
+ return nullptr;
+}
+
+static void rewriteUses(mlir::Value oldVal, mlir::Value newVal,
+ mlir::Operation *targetContract,
+ mlir::PatternRewriter &rewriter) {
+ for (mlir::OpOperand &use : llvm::make_early_inc_range(oldVal.getUses())) {
+
+ mlir::Operation *user = use.getOwner();
+
+ // if (user == targetContract ||
+ if (mlir::isa<mlir::vector::ContractionOp>(user) ||
+ mlir::isa<mlir::scf::ForOp>(user)) {
+ use.set(newVal);
+ }
+ }
+}
+
+void shuffleAfterReadLikeOp(mlir::PatternRewriter &rewriter,
+ mlir::Operation *opA, mlir::Operation *opB,
+ mlir::vector::ContractionOp contractA,
+ mlir::vector::ContractionOp contractB,
+ int64_t nonUnitDimAcc, mlir::VectorType accTy) {
+ mlir::Operation *insertAfter = opA->isBeforeInBlock(opB) ? opB : opA;
+
+ rewriter.setInsertionPointAfter(insertAfter);
+ mlir::Location loc = insertAfter->getLoc();
+
+ auto elemTy = accTy.getElementType();
+ auto flatTy = mlir::VectorType::get(nonUnitDimAcc, elemTy);
+
+ auto castA = mlir::vector::ShapeCastOp::create(rewriter, loc, flatTy,
+ opA->getResult(0));
+
+ auto castB = mlir::vector::ShapeCastOp::create(rewriter, loc, flatTy,
+ opB->getResult(0));
+
+ auto masks = getShuffleMasks(nonUnitDimAcc);
+
+ auto shuffleLo = mlir::vector::ShuffleOp::create(rewriter, loc, flatTy, castA,
+ castB, masks.maskLo);
+
+ auto shuffleHi = mlir::vector::ShuffleOp::create(rewriter, loc, flatTy, castA,
+ castB, masks.maskHi);
+
+ auto newAccA =
+ mlir::vector::ShapeCastOp::create(rewriter, loc, accTy, shuffleLo);
+
+ auto newAccB =
+ mlir::vector::ShapeCastOp::create(rewriter, loc, accTy, shuffleHi);
+
+ rewriteUses(opA->getResult(0), newAccA.getResult(), contractA, rewriter);
+
+ rewriteUses(opB->getResult(0), newAccB.getResult(), contractB, rewriter);
+}
+
+void shuffleBeforeWriteLikeOp(mlir::PatternRewriter &rewriter,
+ mlir::Operation *opA, mlir::Operation *opB,
+ int64_t nonUnitDimAcc, mlir::VectorType accTy) {
+ // Helper to extract vector operand from write-like ops
+ auto getWrittenVector = [](mlir::Operation *op) -> mlir::Value {
+ if (auto write = mlir::dyn_cast<mlir::vector::TransferWriteOp>(op))
+ return write.getVector();
+ if (auto store = mlir::dyn_cast<mlir::vector::StoreOp>(op))
+ return store.getValueToStore();
+ return nullptr;
+ };
+
+ mlir::Value vecA = getWrittenVector(opA);
+ mlir::Value vecB = getWrittenVector(opB);
+
+ assert(vecA && vecB && "expected vector write-like ops");
+
+ // Decide insertion point and location
+ mlir::Operation *insertBefore = opA->isBeforeInBlock(opB) ? opA : opB;
+
+ rewriter.setInsertionPoint(insertBefore);
+ mlir::Location loc = insertBefore->getLoc();
+
+ auto elemTy = accTy.getElementType();
+ auto flatTy = mlir::VectorType::get(nonUnitDimAcc, elemTy);
+
+ // Flatten vectors
+ auto castA = mlir::vector::ShapeCastOp::create(rewriter, loc, flatTy, vecA);
+
+ auto castB = mlir::vector::ShapeCastOp::create(rewriter, loc, flatTy, vecB);
+
+ // TODO: derive shuffle masks instead of hard-coding
+ auto masks = getShuffleMasks(nonUnitDimAcc);
+
+ auto shuffledLo = mlir::vector::ShuffleOp::create(rewriter, loc, flatTy,
+ castA, castB, masks.maskLo);
+
+ auto shuffledHi = mlir::vector::ShuffleOp::create(rewriter, loc, flatTy,
+ castA, castB, masks.maskHi);
+
+ // Cast back to accumulator type
+ auto newVecA =
+ mlir::vector::ShapeCastOp::create(rewriter, loc, accTy, shuffledLo);
+
+ auto newVecB =
+ mlir::vector::ShapeCastOp::create(rewriter, loc, accTy, shuffledHi);
+
+ // Update write operands in place
+ opA->setOperand(0, newVecA.getResult());
+ opB->setOperand(0, newVecB.getResult());
+}
+
+void shuffleNonUnitDimOperand(mlir::PatternRewriter &rewriter,
+ mlir::Operation *opA, mlir::Operation *opB,
+ mlir::vector::ContractionOp contractA,
+ mlir::vector::ContractionOp contractB,
+ int64_t nonUnitDimAcc, mlir::VectorType Ty) {
+ mlir::Operation *insertAfter = opA->isBeforeInBlock(opB) ? opB : opA;
+
+ rewriter.setInsertionPointAfter(insertAfter);
+ mlir::Location loc = insertAfter->getLoc();
+
+ auto elemTy = Ty.getElementType();
+ auto flatTy = mlir::VectorType::get(nonUnitDimAcc, elemTy);
+
+ auto castA = mlir::vector::ShapeCastOp::create(rewriter, loc, flatTy,
+ opA->getResult(0));
+
+ auto castB = mlir::vector::ShapeCastOp::create(rewriter, loc, flatTy,
+ opB->getResult(0));
+
+ static constexpr int64_t maskLo[] = {
+ 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};
+ static constexpr int64_t maskHi[] = {
+ 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};
+
+ auto shuffleLo = mlir::vector::ShuffleOp::create(rewriter, loc, flatTy, castA,
+ castB, maskLo);
+
+ auto shuffleHi = mlir::vector::ShuffleOp::create(rewriter, loc, flatTy, castA,
+ castB, maskHi);
+
+ auto newAccA =
+ mlir::vector::ShapeCastOp::create(rewriter, loc, Ty, shuffleLo);
+
+ auto newAccB =
+ mlir::vector::ShapeCastOp::create(rewriter, loc, Ty, shuffleHi);
+
+ rewriteUses(opA->getResult(0), newAccA.getResult(), contractA, rewriter);
+
+ rewriteUses(opB->getResult(0), newAccB.getResult(), contractB, rewriter);
+}
+
+bool validatePairVectorContract(vector::ContractionOp contractOp,
+ vector::ContractionOp pairContOp,
+ bool rhsHasMultipleNonUnitDims,
+ int64_t nonUnitDimValue) {
+
+ if (!(contractOp.getLhs() == pairContOp.getLhs()) &&
+ !(contractOp.getRhs() == pairContOp.getRhs()))
+ return false;
+
+ if (rhsHasMultipleNonUnitDims &&
+ !(contractOp.getLhs() == pairContOp.getLhs()))
+ return false;
+
+ if (!rhsHasMultipleNonUnitDims &&
+ !(contractOp.getRhs() == pairContOp.getRhs()))
+ return false;
+
+ auto op =
+ rhsHasMultipleNonUnitDims ? contractOp.getRhs() : contractOp.getLhs();
+ auto op1 =
+ rhsHasMultipleNonUnitDims ? pairContOp.getRhs() : pairContOp.getLhs();
+
+ Value srcBuff;
+ SmallVector<OpFoldResult> indexVals;
+ llvm::TypeSwitch<mlir::Operation *>(op.getDefiningOp())
+ .Case<vector::TransferReadOp, vector::LoadOp>([&](auto readOp) {
+ srcBuff = readOp.getOperand(0);
+ indexVals = SmallVector<OpFoldResult>(readOp.getIndices().begin(),
+ readOp.getIndices().end());
+ });
+
+ Value srcBuff1;
+ SmallVector<OpFoldResult> indexVals1;
+ llvm::TypeSwitch<mlir::Operation *>(op1.getDefiningOp())
+ .Case<vector::TransferReadOp, vector::LoadOp>([&](auto readOp) {
+ srcBuff1 = readOp.getOperand(0);
+ indexVals1 = SmallVector<OpFoldResult>(readOp.getIndices().begin(),
+ readOp.getIndices().end());
+ });
+
+ if (!srcBuff || !srcBuff1)
+ return false;
+
+ if (!(srcBuff == srcBuff1))
+ return false;
+
+ for (size_t i = 0; i < indexVals.size(); i++) {
+ if (getConstantIntValue(indexVals[i]) == getConstantIntValue(indexVals1[i]))
+ continue;
+
+ auto value1 = *getConstantIntValue(indexVals1[i]);
+ auto value2 = *getConstantIntValue(indexVals[i]);
+
+ if ((value1 - value2) != nonUnitDimValue)
+ return false;
+ }
+
+ return true;
+}
+
} // namespace x86vector
} // namespace mlir
diff --git a/mlir/test/Dialect/shuffle-bf16-vector-contract-result.mlir b/mlir/test/Dialect/shuffle-bf16-vector-contract-result.mlir
new file mode 100644
index 0000000000000..7b1a6eaa1258a
--- /dev/null
+++ b/mlir/test/Dialect/shuffle-bf16-vector-contract-result.mlir
@@ -0,0 +1,66 @@
+// RUN: mlir-opt %s -transform-interpreter -cse -split-input-file | FileCheck %s
+
+
+!vecA = vector<1x1x1xbf16>
+!vecB = vector<1x1x8xbf16>
+!vecC = vector<1x8xf32>
+!memrefA = memref<1x4x1xbf16>
+!memrefB = memref<1x1x32xbf16>
+!memrefC = memref<2x32xf32>
+#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 @shuffle_VC_output_flat_layout(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %0 = ub.poison : bf16
+ %1 = vector.load %arg0[%c0, %c0, %c0] :
+ !memrefA, !vecA
+ %2 = vector.load %arg1[%c0, %c0, %c0] :
+ !memrefB, !vecB
+ %3 = vector.load %arg1[%c0, %c0, %c8] :
+ !memrefB, !vecB
+ %4 = vector.load %arg2[%c0, %c0] :
+ !memrefC, !vecC
+ %5 = vector.load %arg2[%c0, %c8] :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.store %6, %arg2[%c0, %c0] : !memrefC, !vecC
+
+ vector.store %7, %arg2[%c0, %c8] : !memrefC, !vecC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @shuffle_VC_output_flat_layout
+// CHECK: vector.shuffle
+// CHECK-NEXT: vector.shuffle
+// CHECK: vector.contract
+// CHECK: vector.shuffle
+// CHECK-NEXT: vector.shuffle
+
+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.x86vector.shuffle_bf16_vector_contract_result
+ } : !transform.any_op
+ transform.yield
+ }
+}
>From 7fcfa7772f5e15ebb12653a10ba0c9d061d4bd84 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 6 Jan 2026 06:37:29 -0800
Subject: [PATCH 2/4] fix header format issue
---
.../X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp b/mlir/lib/Dialect/X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp
index 2081de9f2be91..faa1b0cd48131 100644
--- a/mlir/lib/Dialect/X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp
+++ b/mlir/lib/Dialect/X86Vector/Transforms/ShuffleBF16VectorContractResult.cpp
@@ -1,5 +1,4 @@
-//===-
-//ShuffleBF16VectorContractResult.cpp-----------------------------------------===//
+//===- ShuffleBF16VectorContractResult.cpp --------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
>From b73ef356f51547a07efb7b3149062e01d3f675a2 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 6 Jan 2026 20:33:32 -0800
Subject: [PATCH 3/4] added couple of unit tests
---
.../shuffle-bf16-vector-contract-result.mlir | 143 +++++++++++++++++-
1 file changed, 138 insertions(+), 5 deletions(-)
diff --git a/mlir/test/Dialect/shuffle-bf16-vector-contract-result.mlir b/mlir/test/Dialect/shuffle-bf16-vector-contract-result.mlir
index 7b1a6eaa1258a..c109543a8a163 100644
--- a/mlir/test/Dialect/shuffle-bf16-vector-contract-result.mlir
+++ b/mlir/test/Dialect/shuffle-bf16-vector-contract-result.mlir
@@ -1,6 +1,5 @@
// RUN: mlir-opt %s -transform-interpreter -cse -split-input-file | FileCheck %s
-
!vecA = vector<1x1x1xbf16>
!vecB = vector<1x1x8xbf16>
!vecC = vector<1x8xf32>
@@ -49,11 +48,145 @@ func.func @shuffle_VC_output_flat_layout(
}
// CHECK-LABEL: @shuffle_VC_output_flat_layout
-// CHECK: vector.shuffle
-// CHECK-NEXT: vector.shuffle
+// 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: vector.contract
+// 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>
+
+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.x86vector.shuffle_bf16_vector_contract_result
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<1x1xbf16>
+!vecB = vector<1x8xbf16>
+!vecC = vector<1x8xf32>
+!memrefA = memref<4x1xbf16>
+!memrefB = memref<1x32xbf16>
+!memrefC = memref<2x32xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @shuffle_VC_output_flat_layout_transfer_read(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+ %1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg1[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+ %3 = vector.transfer_read %arg1[%c0, %c8], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+ %4 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %5 = vector.transfer_read %arg2[%c0, %c8], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.transfer_write %6, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ vector.transfer_write %7, %arg2[%c0, %c8] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @shuffle_VC_output_flat_layout_transfer_read
+// 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: vector.contract
+// 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>
+
+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.x86vector.shuffle_bf16_vector_contract_result
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<1x2xbf16>
+!vecB = vector<2x16xbf16>
+!vecC = vector<1x16xf32>
+!memrefA = memref<4x2xbf16>
+!memrefB = memref<2x32xbf16>
+!memrefC = memref<2x32xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @shuffle_VC_output_flat_layout_bf16dp(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+ %1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg1[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+ %3 = vector.transfer_read %arg1[%c0, %c16], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+ %4 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %5 = vector.transfer_read %arg2[%c0, %c16], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.transfer_write %6, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ vector.transfer_write %7, %arg2[%c0, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @shuffle_VC_output_flat_layout_bf16dp
+// 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: vector.contract
-// CHECK: vector.shuffle
-// CHECK-NEXT: vector.shuffle
+// 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>
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
>From cc40fb30a3839b89cd55c56e11cff509d465df28 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Thu, 8 Jan 2026 07:25:29 -0800
Subject: [PATCH 4/4] added support for BF16 flat layout to vector.fma using
packed operations
---
.../Transforms/VectorContractBF16ToFMA.cpp | 173 ++++--
.../X86Vector/Utils/X86VectorUtils.cpp | 46 +-
.../vector-contract-bf16-to-fma.mlir | 495 ++++++++++++++++++
3 files changed, 662 insertions(+), 52 deletions(-)
diff --git a/mlir/lib/Dialect/X86Vector/Transforms/VectorContractBF16ToFMA.cpp b/mlir/lib/Dialect/X86Vector/Transforms/VectorContractBF16ToFMA.cpp
index c60d9b91c18e5..eada03977595d 100644
--- a/mlir/lib/Dialect/X86Vector/Transforms/VectorContractBF16ToFMA.cpp
+++ b/mlir/lib/Dialect/X86Vector/Transforms/VectorContractBF16ToFMA.cpp
@@ -8,6 +8,7 @@
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/Dialect/X86Vector/Transforms.h"
@@ -29,7 +30,7 @@ using namespace mlir::x86vector;
// Verifies that the LHS and RHS operands of a vector.contract are load or
// vector.transfer_read operations on a memref source buffer, and checks
// their bounds, dimensions, offsets, and strides.
-static bool validateVectorContractOperands(Value prodOp) {
+static bool validateVectorContractOperands(Value prodOp, bool isVnni) {
Operation *defOp = prodOp.getDefiningOp();
if (!defOp)
return false;
@@ -62,11 +63,13 @@ static bool validateVectorContractOperands(Value prodOp) {
// Return false if the two innermost strides of the memref are not contiguous.
// The x86vector.avx.cvt.packed.even/odd.indexed_to_f32 operations require
// an eight-element tuple of bf16 values to be contiguous.
- if (!llvm::cast<mlir::MemRefType>(srcType).areTrailingDimsContiguous(2))
+ int dimsToCheck = isVnni ? 2 : 1;
+ if (!llvm::cast<mlir::MemRefType>(srcType).areTrailingDimsContiguous(
+ dimsToCheck))
return false;
// Return false if the vnni offset of load or transfer_read is not zero.
- if (getConstantIntValue(indexVals.back()) != 0)
+ if (isVnni && getConstantIntValue(indexVals.back()) != 0)
return false;
return true;
@@ -96,7 +99,8 @@ static bool validateVectorContractOperands(Value prodOp) {
// ```
static SmallVector<memref::SubViewOp>
getSubviewFromVectorInput(Location loc, PatternRewriter &rewriter, Value prodOp,
- ArrayRef<int64_t> nonUnitDimShape, bool isUnitDim) {
+ ArrayRef<int64_t> nonUnitDimShape, bool isUnitDim,
+ bool isVNNI) {
Operation *defOp = prodOp.getDefiningOp();
@@ -122,11 +126,26 @@ getSubviewFromVectorInput(Location loc, PatternRewriter &rewriter, Value prodOp,
}
}
- int vnniDimSize = isUnitDim ? 1 : 2;
+ auto one = rewriter.getIndexAttr(1);
+ llvm::SmallVector<memref::SubViewOp> subviews;
+ if (!isVNNI) {
+ SmallVector<OpFoldResult> strides(indexVals.size(), one);
+ SmallVector<OpFoldResult> sizes(indexVals.size(), one);
+ // Retrive twice the nonUnit dim BF16 element for both even and odd
+ // index elements.
+ if (!isUnitDim)
+ mnDimSize = 2 * mnDimSize;
+ sizes[mnDimIdx] = rewriter.getIndexAttr(mnDimSize);
+ auto subview = memref::SubViewOp::create(rewriter, loc, srcBuff, indexVals,
+ sizes, strides);
+ subviews.push_back(subview);
+ return subviews;
+ }
+
+ int vnniDimSize = isUnitDim ? 1 : 2;
auto nonVNNIDimSize = indexVals.size() - 1;
// Create the size and stride offsets.
- auto one = rewriter.getIndexAttr(1);
SmallVector<OpFoldResult> strides(indexVals.size(), one);
SmallVector<OpFoldResult> sizes(nonVNNIDimSize, one);
@@ -139,7 +158,6 @@ getSubviewFromVectorInput(Location loc, PatternRewriter &rewriter, Value prodOp,
if (isUnitDim)
indexVals[indexVals.size() - 1] = rewriter.getIndexAttr(1);
- llvm::SmallVector<memref::SubViewOp> subviews;
auto subview = memref::SubViewOp::create(rewriter, loc, srcBuff, indexVals,
sizes, strides);
subviews.push_back(subview);
@@ -168,7 +186,7 @@ getSubviewFromVectorInput(Location loc, PatternRewriter &rewriter, Value prodOp,
// Implements outer product contraction as a sequence of BF16-packed
// operation even/odd loads and FMA operations.
//
-// For example:
+// For example (VNNI packed):
// ```
// %1 = vector.load from memref (%m1) -> vector<1x1x2xbf16>
// %2 = vector.load from memref (%m2) -> vector<1x8x2xbf16>
@@ -183,6 +201,24 @@ getSubviewFromVectorInput(Location loc, PatternRewriter &rewriter, Value prodOp,
// %5 = x86vector.avx.cvt.packed.even.indexed_to_f32 %m2 -> vector<8xf32>
// return vector.fma %4, %5, %3
// ```
+//
+// For example (Flat layout):
+// ```
+// %1 = vector.load from memref (%m1) -> vector<1x1xbf16>
+// %2 = vector.load from memref (%m2) -> vector<1x8xbf16>
+// %3 = vector.contract %1, %2, %arg1
+// %4 = vector.load from memref (%m2) -> vector<1x8xbf16>
+// %5 = vector.contract %1, %4, %arg2
+// scf.yield %3, %4
+// ```
+// to
+// ```
+// %1 = x86vector.avx.bcst_to_f32.packed %m1[c0] -> vector<8xf32>
+// %2 = x86vector.avx.cvt.packed.even.indexed_to_f32 %m2 -> vector<8xf32>
+// %3 = vector.fma %1, %2, %arg1
+// %4 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %m2 -> vector<8xf32>
+// %5 = vector.fma %1, %4, %arg2
+// scf.yield %3, %5
struct VectorContractBF16ToFMA
: public OpRewritePattern<vector::ContractionOp> {
using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
@@ -202,11 +238,9 @@ struct VectorContractBF16ToFMA
return rewriter.notifyMatchFailure(contractOp,
"Only BF16 lowering is supported.");
- if (!isInVnniLayout(contractOp.getOperation(),
- contractOp.getIndexingMapsArray(),
- /*blockingFactor=*/2))
- return rewriter.notifyMatchFailure(contractOp,
- "Input matrices not in VNNI format.");
+ bool isVnni = isInVnniLayout(contractOp.getOperation(),
+ contractOp.getIndexingMapsArray(),
+ /*blockingFactor=*/2);
VectorType accTy = dyn_cast<VectorType>(contractOp.getAccType());
if (!accTy)
@@ -216,6 +250,14 @@ struct VectorContractBF16ToFMA
return rewriter.notifyMatchFailure(
contractOp, "Only F32 acumulation supported for BF16 type.");
+ ArrayRef<int64_t> accShape = accTy.getShape();
+ llvm::SmallVector<int64_t> nonUnitDimAcc;
+ llvm::copy_if(accShape, std::back_inserter(nonUnitDimAcc),
+ [](int64_t dim) { return dim != 1; });
+ if (nonUnitDimAcc.size() != 1)
+ return rewriter.notifyMatchFailure(
+ contractOp, "A or B should be a non-unit dim in acc.");
+
ArrayRef<int64_t> lhsShape = lhsTy.getShape();
llvm::SmallVector<int64_t> nonUnitDimLhs;
llvm::copy_if(lhsShape, std::back_inserter(nonUnitDimLhs),
@@ -227,35 +269,38 @@ struct VectorContractBF16ToFMA
llvm::copy_if(rhsShape, std::back_inserter(nonUnitDimRhs),
[](int64_t dim) { return dim != 1; });
- if ((nonUnitDimLhs.size() - 1) > 0 && (nonUnitDimRhs.size() - 1) > 0)
+ if (isVnni && (nonUnitDimLhs.size() - 1) > 0 &&
+ (nonUnitDimRhs.size() - 1) > 0)
return rewriter.notifyMatchFailure(contractOp,
"Excepts unit dimensions for either "
"LHS or RHS shape other than VNNI.");
- if ((nonUnitDimLhs.size() - 1) != 1 && (nonUnitDimRhs.size() - 1) != 1)
+ if (isVnni && (nonUnitDimLhs.size() - 1) != 1 &&
+ (nonUnitDimRhs.size() - 1) != 1)
return rewriter.notifyMatchFailure(
contractOp,
"Excepts a one non-unit A/B dimension for either LHS or RHS shape.");
- ArrayRef<int64_t> accShape = accTy.getShape();
- llvm::SmallVector<int64_t> nonUnitDimAcc;
- llvm::copy_if(accShape, std::back_inserter(nonUnitDimAcc),
- [](int64_t dim) { return dim != 1; });
- if (nonUnitDimAcc.size() != 1)
+ if (!isVnni && nonUnitDimLhs.size() > 0 && nonUnitDimRhs.size() > 0)
+ return rewriter.notifyMatchFailure(contractOp,
+ "Excepts unit dimensions for either "
+ "LHS or RHS shape.");
+
+ if (!isVnni && nonUnitDimLhs.size() != 1 && nonUnitDimRhs.size() != 1)
return rewriter.notifyMatchFailure(
- contractOp, "A or B should be a non-unit dim in acc.");
+ contractOp,
+ "Excepts a one non-unit A/B dimension for either LHS or RHS shape.");
// Non-unit dimensions should match the vector length of BF16.
- unsigned int nonUnitDim = nonUnitDimLhs.size() == 2 ? nonUnitDimLhs.front()
- : nonUnitDimRhs.front();
- if (nonUnitDim != 4 && nonUnitDim != 8 &&
- !(nonUnitDimAcc.front() == nonUnitDim))
+ unsigned int nonUnitDim = nonUnitDimAcc.front();
+
+ if (nonUnitDim != 4 && nonUnitDim != 8)
return rewriter.notifyMatchFailure(
contractOp, "BF16 packed load operation expects non-unit (LHR or "
"RHS) dim and acc dim of size 4/8.");
- if (!validateVectorContractOperands(contractOp.getLhs()) ||
- !validateVectorContractOperands(contractOp.getRhs())) {
+ if (!validateVectorContractOperands(contractOp.getLhs(), isVnni) ||
+ !validateVectorContractOperands(contractOp.getRhs(), isVnni)) {
return rewriter.notifyMatchFailure(
contractOp, "The LHS or RHS is in an invalid format. Either it has "
"false in-bounds, "
@@ -273,7 +318,12 @@ struct VectorContractBF16ToFMA
// vector<1x8x2xbf16>, we create two subview for the LHS and one subview
// for the RHS. In the opposite case (non-unit dimension on the LHS), we
// do vice-versa.
- bool rhsHasMultipleNonUnitDims = (nonUnitDimRhs.size() - 1) > 0;
+
+ bool rhsHasMultipleNonUnitDims = nonUnitDimRhs.size() > 0;
+ if (isVnni) {
+ rhsHasMultipleNonUnitDims = (nonUnitDimRhs.size() - 1) > 0;
+ }
+
// Select which operand is "unit" and which is "non-unit".
Value unitSrc =
rhsHasMultipleNonUnitDims ? contractOp.getLhs() : contractOp.getRhs();
@@ -283,12 +333,38 @@ struct VectorContractBF16ToFMA
ArrayRef<int64_t> nonUnitDimShape =
rhsHasMultipleNonUnitDims ? rhsShape : lhsShape;
+ // Get the pair vector.contract operation. The pair is decided on:
+ // (1) - the unitDim operand Lhs or Rhs should be same,
+ // (2) - the defining source memref should be same for nonUnitDim
+ // operation, (3) - the nonUnit dim offset difference between the
+ // vector.contracts should be 8.
+ vector::ContractionOp pairContractOp;
+ if (!isVnni) {
+ Operation *nextOp = contractOp;
+ while ((nextOp = nextOp->getNextNode())) {
+ auto contOp = dyn_cast<vector::ContractionOp>(nextOp);
+
+ if (!contOp)
+ continue;
+
+ if (validatePairVectorContract(contractOp, contOp,
+ rhsHasMultipleNonUnitDims,
+ nonUnitDimAcc.front())) {
+ pairContractOp = contOp;
+ break;
+ }
+ }
+
+ if (!pairContractOp)
+ return failure();
+ }
+
// Build subviews.
- auto unitDimSubview = getSubviewFromVectorInput(loc, rewriter, unitSrc,
- nonUnitDimShape, true);
+ auto unitDimSubview = getSubviewFromVectorInput(
+ loc, rewriter, unitSrc, nonUnitDimShape, true, isVnni);
auto nonUnitDimSubview = getSubviewFromVectorInput(
- loc, rewriter, nonUnitSrc, nonUnitDimShape, false);
+ loc, rewriter, nonUnitSrc, nonUnitDimShape, false, isVnni);
auto castAcc = vector::ShapeCastOp::create(
rewriter, loc,
@@ -297,6 +373,41 @@ struct VectorContractBF16ToFMA
VectorType dstType =
VectorType::get(nonUnitDimAcc.front(), rewriter.getF32Type());
+ if (!isVnni) {
+ auto loadBcstBF16ElementToF32 = x86vector::BcstToPackedF32Op::create(
+ rewriter, loc, dstType, unitDimSubview[0]);
+ auto loadEvenIdxElementF32 =
+ x86vector::CvtPackedEvenIndexedToF32Op::create(rewriter, loc, dstType,
+ nonUnitDimSubview[0]);
+ auto evenIdxFMA =
+ vector::FMAOp::create(rewriter, loc, loadBcstBF16ElementToF32,
+ loadEvenIdxElementF32, castAcc);
+ auto castEvenFma =
+ vector::ShapeCastOp::create(rewriter, loc, accTy, evenIdxFMA);
+ rewriter.replaceOp(contractOp, castEvenFma);
+
+ rewriter.setInsertionPoint(pairContractOp);
+ auto pairContOpLoc = pairContractOp.getLoc();
+ VectorType accTyPairCont =
+ dyn_cast<VectorType>(pairContractOp.getAccType());
+ auto castAccPairCont = vector::ShapeCastOp::create(
+ rewriter, pairContOpLoc,
+ VectorType::get(nonUnitDimAcc.front(),
+ accTyPairCont.getElementType()),
+ pairContractOp.getAcc());
+
+ auto loadOddIdxElementF32 = x86vector::CvtPackedOddIndexedToF32Op::create(
+ rewriter, pairContOpLoc, dstType, nonUnitDimSubview[0]);
+ auto oddIdxFMA = vector::FMAOp::create(
+ rewriter, pairContOpLoc, loadBcstBF16ElementToF32,
+ loadOddIdxElementF32, castAccPairCont);
+ auto castOddFma = vector::ShapeCastOp::create(rewriter, pairContOpLoc,
+ accTyPairCont, oddIdxFMA);
+ rewriter.replaceOp(pairContractOp, castOddFma);
+
+ return success();
+ }
+
// Load, broadcast, and do FMA for odd indexed BF16 elements.
auto loadBcstOddIdxElementToF32 = x86vector::BcstToPackedF32Op::create(
rewriter, loc, dstType, unitDimSubview[0]);
diff --git a/mlir/lib/Dialect/X86Vector/Utils/X86VectorUtils.cpp b/mlir/lib/Dialect/X86Vector/Utils/X86VectorUtils.cpp
index 65cd0773160ef..38edb8091ce72 100644
--- a/mlir/lib/Dialect/X86Vector/Utils/X86VectorUtils.cpp
+++ b/mlir/lib/Dialect/X86Vector/Utils/X86VectorUtils.cpp
@@ -369,15 +369,16 @@ void shuffleNonUnitDimOperand(mlir::PatternRewriter &rewriter,
rewriteUses(opB->getResult(0), newAccB.getResult(), contractB, rewriter);
}
+// Return true if vector.contract operations matches on below conditions:
+// (1) - the unitDim operand Lhs or Rhs should be same,
+// (2) - the defining source memref should be same for nonUnitDim
+// operation,
+// (3) - the nonUnit dim offset difference between the
+// vector.contracts should be 8.
bool validatePairVectorContract(vector::ContractionOp contractOp,
vector::ContractionOp pairContOp,
bool rhsHasMultipleNonUnitDims,
int64_t nonUnitDimValue) {
-
- if (!(contractOp.getLhs() == pairContOp.getLhs()) &&
- !(contractOp.getRhs() == pairContOp.getRhs()))
- return false;
-
if (rhsHasMultipleNonUnitDims &&
!(contractOp.getLhs() == pairContOp.getLhs()))
return false;
@@ -386,43 +387,46 @@ bool validatePairVectorContract(vector::ContractionOp contractOp,
!(contractOp.getRhs() == pairContOp.getRhs()))
return false;
- auto op =
+ auto nonUnitOperand =
rhsHasMultipleNonUnitDims ? contractOp.getRhs() : contractOp.getLhs();
- auto op1 =
+ auto nonUnitOperandPairContOp =
rhsHasMultipleNonUnitDims ? pairContOp.getRhs() : pairContOp.getLhs();
Value srcBuff;
SmallVector<OpFoldResult> indexVals;
- llvm::TypeSwitch<mlir::Operation *>(op.getDefiningOp())
+ llvm::TypeSwitch<mlir::Operation *>(nonUnitOperand.getDefiningOp())
.Case<vector::TransferReadOp, vector::LoadOp>([&](auto readOp) {
srcBuff = readOp.getOperand(0);
indexVals = SmallVector<OpFoldResult>(readOp.getIndices().begin(),
readOp.getIndices().end());
});
- Value srcBuff1;
- SmallVector<OpFoldResult> indexVals1;
- llvm::TypeSwitch<mlir::Operation *>(op1.getDefiningOp())
+ Value srcBuffPairContOp;
+ SmallVector<OpFoldResult> indexValsPairContOp;
+ llvm::TypeSwitch<mlir::Operation *>(nonUnitOperandPairContOp.getDefiningOp())
.Case<vector::TransferReadOp, vector::LoadOp>([&](auto readOp) {
- srcBuff1 = readOp.getOperand(0);
- indexVals1 = SmallVector<OpFoldResult>(readOp.getIndices().begin(),
- readOp.getIndices().end());
+ srcBuffPairContOp = readOp.getOperand(0);
+ indexValsPairContOp = SmallVector<OpFoldResult>(
+ readOp.getIndices().begin(), readOp.getIndices().end());
});
- if (!srcBuff || !srcBuff1)
+ if (!srcBuff || !srcBuffPairContOp)
return false;
- if (!(srcBuff == srcBuff1))
+ if (!(srcBuff == srcBuffPairContOp))
return false;
for (size_t i = 0; i < indexVals.size(); i++) {
- if (getConstantIntValue(indexVals[i]) == getConstantIntValue(indexVals1[i]))
- continue;
+ auto v0 = getConstantIntValue(indexVals[i]);
+ auto v1 = getConstantIntValue(indexValsPairContOp[i]);
+
+ if (!v0 || !v1)
+ return false;
- auto value1 = *getConstantIntValue(indexVals1[i]);
- auto value2 = *getConstantIntValue(indexVals[i]);
+ if (*v1 == *v0)
+ continue;
- if ((value1 - value2) != nonUnitDimValue)
+ if ((*v1 - *v0) != nonUnitDimValue)
return false;
}
diff --git a/mlir/test/Dialect/X86Vector/vector-contract-bf16-to-fma.mlir b/mlir/test/Dialect/X86Vector/vector-contract-bf16-to-fma.mlir
index e7a70429490d1..05404fa1ef6a5 100644
--- a/mlir/test/Dialect/X86Vector/vector-contract-bf16-to-fma.mlir
+++ b/mlir/test/Dialect/X86Vector/vector-contract-bf16-to-fma.mlir
@@ -233,6 +233,228 @@ module attributes {transform.with_named_sequence} {
// -----
+!vecA = vector<1x1xbf16>
+!vecB = vector<1x8xbf16>
+!vecC = vector<1x8xf32>
+!memrefA = memref<4x1xbf16>
+!memrefB = memref<1x32xbf16>
+!memrefC = memref<2x32xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @matmul_to_fma_flat_layout(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+ %1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg1[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+ %3 = vector.transfer_read %arg1[%c0, %c8], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+ %4 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %5 = vector.transfer_read %arg2[%c0, %c8], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.transfer_write %6, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+ vector.transfer_write %7, %arg2[%c0, %c8] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @matmul_to_fma_flat_layout
+// 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: memref.subview %arg0[%c0, %c0] {{.*}} : memref<4x1xbf16> to memref<1x1xbf16, {{.*}}>
+// CHECK: memref.subview %arg1[%c0, %c0] {{.*}} : memref<1x32xbf16> to memref<1x16xbf16, {{.*}}>
+// CHECK: x86vector.avx.bcst_to_f32.packed {{.*}} : memref<1x1xbf16, strided<[1, 1], offset: ?>>
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32 {{.*}} : memref<1x16xbf16, strided<[32, 1], offset: ?>>
+// CHECK: vector.fma {{.*}} : vector<8xf32>
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32 {{.*}} : memref<1x16xbf16, strided<[32, 1], offset: ?>>
+// CHECK: vector.fma {{.*}} : vector<8xf32>
+// 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>
+
+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.x86vector.shuffle_bf16_vector_contract_result
+ transform.apply_patterns.x86vector.vector_contract_bf16_to_fma
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<1x1xbf16>
+!vecB = vector<1x8xbf16>
+!vecC = vector<1x8xf32>
+!memrefA = memref<4x1xbf16>
+!memrefB = memref<1x32xbf16>
+!memrefC = memref<2x32xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @matmul_to_fma_flat_layout_load(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+ %1 = vector.load %arg0[%c0, %c0] :
+ !memrefA, !vecA
+ %2 = vector.load %arg1[%c0, %c0] :
+ !memrefB, !vecB
+ %3 = vector.load %arg1[%c0, %c8] :
+ !memrefB, !vecB
+ %4 = vector.load %arg2[%c0, %c0] :
+ !memrefC, !vecC
+ %5 = vector.load %arg2[%c0, %c8] :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.store %6, %arg2[%c0, %c0] : !memrefC, !vecC
+ vector.store %7, %arg2[%c0, %c8] : !memrefC, !vecC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @matmul_to_fma_flat_layout_load
+// 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: memref.subview %arg0[%c0, %c0] {{.*}} : memref<4x1xbf16> to memref<1x1xbf16, {{.*}}>
+// CHECK: memref.subview %arg1[%c0, %c0] {{.*}} : memref<1x32xbf16> to memref<1x16xbf16, {{.*}}>
+// CHECK: x86vector.avx.bcst_to_f32.packed {{.*}} : memref<1x1xbf16, strided<[1, 1], offset: ?>>
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32 {{.*}} : memref<1x16xbf16, strided<[32, 1], offset: ?>>
+// CHECK: vector.fma {{.*}} : vector<8xf32>
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32 {{.*}} : memref<1x16xbf16, strided<[32, 1], offset: ?>>
+// CHECK: vector.fma {{.*}} : vector<8xf32>
+// 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>
+
+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.x86vector.shuffle_bf16_vector_contract_result
+ transform.apply_patterns.x86vector.vector_contract_bf16_to_fma
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<8x1xbf16>
+!vecB = vector<1x1xbf16>
+!vecC = vector<8x1xf32>
+!memrefA = memref<32x1xbf16>
+!memrefB = memref<1x4xbf16>
+!memrefC = memref<32x4xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @matmul_to_fma_flat_layout_bcstB(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+
+ %1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg0[%c8, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+
+ %3 = vector.transfer_read %arg1[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+
+ %4 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %5 = vector.transfer_read %arg2[%c8, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %2, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.transfer_write %6, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+ vector.transfer_write %7, %arg2[%c8, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @matmul_to_fma_flat_layout_bcstB
+// 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: memref.subview %arg1[%c0, %c0] {{.*}} : memref<1x4xbf16> to memref<1x1xbf16, {{.*}}>
+// CHECK: memref.subview %arg0[%c0, %c0] {{.*}} : memref<32x1xbf16> to memref<16x1xbf16, {{.*}}>
+// CHECK: x86vector.avx.bcst_to_f32.packed {{.*}} : memref<1x1xbf16, strided<[4, 1], offset: ?>>
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32 {{.*}} : memref<16x1xbf16, strided<[1, 1], offset: ?>>
+// CHECK: vector.fma {{.*}} : vector<8xf32>
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32 {{.*}} : memref<16x1xbf16, strided<[1, 1], offset: ?>>
+// CHECK: vector.fma {{.*}} : vector<8xf32>
+// 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>
+
+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.x86vector.shuffle_bf16_vector_contract_result
+ transform.apply_patterns.x86vector.vector_contract_bf16_to_fma
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
!vecA = vector<1x1x2xbf16>
!vecB = vector<1x8x2xbf16>
!vecC = vector<1x8xf32>
@@ -413,6 +635,279 @@ module attributes {transform.with_named_sequence} {
// -----
+!vecA = vector<1x1xbf16>
+!vecB = vector<1x8xbf16>
+!vecC = vector<1x8xf32>
+!memrefA = memref<4x1xbf16>
+!memrefB = memref<1x32xbf16>
+!memrefC = memref<2x32xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @negative_offset_diff_is_not_8(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+ %1 = vector.load %arg0[%c0, %c0] :
+ !memrefA, !vecA
+ %2 = vector.load %arg1[%c0, %c0] :
+ !memrefB, !vecB
+ %3 = vector.load %arg1[%c0, %c16] :
+ !memrefB, !vecB
+ %4 = vector.load %arg2[%c0, %c0] :
+ !memrefC, !vecC
+ %5 = vector.load %arg2[%c0, %c16] :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.store %6, %arg2[%c0, %c0] : !memrefC, !vecC
+ vector.store %7, %arg2[%c0, %c16] : !memrefC, !vecC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_offset_diff_is_not_8
+// CHECK-NOT: x86vector.avx.bcst_to_f32.packed
+// CHECK-NOT: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK-NOT: vector.fma {{.*}} : vector<8xf32>
+// CHECK-NOT: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// 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.x86vector.vector_contract_bf16_to_fma
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<1x1xbf16>
+!vecB = vector<1x8xbf16>
+!vecC = vector<1x8xf32>
+!memrefA = memref<4x1xbf16>
+!memrefB = memref<1x32xbf16>
+!memrefC = memref<2x32xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @negative_vector_contracts_not_in_order(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+ %1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg1[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+ %3 = vector.transfer_read %arg1[%c0, %c8], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+ %4 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %5 = vector.transfer_read %arg2[%c0, %c8], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %5
+ : !vecA, !vecB into !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %4
+ : !vecA, !vecB into !vecC
+
+ vector.transfer_write %6, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+ vector.transfer_write %7, %arg2[%c0, %c8] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_vector_contracts_not_in_order
+// CHECK-NOT: x86vector.avx.bcst_to_f32.packed
+// CHECK-NOT: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK-NOT: vector.fma {{.*}} : vector<8xf32>
+// CHECK-NOT: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// 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.x86vector.shuffle_bf16_vector_contract_result
+ transform.apply_patterns.x86vector.vector_contract_bf16_to_fma
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<8x1xbf16>
+!vecB = vector<1x1xbf16>
+!vecC = vector<8x1xf32>
+!memrefA = memref<32x1xbf16>
+!memrefB = memref<1x4xbf16>
+!memrefC = memref<32x4xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @negative_flat_layout_dynamic_index(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC, %arg3: index) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+
+ %1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg0[%arg3, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+
+ %3 = vector.transfer_read %arg1[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+
+ %4 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %5 = vector.transfer_read %arg2[%c8, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %2, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.transfer_write %6, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+ vector.transfer_write %7, %arg2[%c8, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_flat_layout_dynamic_index
+// CHECK-NOT: x86vector.avx.bcst_to_f32.packed
+// CHECK-NOT: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK-NOT: vector.fma {{.*}} : vector<8xf32>
+// CHECK-NOT: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// 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.x86vector.shuffle_bf16_vector_contract_result
+ transform.apply_patterns.x86vector.vector_contract_bf16_to_fma
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<8x2xbf16>
+!vecB = vector<2x1xbf16>
+!vecC = vector<8x1xf32>
+!memrefA = memref<32x2xbf16>
+!memrefB = memref<2x4xbf16>
+!memrefC = memref<32x4xf32>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @negative_non_unit_K_dim(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %0 = ub.poison : bf16
+ %32 = ub.poison : f32
+
+ %1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg0[%c8, %c0], %0 {in_bounds = [true, true]} :
+ !memrefA, !vecA
+
+ %3 = vector.transfer_read %arg1[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefB, !vecB
+
+ %4 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %5 = vector.transfer_read %arg2[%c8, %c0], %32 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %6 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %3, %4
+ : !vecA, !vecB into !vecC
+
+ %7 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %2, %3, %5
+ : !vecA, !vecB into !vecC
+
+ vector.transfer_write %6, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+ vector.transfer_write %7, %arg2[%c8, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_non_unit_K_dim
+// CHECK-NOT: x86vector.avx.bcst_to_f32.packed
+// CHECK-NOT: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK-NOT: vector.fma {{.*}} : vector<8xf32>
+// CHECK-NOT: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// 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.x86vector.shuffle_bf16_vector_contract_result
+ transform.apply_patterns.x86vector.vector_contract_bf16_to_fma
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
!vecA = vector<1x1x1x2xbf16>
!vecB = vector<1x1x16x2xbf16>
!vecC = vector<1x1x16xf32>
More information about the Mlir-commits
mailing list