[Mlir-commits] [llvm] [mlir] [mlir][X86] Fold VNNI operand unpacking into contractions (PR #232800)
Frank Schlimbach
llvmlistbot at llvm.org
Mon Oct 26 04:48:44 PDT 2026
https://github.com/fschlimb updated https://github.com/llvm/llvm-project/pull/232800
>From ea841c92f22a753e4240920f7c9563b16df53ce2 Mon Sep 17 00:00:00 2001
From: fschlimb <frank.schlimbach at intel.com>
Date: Tue, 13 Oct 2026 02:45:46 -0700
Subject: [PATCH 1/4] [mlir][X86] Fold VNNI operand unpacking into contractions
---
.../X86/TransformOps/X86TransformOps.td | 10 +
mlir/include/mlir/Dialect/X86/Transforms.h | 15 +-
.../Dialect/X86/TransformOps/CMakeLists.txt | 3 +
.../X86/TransformOps/X86TransformOps.cpp | 11 +
.../lib/Dialect/X86/Transforms/CMakeLists.txt | 3 +
.../Transforms/PrepareVectorContractLoops.cpp | 316 +++++++++++++++
.../X86/vector-contract-prepare-loops.mlir | 379 ++++++++++++++++++
.../CPU/X86/prepare-contract-loops.mlir | 87 ++++
.../llvm-project-overlay/mlir/BUILD.bazel | 9 +
9 files changed, 827 insertions(+), 6 deletions(-)
create mode 100644 mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
create mode 100644 mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
create mode 100644 mlir/test/Integration/Dialect/Vector/CPU/X86/prepare-contract-loops.mlir
diff --git a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
index a5afa9d3a6db1d..b5859f6d27adea 100644
--- a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
+++ b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
@@ -93,6 +93,16 @@ def ApplyMoveAccumulatorForContractLoopPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplyPrepareVectorContractLoopsPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.x86.prepare_vector_contract_loops",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Collect patterns that prepare vector contraction loops for x86 lowering.
+ }];
+
+ let assemblyFormat = "attr-dict";
+}
+
def ApplyVectorContractMultiLevelUnrollForAMXPatternsOp : Op<Transform_Dialect,
"apply_patterns.x86.vector_contract_multi_level_unroll_for_amx",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
diff --git a/mlir/include/mlir/Dialect/X86/Transforms.h b/mlir/include/mlir/Dialect/X86/Transforms.h
index a682d65014c756..c2444f53584880 100644
--- a/mlir/include/mlir/Dialect/X86/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86/Transforms.h
@@ -12,14 +12,12 @@
#include "mlir/IR/Value.h"
namespace mlir {
-
class ImplicitLocOpBuilder;
class LLVMConversionTarget;
class LLVMTypeConverter;
class RewritePatternSet;
namespace x86 {
-
/// Helper class to factor out the creation and extraction of masks from nibs.
struct MaskHelper {
/// b0 captures the lowest bit, b7 captures the highest bit.
@@ -115,6 +113,15 @@ void populateVectorContractToAMXDotProductPatterns(RewritePatternSet &patterns);
void populateMoveAccumulatorForContractLoopPatterns(
RewritePatternSet &patterns);
+// Clients must register the affine, arith and scf value-bounds models before
+// applying these patterns to loops with computed read indices.
+void populatePrepareVectorContractLoopsPatterns(RewritePatternSet &patterns);
+
+// Feature flags to control the multi-level unroll patttern.
+struct MultiLevelUnrollOptions {
+ bool enableAMX = false; // Prepare for AMX tile dot-product
+};
+
// A set of patterns that tile and unroll contractions in order to prepare them
// for the AMX nanokernel lowerings.
void populateVectorContractMultiLevelUnrollForAMXPatterns(
@@ -133,7 +140,6 @@ void populateVectorContractMultiLevelUnrollForAMXPatterns(
//===----------------------------------------------------------------------===//
namespace avx2 {
-
namespace inline_asm {
//===----------------------------------------------------------------------===//
/// Methods in the inline_asm namespace emit calls to LLVM::InlineAsmOp.
@@ -141,7 +147,6 @@ namespace inline_asm {
/// If bit i of `mask` is zero, take f32 at i from v1 else take it from v2.
Value mm256BlendPsAsm(ImplicitLocOpBuilder &b, Value v1, Value v2,
uint8_t mask);
-
} // namespace inline_asm
namespace intrin {
@@ -210,7 +215,6 @@ struct LoweringOptions {
void populateSpecializedTransposeLoweringPatterns(
RewritePatternSet &patterns, LoweringOptions options = LoweringOptions(),
int benefit = 10);
-
} // namespace avx2
} // namespace x86
@@ -225,7 +229,6 @@ void configureX86LegalizeForExportTarget(LLVMConversionTarget &target);
/// Register LLVM conversion interface for X86 dialect.
void registerConvertX86ToLLVMInterface(DialectRegistry ®istry);
-
} // namespace mlir
#endif // MLIR_DIALECT_X86_TRANSFORMS_H
diff --git a/mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt
index 5fe2fbc233eae7..b9036b78c3ac0a 100644
--- a/mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt
@@ -5,9 +5,12 @@ add_mlir_dialect_library(MLIRX86TransformOps
MLIRX86TransformOpsIncGen
LINK_LIBS PUBLIC
+ MLIRAffineDialect
+ MLIRArithDialect
MLIRIR
MLIRLLVMCommonConversion
MLIRLLVMDialect
+ MLIRMemRefDialect
MLIRVectorDialect
MLIRSideEffectInterfaces
MLIRTransformDialect
diff --git a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
index 1266e924b703d0..610c3dcff63348 100644
--- a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
+++ b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
@@ -8,7 +8,10 @@
#include "mlir/Dialect/X86/TransformOps/X86TransformOps.h"
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -57,6 +60,11 @@ void mlir::transform::ApplyMoveAccumulatorForContractLoopPatternsOp::
x86::populateMoveAccumulatorForContractLoopPatterns(patterns);
}
+void mlir::transform::ApplyPrepareVectorContractLoopsPatternsOp::
+ populatePatterns(RewritePatternSet &patterns) {
+ x86::populatePrepareVectorContractLoopsPatterns(patterns);
+}
+
void mlir::transform::ApplyVectorContractMultiLevelUnrollForAMXPatternsOp::
populatePatterns(RewritePatternSet &patterns) {
x86::populateVectorContractMultiLevelUnrollForAMXPatterns(patterns);
@@ -74,6 +82,9 @@ class X86TransformDialectExtension
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(X86TransformDialectExtension)
X86TransformDialectExtension() {
+ declareGeneratedDialect<affine::AffineDialect>();
+ declareGeneratedDialect<arith::ArithDialect>();
+ declareGeneratedDialect<memref::MemRefDialect>();
declareGeneratedDialect<x86::X86Dialect>();
declareGeneratedDialect<LLVM::LLVMDialect>();
registerTransformOps<
diff --git a/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt b/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
index f2cc234f3dc3ad..0e678fda61c37a 100644
--- a/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
@@ -9,9 +9,11 @@ add_mlir_dialect_library(MLIRX86Transforms
VectorContractToAMXDotProduct.cpp
MoveAccumulatorForContractLoop.cpp
VectorContractMultiLevelUnroll.cpp
+ PrepareVectorContractLoops.cpp
LINK_LIBS PUBLIC
MLIRArithDialect
+ MLIRDialectUtils
MLIRIR
MLIRLinalgDialect
MLIRLLVMCommonConversion
@@ -21,6 +23,7 @@ add_mlir_dialect_library(MLIRX86Transforms
MLIRUBDialect
MLIRVectorDialect
MLIRVectorTransforms
+ MLIRTransformUtils
MLIRVectorUtils
MLIRX86Dialect
MLIRX86Utils
diff --git a/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp b/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
new file mode 100644
index 00000000000000..88f9d59287cd35
--- /dev/null
+++ b/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
@@ -0,0 +1,316 @@
+//===- PrepareVectorContractLoops.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/Arith/IR/Arith.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/X86/Transforms.h"
+#include "mlir/Dialect/X86/Utils/X86Utils.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "x86-prepare-vector-contract-loops"
+
+using namespace mlir;
+
+namespace {
+vector::ContractionOp getCandidate(scf::ForOp loop) {
+ if (loop.getNumRegionIterArgs() != 1) {
+ LLVM_DEBUG(llvm::dbgs() << "reject loop at " << loop.getLoc()
+ << ": expected one iter_arg\n");
+ return {};
+ }
+
+ auto yield = cast<scf::YieldOp>(loop.getBody()->getTerminator());
+ for (Operation &op : loop.getBody()->without_terminator()) {
+ auto contract = dyn_cast<vector::ContractionOp>(&op);
+ if (!contract)
+ continue;
+ if (contract.getKind() != vector::CombiningKind::ADD ||
+ contract.getAcc() != loop.getRegionIterArgs().front() ||
+ !contract.getAcc().hasOneUse() || !contract.getResult().hasOneUse() ||
+ yield->getNumOperands() != 1 ||
+ yield->getOperand(0) != contract.getResult()) {
+ LLVM_DEBUG(llvm::dbgs() << "reject contraction at " << contract.getLoc()
+ << ": not a single loop-carried accumulation\n");
+ continue;
+ }
+
+ auto iterators = contract.getIteratorTypesArray();
+ auto maps = contract.getIndexingMapsArray();
+ MLIRContext *ctx = contract.getContext();
+ auto m = getAffineDimExpr(0, ctx);
+ auto n = getAffineDimExpr(1, ctx);
+ auto k = getAffineDimExpr(2, ctx);
+ auto par = vector::IteratorType::parallel;
+ auto red = vector::IteratorType::reduction;
+ if (llvm::equal(iterators, ArrayRef{par, par, red}) &&
+ maps[0] == AffineMap::get(3, 0, {m, k}, ctx) &&
+ maps[1] == AffineMap::get(3, 0, {k, n}, ctx) &&
+ maps[2] == AffineMap::get(3, 0, {m, n}, ctx))
+ return contract;
+
+ auto v = getAffineDimExpr(3, ctx);
+ if (x86::isInVnniLayout(contract, maps) &&
+ llvm::equal(iterators, ArrayRef{par, par, red, red}) &&
+ maps[0] == AffineMap::get(4, 0, {m, k, v}, ctx) &&
+ maps[1] == AffineMap::get(4, 0, {k, n, v}, ctx) &&
+ maps[2] == AffineMap::get(4, 0, {m, n}, ctx))
+ return contract;
+ LLVM_DEBUG(llvm::dbgs() << "reject contraction at " << contract.getLoc()
+ << ": noncanonical matmul layout or iterators\n");
+ }
+ return {};
+}
+
+bool isSimpleReadOutsideLoop(vector::TransferReadOp read, scf::ForOp loop) {
+ return read && read->getParentOp() == loop &&
+ isa<MemRefType>(read.getBase().getType()) &&
+ loop.isDefinedOutsideOfLoop(read.getBase()) &&
+ read.getPermutationMap().isMinorIdentity() && !read.getMask() &&
+ !read.hasOutOfBoundsDim();
+}
+
+struct ExpandedRead {
+ vector::TransferReadOp read;
+ MemRefType type;
+ SmallVector<ReassociationIndices> reassociation;
+ int64_t lastIndex;
+ // Static view of a dynamic last dimension, or null.
+ MemRefType view;
+};
+
+std::optional<ExpandedRead> getExpandableRead(vector::TransferReadOp read,
+ unsigned factor,
+ scf::ForOp loop) {
+ if (!isSimpleReadOutsideLoop(read, loop))
+ return std::nullopt;
+ auto type = cast<MemRefType>(read.getBase().getType());
+ if (type.getElementType() != read.getVectorType().getElementType() ||
+ type.getRank() < 2)
+ return std::nullopt;
+ auto lastIndex = getConstantIntValue(read.getIndices().back());
+ if (!lastIndex || *lastIndex < 0 || *lastIndex % factor != 0)
+ return std::nullopt;
+
+ MemRefType view;
+ if (type.isDynamicDim(type.getRank() - 1)) {
+ SmallVector<int64_t> strides;
+ int64_t offset;
+ if (read.getVectorType().isScalable() ||
+ failed(type.getStridesAndOffset(strides, offset)))
+ return std::nullopt;
+ // The in-bounds read covers this prefix of the last dimension.
+ SmallVector<int64_t> shape(type.getShape());
+ shape.back() = *lastIndex + read.getVectorType().getShape().back();
+ view = type = MemRefType::get(
+ shape, type.getElementType(),
+ StridedLayoutAttr::get(type.getContext(), offset, strides),
+ type.getMemorySpace());
+ }
+ if (type.getShape().back() % factor != 0)
+ return std::nullopt;
+
+ SmallVector<int64_t> shape(type.getShape());
+ shape.back() /= factor;
+ shape.push_back(factor);
+ SmallVector<ReassociationIndices> reassociation;
+ for (int64_t i = 0; i < type.getRank(); ++i)
+ reassociation.push_back(i == type.getRank() - 1
+ ? ReassociationIndices{i, i + 1}
+ : ReassociationIndices{i});
+ FailureOr<MemRefType> expanded =
+ memref::ExpandShapeOp::computeExpandedType(type, shape, reassociation);
+ if (failed(expanded))
+ return std::nullopt;
+ return ExpandedRead{read, *expanded, reassociation, *lastIndex / factor,
+ view};
+}
+
+Value buildExpandedRead(PatternRewriter &rewriter, Location loc,
+ ExpandedRead &info, VectorType resultType) {
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(info.read->getParentOp());
+ Value base = info.read.getBase();
+ if (info.view) {
+ // Unlike subview, this view stays valid when the loop runs zero times.
+ auto metadata =
+ memref::ExtractStridedMetadataOp::create(rewriter, loc, base);
+ SmallVector<OpFoldResult> sizes = metadata.getConstifiedMixedSizes();
+ sizes.back() = rewriter.getIndexAttr(info.view.getShape().back());
+ base = memref::ReinterpretCastOp::create(
+ rewriter, loc, info.view, base, metadata.getConstifiedMixedOffset(),
+ sizes, metadata.getConstifiedMixedStrides());
+ }
+ Value expanded = memref::ExpandShapeOp::create(rewriter, loc, info.type, base,
+ info.reassociation);
+ Value lastIndex =
+ arith::ConstantIndexOp::create(rewriter, loc, info.lastIndex);
+ Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ rewriter.setInsertionPoint(info.read);
+ SmallVector<Value> indices(info.read.getIndices());
+ indices.back() = lastIndex;
+ indices.push_back(zero);
+ auto packed = vector::TransferReadOp::create(
+ rewriter, loc, resultType, expanded, indices, info.read.getPadding());
+ packed.setInBoundsAttr(
+ rewriter.getBoolArrayAttr(SmallVector<bool>(resultType.getRank(), true)));
+ return packed;
+}
+
+// Fold RHS unpacking into a VNNI contraction:
+// - Match a packed RHS read transposed and flattened for a flat matmul.
+// - Split the LHS memref's K dimension and reread it in packed lanes.
+// A dynamic K dimension is first narrowed to the read's static extent.
+// - Contract the packed reads directly over K and the VNNI lane dimension.
+// For bf16, 32x32 LHS and RHS become 32x16x2 LHS and 16x32x2 RHS.
+struct FoldVNNIUnpacking : OpRewritePattern<vector::ContractionOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::ContractionOp contract,
+ PatternRewriter &rewriter) const override {
+ auto loop = dyn_cast<scf::ForOp>(contract->getParentOp());
+ if (!loop)
+ return rewriter.notifyMatchFailure(contract,
+ "not directly in a for loop");
+ if (getCandidate(loop) != contract)
+ return rewriter.notifyMatchFailure(
+ contract, "not the loop's canonical contraction");
+
+ auto lhsType = contract.getLhsType();
+ auto rhsType = contract.getRhsType();
+ auto accType = dyn_cast<VectorType>(contract.getAccType());
+ Type element = lhsType.getElementType();
+ if (!accType || lhsType.getRank() != 2 || rhsType.getRank() != 2 ||
+ accType.getRank() != 2)
+ return rewriter.notifyMatchFailure(
+ contract, "expected 2D operands and accumulator");
+ if (lhsType.isScalable() || rhsType.isScalable())
+ return rewriter.notifyMatchFailure(contract, "scalable operand vectors");
+ if (element != rhsType.getElementType())
+ return rewriter.notifyMatchFailure(contract,
+ "operand element types differ");
+ if (!(element.isBF16() || element.isInteger(8)))
+ return rewriter.notifyMatchFailure(contract,
+ "expected bf16 or i8 operands");
+ unsigned factor = 32 / element.getIntOrFloatBitWidth();
+ if (lhsType.getDimSize(1) != rhsType.getDimSize(0))
+ return rewriter.notifyMatchFailure(contract, "K dimensions differ");
+ if (lhsType.getDimSize(1) % factor != 0)
+ return rewriter.notifyMatchFailure(
+ contract, "K dimension is not divisible by VNNI factor");
+
+ auto shapeCast = contract.getRhs().getDefiningOp<vector::ShapeCastOp>();
+ auto transpose =
+ shapeCast ? shapeCast.getSource().getDefiningOp<vector::TransposeOp>()
+ : vector::TransposeOp();
+ if (!transpose || transpose->getParentOp() != loop)
+ return rewriter.notifyMatchFailure(
+ contract, "RHS is not an in-loop transpose followed by shape_cast");
+ VectorType rhsPackedType =
+ VectorType::get({lhsType.getDimSize(1) / factor, rhsType.getDimSize(1),
+ int64_t(factor)},
+ element);
+ auto sourceType = cast<VectorType>(transpose.getVector().getType());
+ if (sourceType.getRank() < rhsPackedType.getRank())
+ return rewriter.notifyMatchFailure(contract,
+ "RHS source rank is too small");
+ unsigned leading = sourceType.getRank() - rhsPackedType.getRank();
+ if (!llvm::all_of(sourceType.getShape().take_front(leading),
+ [](int64_t dim) { return dim == 1; }) ||
+ !llvm::equal(sourceType.getShape().drop_front(leading),
+ rhsPackedType.getShape()))
+ return rewriter.notifyMatchFailure(contract,
+ "RHS source is not VNNI packed");
+ for (unsigned i = 0; i < leading + 1; ++i)
+ if (transpose.getPermutation()[i] != i)
+ return rewriter.notifyMatchFailure(
+ contract, "transpose reorders leading or K dimensions");
+ if (transpose.getPermutation()[leading + 1] != leading + 2 ||
+ transpose.getPermutation()[leading + 2] != leading + 1)
+ return rewriter.notifyMatchFailure(
+ contract, "transpose does not swap N and VNNI lanes");
+
+ auto lhsRead = contract.getLhs().getDefiningOp<vector::TransferReadOp>();
+ auto lhs = getExpandableRead(lhsRead, factor, loop);
+ auto broadcast = transpose.getVector().getDefiningOp<vector::BroadcastOp>();
+ Value packedSource =
+ broadcast ? broadcast.getSource() : transpose.getVector();
+ if (broadcast && packedSource.getType() != rhsPackedType)
+ return rewriter.notifyMatchFailure(contract,
+ "broadcast source is not VNNI packed");
+ auto packedCast = packedSource.getDefiningOp<vector::ShapeCastOp>();
+ auto rhsRead =
+ packedCast
+ ? packedCast.getSource().getDefiningOp<vector::TransferReadOp>()
+ : packedSource.getDefiningOp<vector::TransferReadOp>();
+ if (!lhs)
+ return rewriter.notifyMatchFailure(
+ contract, "LHS read cannot be expanded into VNNI layout");
+ if (!isSimpleReadOutsideLoop(rhsRead, loop))
+ return rewriter.notifyMatchFailure(contract,
+ "RHS is not a simple transfer read");
+ std::optional<ExpandedRead> splitRhs;
+ if (packedCast) {
+ auto flatType = rhsRead.getVectorType();
+ if (flatType.getRank() != 2 ||
+ flatType.getDimSize(0) != rhsPackedType.getDimSize(0) ||
+ flatType.getDimSize(1) % factor != 0 ||
+ flatType.getDimSize(1) / factor != rhsPackedType.getDimSize(1))
+ return rewriter.notifyMatchFailure(
+ contract, "flat RHS shape is not VNNI compatible");
+ if (!(splitRhs = getExpandableRead(rhsRead, factor, loop)))
+ return rewriter.notifyMatchFailure(
+ contract, "RHS read cannot be expanded into VNNI layout");
+ } else {
+ if (rhsRead.getVectorType() != rhsPackedType)
+ return rewriter.notifyMatchFailure(contract,
+ "RHS read is not VNNI packed");
+ if (getConstantIntValue(rhsRead.getIndices().back()) != 0)
+ return rewriter.notifyMatchFailure(
+ contract, "packed RHS has nonzero inner offset");
+ }
+
+ Location loc = contract.getLoc();
+ VectorType lhsPackedType =
+ VectorType::get({lhsType.getDimSize(0), lhsType.getDimSize(1) / factor,
+ int64_t(factor)},
+ element);
+ Value packedLhs = buildExpandedRead(rewriter, loc, *lhs, lhsPackedType);
+ Value packedRhs =
+ splitRhs ? buildExpandedRead(rewriter, loc, *splitRhs, rhsPackedType)
+ : rhsRead.getResult();
+ rewriter.setInsertionPoint(contract);
+ auto ctx = rewriter.getContext();
+ auto m = rewriter.getAffineDimExpr(0);
+ auto n = rewriter.getAffineDimExpr(1);
+ auto k = rewriter.getAffineDimExpr(2);
+ auto v = rewriter.getAffineDimExpr(3);
+ auto vnniMaps =
+ rewriter.getAffineMapArrayAttr({AffineMap::get(4, 0, {m, k, v}, ctx),
+ AffineMap::get(4, 0, {k, n, v}, ctx),
+ AffineMap::get(4, 0, {m, n}, ctx)});
+ auto iterators = rewriter.getArrayAttr(
+ {vector::IteratorTypeAttr::get(ctx, vector::IteratorType::parallel),
+ vector::IteratorTypeAttr::get(ctx, vector::IteratorType::parallel),
+ vector::IteratorTypeAttr::get(ctx, vector::IteratorType::reduction),
+ vector::IteratorTypeAttr::get(ctx, vector::IteratorType::reduction)});
+ rewriter.replaceOpWithNewOp<vector::ContractionOp>(
+ contract, packedLhs, packedRhs, contract.getAcc(), vnniMaps, iterators,
+ contract.getKind(), contract.getFastmath());
+ return success();
+ }
+};
+} // namespace
+
+void x86::populatePrepareVectorContractLoopsPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<FoldVNNIUnpacking>(patterns.getContext());
+}
\ No newline at end of file
diff --git a/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir b/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
new file mode 100644
index 00000000000000..865aba34cab350
--- /dev/null
+++ b/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
@@ -0,0 +1,379 @@
+// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter=entry-point=__transform_nano | FileCheck %s --check-prefix=NANO
+
+#lhs = affine_map<(m, n, k) -> (m, k)>
+#rhs = affine_map<(m, n, k) -> (k, n)>
+#acc = affine_map<(m, n, k) -> (m, n)>
+
+// CHECK-DAG: #[[$VNNI_LHS:[A-Za-z0-9_]+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+// CHECK-DAG: #[[$VNNI_RHS:[A-Za-z0-9_]+]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+// CHECK-DAG: #[[$VNNI_ACC:[A-Za-z0-9_]+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+
+// fold VNNI unpacking of packed bf16 rhs read into contraction, keeping fastmath
+// CHECK-LABEL: func.func @direct_vnni_bf16(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x32xbf16>, %[[B:[^,]+]]: memref<?x16x32x2xbf16>
+// NANO-LABEL: func.func @direct_vnni_bf16(
+// NANO-NOT: vector.contract
+// NANO: return
+// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x2xbf16>
+// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A_PACKED]][%[[K]], {{.*}} vector<32x16x2xbf16>
+// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B]][%[[K]], {{.*}} vector<16x32x2xbf16>
+// CHECK: vector.contract {fastmath = #arith.fastmath<reassoc>, indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+func.func @direct_vnni_bf16(%a: memref<?x32x32xbf16>,
+ %b: memref<?x16x32x2xbf16>,
+ %init: vector<32x32xf32>, %ub: index,
+ %m: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %m, %zero], %pad {in_bounds = [true, true]} : memref<?x32x32xbf16>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x2xbf16>, vector<16x32x2xbf16>
+ %transpose = vector.transpose %rhs, [0, 2, 1] : vector<16x32x2xbf16> to vector<16x2x32xbf16>
+ %flat = vector.shape_cast %transpose : vector<16x2x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>, fastmath = #arith.fastmath<reassoc>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %dot : vector<32x32xf32>
+ }
+ return %result : vector<32x32xf32>
+}
+
+// fold VNNI unpacking of packed bf16 rhs read with leading unit dims
+// CHECK-LABEL: func.func @leading_unit_vnni(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x32xbf16>, %[[B:[^,]+]]: memref<?x1x1x16x32x2xbf16>
+// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x2xbf16>
+// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A_PACKED]][%[[K]], {{.*}} vector<32x16x2xbf16>
+// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B]][%[[K]], {{.*}} vector<16x32x2xbf16>
+// CHECK: vector.contract {indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+// NANO-LABEL: func.func @leading_unit_vnni(
+// NANO-NOT: vector.contract
+// NANO: return
+func.func @leading_unit_vnni(%a: memref<?x32x32xbf16>,
+ %b: memref<?x1x1x16x32x2xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x32x32xbf16>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero, %zero, %zero], %pad {in_bounds = [true, true, true, true, true]} : memref<?x1x1x16x32x2xbf16>, vector<1x1x16x32x2xbf16>
+ %transpose = vector.transpose %rhs, [0, 1, 2, 4, 3] : vector<1x1x16x32x2xbf16> to vector<1x1x16x2x32xbf16>
+ %flat = vector.shape_cast %transpose : vector<1x1x16x2x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %dot : vector<32x32xf32>
+ }
+ return %result : vector<32x32xf32>
+}
+
+// fold VNNI unpacking of flat bf16 rhs read split into lanes by shape_cast
+// CHECK-LABEL: func.func @split_vnni_bf16(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x32xbf16>, %[[B:[^,]+]]: memref<?x16x64xbf16>
+// NANO-LABEL: func.func @split_vnni_bf16(
+// NANO-NOT: vector.contract
+// NANO: return
+// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x2xbf16>
+// CHECK: %[[B_PACKED:.+]] = memref.expand_shape %[[B]] {{.*}} into memref<?x16x32x2xbf16>
+// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A_PACKED]][%[[K]], {{.*}} vector<32x16x2xbf16>
+// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B_PACKED]][%[[K]], {{.*}} vector<16x32x2xbf16>
+// CHECK: vector.contract {indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+func.func @split_vnni_bf16(%a: memref<?x32x32xbf16>,
+ %b: memref<?x16x64xbf16>,
+ %init: vector<32x32xf32>, %ub: index,
+ %m: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %m, %zero], %pad {in_bounds = [true, true]} : memref<?x32x32xbf16>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x64xbf16>
+ %packed = vector.shape_cast %rhs : vector<16x64xbf16> to vector<16x32x2xbf16>
+ %transpose = vector.transpose %packed, [0, 2, 1] : vector<16x32x2xbf16> to vector<16x2x32xbf16>
+ %flat = vector.shape_cast %transpose : vector<16x2x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %dot : vector<32x32xf32>
+ }
+ return %result : vector<32x32xf32>
+}
+
+// fold VNNI unpacking of packed i8 rhs read into contraction
+// CHECK-LABEL: func.func @direct_vnni_int8(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x64xi8>, %[[B:[^,]+]]: memref<?x16x32x4xi8>
+// NANO-LABEL: func.func @direct_vnni_int8(
+// NANO-NOT: vector.contract
+// NANO: return
+// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x4xi8>
+// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A_PACKED]][%[[K]], {{.*}} vector<32x16x4xi8>
+// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B]][%[[K]], {{.*}} vector<16x32x4xi8>
+// CHECK: vector.contract {indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x4xi8>, vector<16x32x4xi8> into vector<32x32xi32>
+func.func @direct_vnni_int8(%a: memref<?x32x64xi8>,
+ %b: memref<?x16x32x4xi8>,
+ %init: vector<32x32xi32>, %ub: index,
+ %m: index) -> vector<32x32xi32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0 : i8
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xi32>) {
+ %lhs = vector.transfer_read %a[%k, %m, %zero], %pad {in_bounds = [true, true]} : memref<?x32x64xi8>, vector<32x64xi8>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x4xi8>, vector<16x32x4xi8>
+ %transpose = vector.transpose %rhs, [0, 2, 1] : vector<16x32x4xi8> to vector<16x4x32xi8>
+ %flat = vector.shape_cast %transpose : vector<16x4x32xi8> to vector<64x32xi8>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x64xi8>, vector<64x32xi8> into vector<32x32xi32>
+ scf.yield %dot : vector<32x32xi32>
+ }
+ return %result : vector<32x32xi32>
+}
+
+// fold VNNI unpacking of flat i8 rhs read split into lanes by shape_cast
+// CHECK-LABEL: func.func @split_vnni_int8(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x64xi8>, %[[B:[^,]+]]: memref<?x16x128xi8>
+// NANO-LABEL: func.func @split_vnni_int8(
+// NANO-NOT: vector.contract
+// NANO: return
+// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x4xi8>
+// CHECK: %[[B_PACKED:.+]] = memref.expand_shape %[[B]] {{.*}} into memref<?x16x32x4xi8>
+// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A_PACKED]][%[[K]], {{.*}} vector<32x16x4xi8>
+// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B_PACKED]][%[[K]], {{.*}} vector<16x32x4xi8>
+// CHECK: vector.contract {indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x4xi8>, vector<16x32x4xi8> into vector<32x32xi32>
+func.func @split_vnni_int8(%a: memref<?x32x64xi8>,
+ %b: memref<?x16x128xi8>,
+ %init: vector<32x32xi32>, %ub: index,
+ %m: index) -> vector<32x32xi32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0 : i8
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xi32>) {
+ %lhs = vector.transfer_read %a[%k, %m, %zero], %pad {in_bounds = [true, true]} : memref<?x32x64xi8>, vector<32x64xi8>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x16x128xi8>, vector<16x128xi8>
+ %packed = vector.shape_cast %rhs : vector<16x128xi8> to vector<16x32x4xi8>
+ %transpose = vector.transpose %packed, [0, 2, 1] : vector<16x32x4xi8> to vector<16x4x32xi8>
+ %flat = vector.shape_cast %transpose : vector<16x4x32xi8> to vector<64x32xi8>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x64xi8>, vector<64x32xi8> into vector<32x32xi32>
+ scf.yield %dot : vector<32x32xi32>
+ }
+ return %result : vector<32x32xi32>
+}
+
+// cast away leading unit dims of reads so that the contraction can be unrolled
+// CHECK-LABEL: func.func @unit_dimension_reads(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x?x32x32xbf16>, %[[B:[^,]+]]: memref<?x?x32x32xbf16>
+// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][%[[K]], {{.*}} vector<32x32xbf16>
+// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B]][%[[K]], {{.*}} vector<32x32xbf16>
+// CHECK: vector.contract {{.*}} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+// NANO-LABEL: func.func @unit_dimension_reads(
+// NANO-NOT: vector.contract
+// NANO: return
+func.func @unit_dimension_reads(%a: memref<?x?x32x32xbf16>,
+ %b: memref<?x?x32x32xbf16>,
+ %init: vector<32x32xf32>, %ub: index,
+ %m: index, %n: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 32 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %m, %zero, %zero], %pad {in_bounds = [true, true, true, true]} : memref<?x?x32x32xbf16>, vector<1x1x32x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %n, %zero, %zero], %pad {in_bounds = [true, true, true, true]} : memref<?x?x32x32xbf16>, vector<1x1x32x32xbf16>
+ %left = vector.shape_cast %lhs : vector<1x1x32x32xbf16> to vector<32x32xbf16>
+ %right = vector.shape_cast %rhs : vector<1x1x32x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %right, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %dot : vector<32x32xf32>
+ }
+ return %result : vector<32x32xf32>
+}
+
+// fold VNNI unpacking through a static view of a dynamic lhs innermost dim
+// CHECK-LABEL: func.func @dynamic_inner_k(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x?xbf16>
+// NANO-LABEL: func.func @dynamic_inner_k(
+// NANO-NOT: vector.contract
+// NANO: return
+// CHECK: %[[VIEW:.+]] = memref.reinterpret_cast %[[A]] {{.*}} sizes: [%{{.+}}, 32, 32]{{.*}} to memref<?x32x32xbf16, strided<[?, ?, 1]>>
+// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[VIEW]] {{.*}} into memref<?x32x16x2xbf16, strided<[?, ?, 2, 1]>>
+// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A_PACKED]][%[[K]], {{.*}} vector<32x16x2xbf16>
+// CHECK: vector.contract {indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %{{.+}}, %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+func.func @dynamic_inner_k(%a: memref<?x32x?xbf16>,
+ %b: memref<?x16x32x2xbf16>,
+ %init: vector<32x32xf32>, %ub: index,
+ %m: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %m, %zero], %pad {in_bounds = [true, true]} : memref<?x32x?xbf16>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x2xbf16>, vector<16x32x2xbf16>
+ %transpose = vector.transpose %rhs, [0, 2, 1] : vector<16x32x2xbf16> to vector<16x2x32xbf16>
+ %flat = vector.shape_cast %transpose : vector<16x2x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %dot : vector<32x32xf32>
+ }
+ return %result : vector<32x32xf32>
+}
+
+// negative test: lhs innermost read offset 1 is not a multiple of VNNI factor 2 -> skip
+// CHECK-LABEL: func.func @misaligned_inner_k(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x64xbf16>
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x32xbf16>
+// CHECK: %[[T:.+]] = vector.transpose
+// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x2x32xbf16> to vector<32x32xbf16>
+// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @misaligned_inner_k(%a: memref<?x32x64xbf16>,
+ %b: memref<?x16x32x2xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %one = arith.constant 1 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %zero, %one], %pad {in_bounds = [true, true]} : memref<?x32x64xbf16>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x2xbf16>, vector<16x32x2xbf16>
+ %transpose = vector.transpose %rhs, [0, 2, 1] : vector<16x32x2xbf16> to vector<16x2x32xbf16>
+ %flat = vector.shape_cast %transpose : vector<16x2x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %dot : vector<32x32xf32>
+ }
+ return %result : vector<32x32xf32>
+}
+
+// negative test: lhs innermost dim 33 is not divisible by VNNI factor 2 -> skip
+// CHECK-LABEL: func.func @nondivisible_inner_k(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x33xbf16>
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x32xbf16>
+// CHECK: %[[T:.+]] = vector.transpose
+// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x2x32xbf16> to vector<32x32xbf16>
+// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @nondivisible_inner_k(%a: memref<?x32x33xbf16>,
+ %b: memref<?x16x32x2xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x32x33xbf16>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x2xbf16>, vector<16x32x2xbf16>
+ %transpose = vector.transpose %rhs, [0, 2, 1] : vector<16x32x2xbf16> to vector<16x2x32xbf16>
+ %flat = vector.shape_cast %transpose : vector<16x2x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %dot : vector<32x32xf32>
+ }
+ return %result : vector<32x32xf32>
+}
+
+// negative test: rhs VNNI lanes are reordered before unpacking -> skip
+// CHECK-LABEL: func.func @interleaved_lanes(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x64xi8>
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x64xi8>
+// CHECK: %[[EVEN:[^,]+]], %[[ODD:.+]] = vector.deinterleave
+// CHECK: %[[LANES:.+]] = vector.interleave %[[ODD]], %[[EVEN]]
+// CHECK: %[[T:.+]] = vector.transpose %[[LANES]], [0, 2, 1]
+// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x4x32xi8> to vector<64x32xi8>
+// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<32x64xi8>, vector<64x32xi8> into vector<32x32xi32>
+func.func @interleaved_lanes(%a: memref<?x32x64xi8>,
+ %b: memref<?x16x32x4xi8>,
+ %init: vector<32x32xi32>, %ub: index) -> vector<32x32xi32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0 : i8
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xi32>) {
+ %lhs = vector.transfer_read %a[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x32x64xi8>, vector<32x64xi8>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x4xi8>, vector<16x32x4xi8>
+ %even, %odd = vector.deinterleave %rhs : vector<16x32x4xi8> -> vector<16x32x2xi8>
+ %reordered = vector.interleave %odd, %even : vector<16x32x2xi8> -> vector<16x32x4xi8>
+ %transpose = vector.transpose %reordered, [0, 2, 1] : vector<16x32x4xi8> to vector<16x4x32xi8>
+ %flat = vector.shape_cast %transpose : vector<16x4x32xi8> to vector<64x32xi8>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x64xi8>, vector<64x32xi8> into vector<32x32xi32>
+ scf.yield %dot : vector<32x32xi32>
+ }
+ return %result : vector<32x32xi32>
+}
+
+// negative test: scalable lhs cannot be packed into a fixed VNNI shape -> skip
+// CHECK-LABEL: func.func @scalable_vnni(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x?x32xbf16>
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<[32]x32xbf16>
+// CHECK: %[[T:.+]] = vector.transpose
+// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x2x32xbf16> to vector<32x32xbf16>
+// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<[32]x32xbf16>, vector<32x32xbf16> into vector<[32]x32xf32>
+func.func @scalable_vnni(%a: memref<?x?x32xbf16>,
+ %b: memref<?x16x32x2xbf16>,
+ %init: vector<[32]x32xf32>, %ub: index,
+ %m: index) -> vector<[32]x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<[32]x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %m, %zero], %pad {in_bounds = [true, true]} : memref<?x?x32xbf16>, vector<[32]x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x2xbf16>, vector<16x32x2xbf16>
+ %transpose = vector.transpose %rhs, [0, 2, 1] : vector<16x32x2xbf16> to vector<16x2x32xbf16>
+ %flat = vector.shape_cast %transpose : vector<16x2x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<[32]x32xbf16>, vector<32x32xbf16> into vector<[32]x32xf32>
+ scf.yield %dot : vector<[32]x32xf32>
+ }
+ return %result : vector<[32]x32xf32>
+}
+
+// negative test: lhs memref with vector elements cannot be expanded into VNNI layout -> skip
+// CHECK-LABEL: func.func @vector_element_lhs(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32xvector<32xbf16>>
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x32xbf16>
+// CHECK: %[[T:.+]] = vector.transpose
+// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x2x32xbf16> to vector<32x32xbf16>
+// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @vector_element_lhs(%a: memref<?x32xvector<32xbf16>>,
+ %b: memref<?x16x32x2xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 16 : index
+ %pad = arith.constant 0.0 : bf16
+ %vpad = arith.constant dense<0.0> : vector<32xbf16>
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%k, %zero], %vpad {in_bounds = [true]} : memref<?x32xvector<32xbf16>>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x2xbf16>, vector<16x32x2xbf16>
+ %transpose = vector.transpose %rhs, [0, 2, 1] : vector<16x32x2xbf16> to vector<16x2x32xbf16>
+ %flat = vector.shape_cast %transpose : vector<16x2x32xbf16> to vector<32x32xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+ scf.yield %dot : vector<32x32xf32>
+ }
+ return %result : vector<32x32xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ %loops = transform.structured.match ops{["scf.for"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_licm to %loops : !transform.any_op
+ transform.apply_patterns to %funcs {
+ transform.apply_patterns.vector.cast_away_vector_leading_one_dim
+ transform.apply_patterns.x86.prepare_vector_contract_loops
+ } : !transform.any_op
+ transform.apply_dce to %funcs : !transform.any_op
+ transform.apply_cse to %funcs : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @__transform_nano(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ %loops = transform.structured.match ops{["scf.for"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_licm to %loops : !transform.any_op
+ transform.apply_patterns to %funcs {
+ transform.apply_patterns.vector.cast_away_vector_leading_one_dim
+ transform.apply_patterns.x86.prepare_vector_contract_loops
+ } : !transform.any_op
+ transform.apply_dce to %funcs : !transform.any_op
+ transform.apply_cse to %funcs : !transform.any_op
+ transform.apply_patterns to %funcs {
+ transform.apply_patterns.x86.vector_contract_multi_level_unroll_for_amx
+ transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
\ No newline at end of file
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/X86/prepare-contract-loops.mlir b/mlir/test/Integration/Dialect/Vector/CPU/X86/prepare-contract-loops.mlir
new file mode 100644
index 00000000000000..d5deb2b253330b
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Vector/CPU/X86/prepare-contract-loops.mlir
@@ -0,0 +1,87 @@
+// RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-vector-to-llvm='vector-contract-lowering=outerproduct' -test-lower-to-llvm | mlir-runner -e main -entry-point-result=void -shared-libs=%mlir_c_runner_utils | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter -test-transform-dialect-erase-schedule -convert-vector-to-llvm='vector-contract-lowering=outerproduct' -test-lower-to-llvm | mlir-runner -e main -entry-point-result=void -shared-libs=%mlir_c_runner_utils | FileCheck %s
+
+#lhs = affine_map<(m, n, k) -> (m, k)>
+#rhs = affine_map<(m, n, k) -> (k, n)>
+#acc = affine_map<(m, n, k) -> (m, n)>
+
+// Generic CPU lowering executes VNNI contractions only without mixed precision.
+func.func @vnni_bf16(%a: memref<1x2x8xbf16>,
+ %b: memref<1x4x4x2xbf16>) -> vector<2x4xbf16> {
+ %zero = arith.constant 0 : index
+ %end = arith.constant 1 : index
+ %pad = arith.constant 0.0 : bf16
+ %init = arith.constant dense<0.0> : vector<2x4xbf16>
+ %result = scf.for %k = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x4xbf16>) {
+ %left = vector.transfer_read %a[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<1x2x8xbf16>, vector<2x8xbf16>
+ %right = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<1x4x4x2xbf16>, vector<4x4x2xbf16>
+ %transpose = vector.transpose %right, [0, 2, 1] : vector<4x4x2xbf16> to vector<4x2x4xbf16>
+ %flat = vector.shape_cast %transpose : vector<4x2x4xbf16> to vector<8x4xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %flat, %accum : vector<2x8xbf16>, vector<8x4xbf16> into vector<2x4xbf16>
+ scf.yield %dot : vector<2x4xbf16>
+ }
+ return %result : vector<2x4xbf16>
+}
+
+func.func @vnni_i8(%a: memref<1x2x8xi8>,
+ %b: memref<1x2x4x4xi8>) -> vector<2x4xi8> {
+ %zero = arith.constant 0 : index
+ %end = arith.constant 1 : index
+ %pad = arith.constant 0 : i8
+ %init = arith.constant dense<0> : vector<2x4xi8>
+ %result = scf.for %k = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x4xi8>) {
+ %left = vector.transfer_read %a[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<1x2x8xi8>, vector<2x8xi8>
+ %right = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<1x2x4x4xi8>, vector<2x4x4xi8>
+ %transpose = vector.transpose %right, [0, 2, 1] : vector<2x4x4xi8> to vector<2x4x4xi8>
+ %flat = vector.shape_cast %transpose : vector<2x4x4xi8> to vector<8x4xi8>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %flat, %accum : vector<2x8xi8>, vector<8x4xi8> into vector<2x4xi8>
+ scf.yield %dot : vector<2x4xi8>
+ }
+ return %result : vector<2x4xi8>
+}
+
+func.func @main() {
+ %zero = arith.constant 0 : index
+ %af = memref.alloc() : memref<1x2x8xbf16>
+ %bf = memref.alloc() : memref<1x4x4x2xbf16>
+ %ones_af = arith.constant dense<[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
+ [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]]> : vector<2x8xbf16>
+ %ones_bf = arith.constant dense<[
+ [[2.0, 3.0], [2.0, 3.0], [2.0, 3.0], [2.0, 3.0]],
+ [[2.0, 3.0], [2.0, 3.0], [2.0, 3.0], [2.0, 3.0]],
+ [[2.0, 3.0], [2.0, 3.0], [2.0, 3.0], [2.0, 3.0]],
+ [[2.0, 3.0], [2.0, 3.0], [2.0, 3.0], [2.0, 3.0]]
+ ]> : vector<4x4x2xbf16>
+ vector.transfer_write %ones_af, %af[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<2x8xbf16>, memref<1x2x8xbf16>
+ vector.transfer_write %ones_bf, %bf[%zero, %zero, %zero, %zero] {in_bounds = [true, true, true]} : vector<4x4x2xbf16>, memref<1x4x4x2xbf16>
+ %bf_result = func.call @vnni_bf16(%af, %bf) : (memref<1x2x8xbf16>, memref<1x4x4x2xbf16>) -> vector<2x4xbf16>
+ %bf_value = vector.extract %bf_result[0, 0] : bf16 from vector<2x4xbf16>
+ vector.print %bf_value : bf16
+ // CHECK: 92
+
+ %a4 = memref.alloc() : memref<1x2x8xi8>
+ %b4 = memref.alloc() : memref<1x2x4x4xi8>
+ %ones_a4 = arith.constant dense<[[1, 2, 3, 4, 5, 6, 7, 8],
+ [1, 2, 3, 4, 5, 6, 7, 8]]> : vector<2x8xi8>
+ %ones_b4 = arith.constant dense<[
+ [[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]],
+ [[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]
+ ]> : vector<2x4x4xi8>
+ vector.transfer_write %ones_a4, %a4[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<2x8xi8>, memref<1x2x8xi8>
+ vector.transfer_write %ones_b4, %b4[%zero, %zero, %zero, %zero] {in_bounds = [true, true, true]} : vector<2x4x4xi8>, memref<1x2x4x4xi8>
+ %i8_result = func.call @vnni_i8(%a4, %b4) : (memref<1x2x8xi8>, memref<1x2x4x4xi8>) -> vector<2x4xi8>
+ %i8_value = vector.extract %i8_result[0, 0] : i8 from vector<2x4xi8>
+ vector.print %i8_value : i8
+ // CHECK-NEXT: 100
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %funcs {
+ transform.apply_patterns.x86.prepare_vector_contract_loops
+ } : !transform.any_op
+ transform.yield
+ }
+}
\ No newline at end of file
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 7558c476e96d1e..a36fa4578b2f2c 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -2636,9 +2636,12 @@ cc_library(
hdrs = glob(["include/mlir/Dialect/X86/TransformOps/*.h"]),
includes = ["include"],
deps = [
+ ":AffineDialect",
+ ":ArithDialect",
":IR",
":LLVMCommonConversion",
":LLVMDialect",
+ ":MemRefDialect",
":TransformDialect",
":TransformDialectInterfaces",
":VectorDialect",
@@ -2671,9 +2674,12 @@ cc_library(
hdrs = ["include/mlir/Dialect/X86/Transforms.h"],
includes = ["include"],
deps = [
+ ":AffineDialect",
":ArithDialect",
":ControlFlowDialect",
":ConvertToLLVMInterface",
+ ":DataLayoutInterfaces",
+ ":DialectUtils",
":IR",
":LLVMCommonConversion",
":LLVMDialect",
@@ -2683,10 +2689,13 @@ cc_library(
":MemRefUtils",
":Pass",
":SCFDialect",
+ ":SCFUtils",
":TransformUtils",
":UBDialect",
":VectorDialect",
":VectorTransforms",
+ ":ValueBoundsOpInterface",
+ ":VectorDialect",
":VectorUtils",
":X86Dialect",
":X86Utils",
>From 978d8140e4fd280402c2d61d33e2aaecf5d58915 Mon Sep 17 00:00:00 2001
From: fschlimb <frank.schlimbach at intel.com>
Date: Fri, 16 Oct 2026 08:32:47 -0700
Subject: [PATCH 2/4] removing duplicate lib
---
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel | 1 -
1 file changed, 1 deletion(-)
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index a36fa4578b2f2c..15a23d2e41a8b1 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -2695,7 +2695,6 @@ cc_library(
":VectorDialect",
":VectorTransforms",
":ValueBoundsOpInterface",
- ":VectorDialect",
":VectorUtils",
":X86Dialect",
":X86Utils",
>From d8662136eac50f7dd04abecfec0b71fa81461958 Mon Sep 17 00:00:00 2001
From: fschlimb <frank.schlimbach at intel.com>
Date: Mon, 26 Oct 2026 03:53:14 -0700
Subject: [PATCH 3/4] revert bazel changes
---
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel | 8 --------
1 file changed, 8 deletions(-)
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 15a23d2e41a8b1..7558c476e96d1e 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -2636,12 +2636,9 @@ cc_library(
hdrs = glob(["include/mlir/Dialect/X86/TransformOps/*.h"]),
includes = ["include"],
deps = [
- ":AffineDialect",
- ":ArithDialect",
":IR",
":LLVMCommonConversion",
":LLVMDialect",
- ":MemRefDialect",
":TransformDialect",
":TransformDialectInterfaces",
":VectorDialect",
@@ -2674,12 +2671,9 @@ cc_library(
hdrs = ["include/mlir/Dialect/X86/Transforms.h"],
includes = ["include"],
deps = [
- ":AffineDialect",
":ArithDialect",
":ControlFlowDialect",
":ConvertToLLVMInterface",
- ":DataLayoutInterfaces",
- ":DialectUtils",
":IR",
":LLVMCommonConversion",
":LLVMDialect",
@@ -2689,12 +2683,10 @@ cc_library(
":MemRefUtils",
":Pass",
":SCFDialect",
- ":SCFUtils",
":TransformUtils",
":UBDialect",
":VectorDialect",
":VectorTransforms",
- ":ValueBoundsOpInterface",
":VectorUtils",
":X86Dialect",
":X86Utils",
>From ca24a80efe9bd286e9b8e955cd252a1ec726064c Mon Sep 17 00:00:00 2001
From: fschlimb <frank.schlimbach at intel.com>
Date: Mon, 26 Oct 2026 04:48:02 -0700
Subject: [PATCH 4/4] negative_ prefix, remove cast-away-leading-unit-dims
---
.../X86/vector-contract-prepare-loops.mlir | 78 +++----------------
1 file changed, 10 insertions(+), 68 deletions(-)
diff --git a/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir b/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
index 865aba34cab350..029e503a6b592b 100644
--- a/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
@@ -38,34 +38,6 @@ func.func @direct_vnni_bf16(%a: memref<?x32x32xbf16>,
return %result : vector<32x32xf32>
}
-// fold VNNI unpacking of packed bf16 rhs read with leading unit dims
-// CHECK-LABEL: func.func @leading_unit_vnni(
-// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x32xbf16>, %[[B:[^,]+]]: memref<?x1x1x16x32x2xbf16>
-// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x2xbf16>
-// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
-// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A_PACKED]][%[[K]], {{.*}} vector<32x16x2xbf16>
-// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B]][%[[K]], {{.*}} vector<16x32x2xbf16>
-// CHECK: vector.contract {indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
-// NANO-LABEL: func.func @leading_unit_vnni(
-// NANO-NOT: vector.contract
-// NANO: return
-func.func @leading_unit_vnni(%a: memref<?x32x32xbf16>,
- %b: memref<?x1x1x16x32x2xbf16>,
- %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
- %zero = arith.constant 0 : index
- %step = arith.constant 16 : index
- %pad = arith.constant 0.0 : bf16
- %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
- %lhs = vector.transfer_read %a[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x32x32xbf16>, vector<32x32xbf16>
- %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero, %zero, %zero], %pad {in_bounds = [true, true, true, true, true]} : memref<?x1x1x16x32x2xbf16>, vector<1x1x16x32x2xbf16>
- %transpose = vector.transpose %rhs, [0, 1, 2, 4, 3] : vector<1x1x16x32x2xbf16> to vector<1x1x16x2x32xbf16>
- %flat = vector.shape_cast %transpose : vector<1x1x16x2x32xbf16> to vector<32x32xbf16>
- %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
- scf.yield %dot : vector<32x32xf32>
- }
- return %result : vector<32x32xf32>
-}
-
// fold VNNI unpacking of flat bf16 rhs read split into lanes by shape_cast
// CHECK-LABEL: func.func @split_vnni_bf16(
// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x32xbf16>, %[[B:[^,]+]]: memref<?x16x64xbf16>
@@ -157,34 +129,6 @@ func.func @split_vnni_int8(%a: memref<?x32x64xi8>,
return %result : vector<32x32xi32>
}
-// cast away leading unit dims of reads so that the contraction can be unrolled
-// CHECK-LABEL: func.func @unit_dimension_reads(
-// CHECK-SAME: %[[A:[^,]+]]: memref<?x?x32x32xbf16>, %[[B:[^,]+]]: memref<?x?x32x32xbf16>
-// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
-// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][%[[K]], {{.*}} vector<32x32xbf16>
-// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B]][%[[K]], {{.*}} vector<32x32xbf16>
-// CHECK: vector.contract {{.*}} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
-// NANO-LABEL: func.func @unit_dimension_reads(
-// NANO-NOT: vector.contract
-// NANO: return
-func.func @unit_dimension_reads(%a: memref<?x?x32x32xbf16>,
- %b: memref<?x?x32x32xbf16>,
- %init: vector<32x32xf32>, %ub: index,
- %m: index, %n: index) -> vector<32x32xf32> {
- %zero = arith.constant 0 : index
- %step = arith.constant 32 : index
- %pad = arith.constant 0.0 : bf16
- %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
- %lhs = vector.transfer_read %a[%k, %m, %zero, %zero], %pad {in_bounds = [true, true, true, true]} : memref<?x?x32x32xbf16>, vector<1x1x32x32xbf16>
- %rhs = vector.transfer_read %b[%k, %n, %zero, %zero], %pad {in_bounds = [true, true, true, true]} : memref<?x?x32x32xbf16>, vector<1x1x32x32xbf16>
- %left = vector.shape_cast %lhs : vector<1x1x32x32xbf16> to vector<32x32xbf16>
- %right = vector.shape_cast %rhs : vector<1x1x32x32xbf16> to vector<32x32xbf16>
- %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %right, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
- scf.yield %dot : vector<32x32xf32>
- }
- return %result : vector<32x32xf32>
-}
-
// fold VNNI unpacking through a static view of a dynamic lhs innermost dim
// CHECK-LABEL: func.func @dynamic_inner_k(
// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x?xbf16>
@@ -215,14 +159,14 @@ func.func @dynamic_inner_k(%a: memref<?x32x?xbf16>,
}
// negative test: lhs innermost read offset 1 is not a multiple of VNNI factor 2 -> skip
-// CHECK-LABEL: func.func @misaligned_inner_k(
+// CHECK-LABEL: func.func @negative_misaligned_inner_k(
// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x64xbf16>
// CHECK-NOT: memref.expand_shape
// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x32xbf16>
// CHECK: %[[T:.+]] = vector.transpose
// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x2x32xbf16> to vector<32x32xbf16>
// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
-func.func @misaligned_inner_k(%a: memref<?x32x64xbf16>,
+func.func @negative_misaligned_inner_k(%a: memref<?x32x64xbf16>,
%b: memref<?x16x32x2xbf16>,
%init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
%zero = arith.constant 0 : index
@@ -241,14 +185,14 @@ func.func @misaligned_inner_k(%a: memref<?x32x64xbf16>,
}
// negative test: lhs innermost dim 33 is not divisible by VNNI factor 2 -> skip
-// CHECK-LABEL: func.func @nondivisible_inner_k(
+// CHECK-LABEL: func.func @negative_nondivisible_inner_k(
// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x33xbf16>
// CHECK-NOT: memref.expand_shape
// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x32xbf16>
// CHECK: %[[T:.+]] = vector.transpose
// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x2x32xbf16> to vector<32x32xbf16>
// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
-func.func @nondivisible_inner_k(%a: memref<?x32x33xbf16>,
+func.func @negative_nondivisible_inner_k(%a: memref<?x32x33xbf16>,
%b: memref<?x16x32x2xbf16>,
%init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
%zero = arith.constant 0 : index
@@ -266,7 +210,7 @@ func.func @nondivisible_inner_k(%a: memref<?x32x33xbf16>,
}
// negative test: rhs VNNI lanes are reordered before unpacking -> skip
-// CHECK-LABEL: func.func @interleaved_lanes(
+// CHECK-LABEL: func.func @negative_interleaved_lanes(
// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x64xi8>
// CHECK-NOT: memref.expand_shape
// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x64xi8>
@@ -275,7 +219,7 @@ func.func @nondivisible_inner_k(%a: memref<?x32x33xbf16>,
// CHECK: %[[T:.+]] = vector.transpose %[[LANES]], [0, 2, 1]
// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x4x32xi8> to vector<64x32xi8>
// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<32x64xi8>, vector<64x32xi8> into vector<32x32xi32>
-func.func @interleaved_lanes(%a: memref<?x32x64xi8>,
+func.func @negative_interleaved_lanes(%a: memref<?x32x64xi8>,
%b: memref<?x16x32x4xi8>,
%init: vector<32x32xi32>, %ub: index) -> vector<32x32xi32> {
%zero = arith.constant 0 : index
@@ -295,14 +239,14 @@ func.func @interleaved_lanes(%a: memref<?x32x64xi8>,
}
// negative test: scalable lhs cannot be packed into a fixed VNNI shape -> skip
-// CHECK-LABEL: func.func @scalable_vnni(
+// CHECK-LABEL: func.func @negative_scalable_vnni(
// CHECK-SAME: %[[A:[^,]+]]: memref<?x?x32xbf16>
// CHECK-NOT: memref.expand_shape
// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<[32]x32xbf16>
// CHECK: %[[T:.+]] = vector.transpose
// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x2x32xbf16> to vector<32x32xbf16>
// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<[32]x32xbf16>, vector<32x32xbf16> into vector<[32]x32xf32>
-func.func @scalable_vnni(%a: memref<?x?x32xbf16>,
+func.func @negative_scalable_vnni(%a: memref<?x?x32xbf16>,
%b: memref<?x16x32x2xbf16>,
%init: vector<[32]x32xf32>, %ub: index,
%m: index) -> vector<[32]x32xf32> {
@@ -321,14 +265,14 @@ func.func @scalable_vnni(%a: memref<?x?x32xbf16>,
}
// negative test: lhs memref with vector elements cannot be expanded into VNNI layout -> skip
-// CHECK-LABEL: func.func @vector_element_lhs(
+// CHECK-LABEL: func.func @negative_vector_element_lhs(
// CHECK-SAME: %[[A:[^,]+]]: memref<?x32xvector<32xbf16>>
// CHECK-NOT: memref.expand_shape
// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x32xbf16>
// CHECK: %[[T:.+]] = vector.transpose
// CHECK: %[[FLAT:.+]] = vector.shape_cast %[[T]] : vector<16x2x32xbf16> to vector<32x32xbf16>
// CHECK: vector.contract {{.*}} %[[LHS]], %[[FLAT]], %{{.+}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
-func.func @vector_element_lhs(%a: memref<?x32xvector<32xbf16>>,
+func.func @negative_vector_element_lhs(%a: memref<?x32xvector<32xbf16>>,
%b: memref<?x16x32x2xbf16>,
%init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
%zero = arith.constant 0 : index
@@ -352,7 +296,6 @@ module attributes {transform.with_named_sequence} {
%loops = transform.structured.match ops{["scf.for"]} in %root : (!transform.any_op) -> !transform.any_op
transform.apply_licm to %loops : !transform.any_op
transform.apply_patterns to %funcs {
- transform.apply_patterns.vector.cast_away_vector_leading_one_dim
transform.apply_patterns.x86.prepare_vector_contract_loops
} : !transform.any_op
transform.apply_dce to %funcs : !transform.any_op
@@ -365,7 +308,6 @@ module attributes {transform.with_named_sequence} {
%loops = transform.structured.match ops{["scf.for"]} in %root : (!transform.any_op) -> !transform.any_op
transform.apply_licm to %loops : !transform.any_op
transform.apply_patterns to %funcs {
- transform.apply_patterns.vector.cast_away_vector_leading_one_dim
transform.apply_patterns.x86.prepare_vector_contract_loops
} : !transform.any_op
transform.apply_dce to %funcs : !transform.any_op
More information about the Mlir-commits
mailing list