[Mlir-commits] [mlir] [mlir][X86] Prepare contract interleave (PR #236750)
Frank Schlimbach
llvmlistbot at llvm.org
Thu Oct 29 05:18:02 PDT 2026
https://github.com/fschlimb created https://github.com/llvm/llvm-project/pull/236750
Follow-up to #232800, which folds transpose + shape_cast unpacking of a
VNNI-packed RHS. Some frontends instead unpack by splitting the source
into VNNI lanes, transposing each lane, and joining the lanes back.
triton-cpu does this with deinterleave/interleave, and cutile-cpu with
transpose + extract and insert_strided_slice. Neither form was matched,
so the contraction never reached AMX lowering.
This matches both forms. For i8, factor 4 can be expressed as two binary
levels or a single 4-way step. We only fold when every lane comes from
the same source in VNNI order.
Also simplifying checks and hardening integration tests.
Assisted by AI.
>From 862c3b4f9e4e18826c72ff5c789188d8214ade32 Mon Sep 17 00:00:00 2001
From: fschlimb <frank.schlimbach at intel.com>
Date: Tue, 27 Oct 2026 07:24:31 -0700
Subject: [PATCH 1/2] [mlir][X86] Fold split-and-join VNNI unpacking into
contractions
Recognize bf16 even/odd split-and-join unpacking, including
transpose/extract and slice-insert forms.
Add regression and execution tests for supported forms and safety checks.
Co-authored-by: Copilot
---
.../X86/TransformOps/X86TransformOps.td | 8 +-
mlir/include/mlir/Dialect/X86/Transforms.h | 4 +-
.../Transforms/PrepareVectorContractLoops.cpp | 265 +++++++++++----
.../X86/vector-contract-prepare-loops.mlir | 311 +++++++++++++++++-
.../CPU/X86/prepare-contract-loops.mlir | 75 +++++
5 files changed, 588 insertions(+), 75 deletions(-)
diff --git a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
index b5859f6d27adea..19ee3d71957d54 100644
--- a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
+++ b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
@@ -97,7 +97,13 @@ 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.
+ Apply patterns that rewrite bf16 and i8 matmul loops to contract packed
+ VNNI operands directly, avoiding RHS unpacking. Supported unpacking uses
+ shape casts and transposes, or bf16 even/odd split-and-join sequences.
+
+ Before applying these patterns, remove leading unit dimensions from
+ vectors with
+ `apply_patterns.vector.cast_away_vector_leading_one_dim`.
}];
let assemblyFormat = "attr-dict";
diff --git a/mlir/include/mlir/Dialect/X86/Transforms.h b/mlir/include/mlir/Dialect/X86/Transforms.h
index 6d28cf5b081d6b..60464b72e9a56c 100644
--- a/mlir/include/mlir/Dialect/X86/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86/Transforms.h
@@ -113,8 +113,8 @@ void populateVectorContractToAMXDotProductPatterns(RewritePatternSet &patterns);
void populateMoveAccumulatorForContractLoopPatterns(
RewritePatternSet &patterns);
-// Adds patterns to prepare vector contractions inside scf.for loops for
-// subsequent lowering (e.g. multi-level-unroll).
+// Adds patterns to fold RHS unpacking into VNNI vector contractions inside
+// scf.for loops for subsequent lowering (e.g. multi-level-unroll).
void populatePrepareVectorContractLoopsPatterns(RewritePatternSet &patterns);
// A set of patterns that tile and unroll contractions in order to prepare them
diff --git a/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp b/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
index f5eec9fb270895..1097db2fa14555 100644
--- a/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
@@ -12,8 +12,8 @@
#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/ADT/SmallVectorExtras.h"
#include "llvm/Support/DebugLog.h"
#define DEBUG_TYPE "x86-prepare-vector-contract-loops"
@@ -34,7 +34,6 @@ static vector::ContractionOp getCandidate(scf::ForOp loop) {
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()) {
LDBG() << "reject contraction at " << contract.getLoc()
<< ": not a single loop-carried accumulation";
@@ -54,20 +53,201 @@ static vector::ContractionOp getCandidate(scf::ForOp loop) {
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;
LDBG() << "reject contraction at " << contract.getLoc()
<< ": noncanonical matmul layout or iterators";
}
return {};
}
+// Look through reshapes/broadcasts without changing the row-major element
+// order. A broadcast is safe to skip if it adds unit dimensions only.
+static Value skipReshapes(Value v) {
+ while (isa_and_present<vector::ShapeCastOp, vector::BroadcastOp>(
+ v.getDefiningOp())) {
+ Value source = v.getDefiningOp()->getOperand(0);
+ auto sourceType = dyn_cast<VectorType>(source.getType());
+ if (!sourceType || sourceType.getNumElements() !=
+ cast<VectorType>(v.getType()).getNumElements())
+ break;
+ v = source;
+ }
+ return v;
+}
+
+// Find the source and parity of an even/odd split: v[i] = src[2 * i + j].
+// Match either pattern below (unit dimensions may also be present):
+//
+// %even, %odd = vector.deinterleave %src
+// : vector<4x8xbf16> -> vector<4x4xbf16>
+//
+// or:
+//
+// %packed = vector.shape_cast %src
+// : vector<4x8xbf16> to vector<4x4x2xbf16>
+// %split = vector.transpose %packed, [2, 0, 1]
+// : vector<4x4x2xbf16> to vector<2x4x4xbf16>
+// %even = vector.extract %split[0]
+// : vector<4x4xbf16> from vector<2x4x4xbf16>
+// %odd = vector.extract %split[1]
+// : vector<4x4xbf16> from vector<2x4x4xbf16>
+static std::optional<std::pair<Value, int64_t>> matchEvenOrOdd(Value v) {
+ v = skipReshapes(v);
+ if (auto deinterleave = v.getDefiningOp<vector::DeinterleaveOp>())
+ return std::pair(skipReshapes(deinterleave.getSource()),
+ int64_t(cast<OpResult>(v).getResultNumber()));
+
+ auto extract = v.getDefiningOp<vector::ExtractOp>();
+ auto transpose =
+ extract ? extract.getSource().getDefiningOp<vector::TransposeOp>()
+ : vector::TransposeOp();
+ auto type = dyn_cast<VectorType>(v.getType());
+ if (!transpose || !type || extract.hasDynamicPosition() ||
+ transpose.getSourceVectorType().isScalable())
+ return std::nullopt;
+
+ ArrayRef<int64_t> shape = transpose.getSourceVectorType().getShape();
+ auto isNonUnit = [&](int64_t dim) { return shape[dim] != 1; };
+ SmallVector<int64_t> order =
+ llvm::filter_to_vector(llvm::seq<int64_t>(0, shape.size()), isNonUnit);
+ if (order.empty() || shape[order.back()] != 2)
+ return std::nullopt;
+
+ std::rotate(order.begin(), std::prev(order.end()), order.end());
+ ArrayRef<int64_t> perm = transpose.getPermutation();
+ if (llvm::filter_to_vector(perm, isNonUnit) != order)
+ return std::nullopt;
+
+ // The dimension of size 2 is now first among the non-unit dimensions.
+ // Keeping exactly half the elements means the extract indexes that
+ // dimension; any other dimensions it indexes must be unit dimensions.
+ if (2 * type.getNumElements() !=
+ transpose.getSourceVectorType().getNumElements())
+ return std::nullopt;
+
+ int64_t j =
+ extract
+ .getStaticPosition()[llvm::find(perm, order.front()) - perm.begin()];
+ if (j != 0 && j != 1)
+ return std::nullopt;
+
+ return std::pair(skipReshapes(transpose.getVector()), j);
+}
+
+// Find the even and odd inputs of either pattern below. The slice inserts may
+// appear in either order.
+//
+// %joined = vector.interleave %even, %odd
+// : vector<4x1xbf16> -> vector<4x2xbf16>
+//
+// or:
+//
+// %poison = ub.poison : vector<4x2xbf16>
+// %partial = vector.insert_strided_slice %odd, %poison
+// offsets = [0, 1], strides = [1, 1]
+// : vector<4x1xbf16> into vector<4x2xbf16>
+// %joined = vector.insert_strided_slice %even, %partial
+// offsets = [0, 0], strides = [1, 1]
+// : vector<4x1xbf16> into vector<4x2xbf16>
+static std::optional<std::array<Value, 2>> matchInterleave(Value v) {
+ v = skipReshapes(v);
+ if (auto interleave = v.getDefiningOp<vector::InterleaveOp>())
+ return std::array<Value, 2>{interleave.getLhs(), interleave.getRhs()};
+
+ // Both slices must be complete and disjoint so the original destination
+ // contributes no elements to the result.
+ std::array<Value, 2> halves;
+ for (int i = 0; i < 2; ++i) {
+ auto insert = v.getDefiningOp<vector::InsertStridedSliceOp>();
+ if (!insert)
+ return std::nullopt;
+
+ ArrayRef<int64_t> shape = insert.getDestVectorType().getShape();
+ SmallVector<int64_t> slice(shape.drop_back());
+ slice.push_back(1);
+ auto offsets = extractFromIntegerArrayAttr<int64_t>(insert.getOffsets());
+ int64_t j = offsets.back();
+ if (shape.back() != 2 ||
+ insert.getSourceVectorType().getShape() != ArrayRef(slice) || halves[j])
+ return std::nullopt;
+
+ halves[j] = insert.getValueToStore();
+ v = insert.getDest();
+ }
+ return halves;
+}
+
+// Find the packed RHS so the contraction can use it without unpacking.
+// For a K x N RHS and VNNI factor f, require
+// rhs[f * r + j][c] = src[r][c][j], with src viewed as K/f x N x f.
+// Match transpose/shape_cast unpacking (shown for f = 2):
+//
+// %packed = vector.shape_cast %src
+// : vector<4x6xbf16> to vector<4x3x2xbf16>
+// %transposed = vector.transpose %packed, [0, 2, 1]
+// : vector<4x3x2xbf16> to vector<4x2x3xbf16>
+// %rhs = vector.shape_cast %transposed
+// : vector<4x2x3xbf16> to vector<8x3xbf16>
+//
+// or, for f = 2, split/transpose/join/transpose unpacking:
+//
+// %even, %odd = vector.deinterleave %src
+// : vector<4x6xbf16> -> vector<4x3xbf16>
+// %even_t = vector.transpose %even, [1, 0]
+// : vector<4x3xbf16> to vector<3x4xbf16>
+// %odd_t = vector.transpose %odd, [1, 0]
+// : vector<4x3xbf16> to vector<3x4xbf16>
+// %joined = vector.interleave %even_t, %odd_t
+// : vector<3x4xbf16> -> vector<3x8xbf16>
+// %rhs = vector.transpose %joined, [1, 0]
+// : vector<3x8xbf16> to vector<8x3xbf16>
+static Value matchUnpackedRhs(Value rhs, int64_t factor) {
+ auto type = cast<VectorType>(rhs.getType());
+ int64_t k = type.getDimSize(0) / factor, n = type.getDimSize(1);
+ auto transpose = skipReshapes(rhs).getDefiningOp<vector::TransposeOp>();
+ if (!transpose)
+ return {};
+
+ ArrayRef<int64_t> shape = transpose.getSourceVectorType().getShape();
+ ArrayRef<int64_t> perm = transpose.getPermutation();
+ if (shape.size() >= 3) {
+ size_t leading = shape.size() - 3;
+ auto expected = llvm::to_vector(llvm::seq<int64_t>(0, shape.size()));
+ std::swap(expected[leading + 1], expected[leading + 2]);
+ bool isUnpacking =
+ llvm::all_of(shape.take_front(leading),
+ [](int64_t dim) { return dim == 1; }) &&
+ shape.drop_front(leading) == ArrayRef<int64_t>{k, n, factor} &&
+ perm == ArrayRef(expected);
+ return isUnpacking ? skipReshapes(transpose.getVector()) : Value();
+ }
+
+ const int64_t swap[] = {1, 0};
+ if (factor != 2 || perm != ArrayRef(swap) ||
+ shape != ArrayRef<int64_t>{n, 2 * k})
+ return {};
+
+ std::optional<std::array<Value, 2>> halves =
+ matchInterleave(transpose.getVector());
+ if (!halves)
+ return {};
+
+ Value src;
+ for (auto [j, half] : llvm::enumerate(*halves)) {
+ auto transpose = skipReshapes(half).getDefiningOp<vector::TransposeOp>();
+ if (!transpose || transpose.getPermutation() != ArrayRef(swap) ||
+ transpose.getSourceVectorType().getShape() != ArrayRef<int64_t>{k, n})
+ return {};
+
+ auto evenOrOdd = matchEvenOrOdd(transpose.getVector());
+ if (!evenOrOdd || evenOrOdd->second != int64_t(j) ||
+ (src && evenOrOdd->first != src))
+ return {};
+
+ src = evenOrOdd->first;
+ }
+ return src;
+}
+
static bool isSimpleReadOutsideLoop(vector::TransferReadOp read,
scf::ForOp loop) {
return read && read->getParentOp() == loop &&
@@ -171,10 +351,10 @@ static Value buildExpandedRead(PatternRewriter &rewriter, Location loc,
namespace {
// 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.
+// - Match a packed RHS read unpacked for a flat matmul.
+// - Split the LHS memref's K dimension and reread it in VNNI layout.
// 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.
+// - Contract the packed reads directly over K and the VNNI dimension.
// For bf16, 32x32 LHS and RHS become 32x16x2 LHS and 16x32x2 RHS.
struct FoldVNNIUnpacking : OpRewritePattern<vector::ContractionOp> {
using OpRewritePattern::OpRewritePattern;
@@ -191,71 +371,28 @@ struct FoldVNNIUnpacking : OpRewritePattern<vector::ContractionOp> {
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");
+ Value packedSource = matchUnpackedRhs(contract.getRhs(), factor);
+ if (!packedSource)
+ return rewriter.notifyMatchFailure(contract,
+ "RHS does not unpack a packed vector");
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>();
+ auto rhsRead = packedSource.getDefiningOp<vector::TransferReadOp>();
if (!lhs)
return rewriter.notifyMatchFailure(
contract, "LHS read cannot be expanded into VNNI layout");
@@ -263,7 +400,7 @@ struct FoldVNNIUnpacking : OpRewritePattern<vector::ContractionOp> {
return rewriter.notifyMatchFailure(contract,
"RHS is not a simple transfer read");
std::optional<ExpandedRead> splitRhs;
- if (packedCast) {
+ if (rhsRead.getVectorType() != rhsPackedType) {
auto flatType = rhsRead.getVectorType();
if (flatType.getRank() != 2 ||
flatType.getDimSize(0) != rhsPackedType.getDimSize(0) ||
@@ -275,9 +412,6 @@ struct FoldVNNIUnpacking : OpRewritePattern<vector::ContractionOp> {
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");
@@ -310,6 +444,7 @@ struct FoldVNNIUnpacking : OpRewritePattern<vector::ContractionOp> {
rewriter.replaceOpWithNewOp<vector::ContractionOp>(
contract, packedLhs, packedRhs, contract.getAcc(), vnniMaps, iterators,
contract.getKind(), contract.getFastmath());
+ LDBG() << "Replaced contract with VNNI-compatible contraction";
return success();
}
};
diff --git a/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir b/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
index 3b9682575011db..2bcac9a35caa78 100644
--- a/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
@@ -14,6 +14,9 @@
// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x32xbf16>, %[[B:[^,]+]]: memref<?x16x32x2xbf16>
// NANO-LABEL: func.func @direct_vnni_bf16(
// NANO-NOT: vector.contract
+// NANO-COUNT-4: x86.amx.tile_mulf
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_mulf
// NANO: return
// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x2xbf16>
// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
@@ -38,11 +41,14 @@ func.func @direct_vnni_bf16(%a: memref<?x32x32xbf16>,
return %result : vector<32x32xf32>
}
-// fold VNNI unpacking of flat bf16 rhs read split into lanes by shape_cast
+// fold VNNI unpacking of flat bf16 rhs read reshaped into VNNI layout
// 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-COUNT-4: x86.amx.tile_mulf
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_mulf
// NANO: return
// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x2xbf16>
// CHECK: %[[B_PACKED:.+]] = memref.expand_shape %[[B]] {{.*}} into memref<?x16x32x2xbf16>
@@ -69,11 +75,293 @@ func.func @split_vnni_bf16(%a: memref<?x32x32xbf16>,
return %result : vector<32x32xf32>
}
+// Fold deinterleave/transpose/interleave/transpose unpacking with strided
+// memrefs and an i32 loop index.
+// CHECK-LABEL: func.func @interleaved_vnni_bf16(
+// CHECK-SAME: %[[A:[^:]+]]: memref<?x?x32x32xbf16, strided<[?, 1024, 32, 1]>>, %[[B:[^:]+]]: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>
+// NANO-LABEL: func.func @interleaved_vnni_bf16(
+// NANO: vector.contract {{.*}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+// NANO: return
+// CHECK-DAG: %[[AP:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x?x32x16x2xbf16, strided<[?, 1024, 32, 2, 1]>>
+// CHECK-DAG: %[[BP:.+]] = memref.expand_shape %[[B]] {{.*}} into memref<?x?x16x32x2xbf16, strided<[?, 1024, 64, 2, 1]>>
+// CHECK: scf.for {{.*}} : i32 {
+// CHECK: %[[K:.+]] = arith.index_cast {{.*}} : i32 to index
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[AP]][%{{.+}}, %[[K]], {{.*}} vector<32x16x2xbf16>
+// CHECK-NEXT: %[[RHS:.+]] = vector.transfer_read %[[BP]][%{{.+}}, %[[K]], {{.*}} vector<16x32x2xbf16>
+// CHECK-NEXT: %[[DOT:.+]] = 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>
+// CHECK-NEXT: scf.yield %[[DOT]]
+func.func @interleaved_vnni_bf16(%a: memref<?x?x32x32xbf16, strided<[?, 1024, 32, 1]>>,
+ %b: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>,
+ %init: vector<32x32xf32>, %lb: i32, %ub: i32,
+ %m: index, %n: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : i32
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %iv = %lb to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) : i32 {
+ %k = arith.index_cast %iv : i32 to index
+ %lhs = vector.transfer_read %a[%m, %k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x?x32x32xbf16, strided<[?, 1024, 32, 1]>>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%n, %k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>, vector<16x64xbf16>
+ %even, %odd = vector.deinterleave %rhs : vector<16x64xbf16> -> vector<16x32xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %interleaved = vector.interleave %even_t, %odd_t : vector<32x16xbf16> -> vector<32x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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>
+}
+
+// Also handle non-square operands and aligned nonzero inner read offsets.
+// CHECK-LABEL: func.func @interleaved_vnni_rectangular(
+// CHECK: %[[ONE:.+]] = arith.constant 1 : index
+// CHECK: memref.expand_shape {{.*}} into memref<?x16x32x2xbf16>
+// CHECK: memref.expand_shape {{.*}} into memref<?x16x32x2xbf16>
+// CHECK: scf.for %[[K:[^ ]+]] =
+// CHECK: vector.transfer_read %{{.+}}[%[[K]], %{{.+}}, %[[ONE]], %{{.+}}]{{.*}} vector<16x16x2xbf16>
+// CHECK: vector.transfer_read %{{.+}}[%[[K]], %{{.+}}, %[[ONE]], %{{.+}}]{{.*}} vector<16x16x2xbf16>
+// CHECK-NOT: vector.deinterleave
+// CHECK-NOT: vector.interleave
+// CHECK-NOT: vector.transpose
+// CHECK: vector.contract {{.*}} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+func.func @interleaved_vnni_rectangular(%a: memref<?x16x64xbf16>,
+ %b: memref<?x16x64xbf16>,
+ %init: vector<16x16xf32>,
+ %ub: index) -> vector<16x16xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %offset = arith.constant 2 : index
+ %pad = arith.constant 0.0 : bf16
+ %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<16x16xf32>) {
+ %lhs = vector.transfer_read %a[%k, %zero, %offset], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x32xbf16>
+ %rhs = vector.transfer_read %b[%k, %zero, %offset], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x32xbf16>
+ %even, %odd = vector.deinterleave %rhs : vector<16x32xbf16> -> vector<16x16xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x16xbf16> to vector<16x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x16xbf16> to vector<16x16xbf16>
+ %interleaved = vector.interleave %even_t, %odd_t : vector<16x16xbf16> -> vector<16x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+ scf.yield %dot : vector<16x16xf32>
+ }
+ return %result : vector<16x16xf32>
+}
+
+// Swapping the even and odd elements changes the packed element order.
+// CHECK-LABEL: func.func @negative_interleaved_swapped(
+// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @negative_interleaved_swapped(%a: memref<?x32x32xbf16>, %b: memref<?x16x64xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0.0 : bf16
+ %r0 = 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], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x64xbf16>
+ %even, %odd = vector.deinterleave %rhs : vector<16x64xbf16> -> vector<16x32xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %interleaved = vector.interleave %odd_t, %even_t : vector<32x16xbf16> -> vector<32x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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 %r0 : vector<32x32xf32>
+}
+
+// The even and odd elements must come from the same packed source.
+// CHECK-LABEL: func.func @negative_interleaved_different_sources(
+// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @negative_interleaved_different_sources(%a: memref<?x32x32xbf16>, %b: memref<?x16x64xbf16>,
+ %c: memref<?x16x64xbf16>, %init: vector<32x32xf32>,
+ %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0.0 : bf16
+ %r1 = 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], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x64xbf16>
+ %rhs_c = vector.transfer_read %c[%k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x64xbf16>
+ %even, %unused = vector.deinterleave %rhs : vector<16x64xbf16> -> vector<16x32xbf16>
+ %unused_c, %odd = vector.deinterleave %rhs_c : vector<16x64xbf16> -> vector<16x32xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %interleaved = vector.interleave %even_t, %odd_t : vector<32x16xbf16> -> vector<32x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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 %r1 : vector<32x32xf32>
+}
+
+// A loop-invariant RHS does not provide a packed read to fold.
+// CHECK-LABEL: func.func @negative_interleaved_invariant(
+// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @negative_interleaved_invariant(%a: memref<?x32x32xbf16>, %invariant: vector<16x64xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0.0 : bf16
+ %r2 = 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>
+ %even, %odd = vector.deinterleave %invariant : vector<16x64xbf16> -> vector<16x32xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %interleaved = vector.interleave %even_t, %odd_t : vector<32x16xbf16> -> vector<32x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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 %r2 : vector<32x32xf32>
+}
+
+// Fold transpose/extract and slice-insert unpacking despite interspersed unit
+// dimensions and odd-first insertion.
+// CHECK-LABEL: func.func @split_join_bf16(
+// CHECK-DAG: %[[AP:.+]] = memref.expand_shape {{.*}} into memref<?x?x?x16x2xbf16, strided<[?, ?, ?, 2, 1], offset: ?>>
+// CHECK-DAG: %[[BP:.+]] = memref.expand_shape {{.*}} into memref<?x?x?x32x2xbf16, strided<[?, ?, ?, 2, 1], offset: ?>>
+// CHECK: scf.for %[[IV:[^ ]+]] =
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[AP]][%{{.+}}, %[[IV]], {{.*}} vector<32x16x2xbf16>
+// CHECK-NEXT: %[[RHS:.+]] = vector.transfer_read %[[BP]][%{{.+}}, %[[IV]], {{.*}} vector<16x32x2xbf16>
+// CHECK-NEXT: vector.contract {{.*}} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+// CHECK-NEXT: scf.yield
+// NANO-LABEL: func.func @split_join_bf16(
+// NANO-NOT: vector.contract
+// NANO-COUNT-4: x86.amx.tile_mulf
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_mulf
+// NANO: return
+func.func @split_join_bf16(%a: memref<?x?x?x?xbf16, strided<[?, ?, ?, 1], offset: ?>>,
+ %b: memref<?x?x?x?xbf16, strided<[?, ?, ?, 1], offset: ?>>,
+ %ub: index, %m: index, %n: index) -> vector<32x32xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %poison = ub.poison : vector<32x16x2xbf16>
+ %pad = arith.constant 0.0 : bf16
+ %init = arith.constant dense<0.0> : vector<32x32xf32>
+ %lb = arith.constant 0 : index
+ %result = scf.for %iv = %lb to %ub step %c1 iter_args(%accum = %init) -> (vector<32x32xf32>) {
+ %lhs = vector.transfer_read %a[%m, %iv, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x?xbf16, strided<[?, ?, ?, 1], offset: ?>>, vector<32x32xbf16>
+ %rhs = vector.transfer_read %b[%n, %iv, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x?xbf16, strided<[?, ?, ?, 1], offset: ?>>, vector<16x64xbf16>
+ %rhs4 = vector.broadcast %rhs : vector<16x64xbf16> to vector<1x1x16x64xbf16>
+ %packed = vector.shape_cast %rhs4 : vector<1x1x16x64xbf16> to vector<1x16x1x32x2x1xbf16>
+ %even_odd = vector.transpose %packed, [0, 2, 4, 1, 3, 5] : vector<1x16x1x32x2x1xbf16> to vector<1x1x2x16x32x1xbf16>
+ %even3 = vector.extract %even_odd[0, 0, 0] : vector<16x32x1xbf16> from vector<1x1x2x16x32x1xbf16>
+ %odd3 = vector.extract %even_odd[0, 0, 1] : vector<16x32x1xbf16> from vector<1x1x2x16x32x1xbf16>
+ %even = vector.shape_cast %even3 : vector<16x32x1xbf16> to vector<16x32xbf16>
+ %odd = vector.shape_cast %odd3 : vector<16x32x1xbf16> to vector<16x32xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %even_t3 = vector.shape_cast %even_t : vector<32x16xbf16> to vector<32x16x1xbf16>
+ %odd_t3 = vector.shape_cast %odd_t : vector<32x16xbf16> to vector<32x16x1xbf16>
+ %join0 = vector.insert_strided_slice %odd_t3, %poison offsets = [0, 0, 1], strides = [1, 1, 1] : vector<32x16x1xbf16> into vector<32x16x2xbf16>
+ %join = vector.insert_strided_slice %even_t3, %join0 offsets = [0, 0, 0], strides = [1, 1, 1] : vector<32x16x1xbf16> into vector<32x16x2xbf16>
+ %interleaved = vector.shape_cast %join : vector<32x16x2xbf16> to vector<32x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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>
+}
+
+// Reordering dimensions other than the VNNI dimension changes element order.
+// CHECK-LABEL: func.func @negative_split_join_reordered(
+// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @negative_split_join_reordered(%a: memref<?x32x32xbf16>, %b: memref<?x16x64xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0.0 : bf16
+ %poison = ub.poison : vector<32x16x2xbf16>
+ %r0 = 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], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x64xbf16>
+ %packed = vector.shape_cast %rhs : vector<16x64xbf16> to vector<16x32x2xbf16>
+ %even_odd = vector.transpose %packed, [2, 1, 0] : vector<16x32x2xbf16> to vector<2x32x16xbf16>
+ %even_r = vector.extract %even_odd[0] : vector<32x16xbf16> from vector<2x32x16xbf16>
+ %odd_r = vector.extract %even_odd[1] : vector<32x16xbf16> from vector<2x32x16xbf16>
+ %even = vector.shape_cast %even_r : vector<32x16xbf16> to vector<16x32xbf16>
+ %odd = vector.shape_cast %odd_r : vector<32x16xbf16> to vector<16x32xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %even3 = vector.shape_cast %even_t : vector<32x16xbf16> to vector<32x16x1xbf16>
+ %odd3 = vector.shape_cast %odd_t : vector<32x16xbf16> to vector<32x16x1xbf16>
+ %join0 = vector.insert_strided_slice %even3, %poison offsets = [0, 0, 0], strides = [1, 1, 1] : vector<32x16x1xbf16> into vector<32x16x2xbf16>
+ %join = vector.insert_strided_slice %odd3, %join0 offsets = [0, 0, 1], strides = [1, 1, 1] : vector<32x16x1xbf16> into vector<32x16x2xbf16>
+ %interleaved = vector.shape_cast %join : vector<32x16x2xbf16> to vector<32x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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 %r0 : vector<32x32xf32>
+}
+
+// A single insert leaves the odd elements unspecified.
+// CHECK-LABEL: func.func @negative_split_join_incomplete(
+// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @negative_split_join_incomplete(%a: memref<?x32x32xbf16>, %b: memref<?x16x64xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0.0 : bf16
+ %poison = ub.poison : vector<32x16x2xbf16>
+ %r1 = 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], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x64xbf16>
+ %packed = vector.shape_cast %rhs : vector<16x64xbf16> to vector<16x32x2xbf16>
+ %even_odd = vector.transpose %packed, [2, 0, 1] : vector<16x32x2xbf16> to vector<2x16x32xbf16>
+ %even = vector.extract %even_odd[0] : vector<16x32xbf16> from vector<2x16x32xbf16>
+ %odd = vector.extract %even_odd[1] : vector<16x32xbf16> from vector<2x16x32xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %even3 = vector.shape_cast %even_t : vector<32x16xbf16> to vector<32x16x1xbf16>
+ %odd3 = vector.shape_cast %odd_t : vector<32x16xbf16> to vector<32x16x1xbf16>
+ %join = vector.insert_strided_slice %even3, %poison offsets = [0, 0, 0], strides = [1, 1, 1] : vector<32x16x1xbf16> into vector<32x16x2xbf16>
+ %interleaved = vector.shape_cast %join : vector<32x16x2xbf16> to vector<32x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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 %r1 : vector<32x32xf32>
+}
+
+// Two inserts at the same offset still leave the odd positions unspecified.
+// CHECK-LABEL: func.func @negative_split_join_duplicate(
+// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+func.func @negative_split_join_duplicate(%a: memref<?x32x32xbf16>, %b: memref<?x16x64xbf16>,
+ %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0.0 : bf16
+ %poison = ub.poison : vector<32x16x2xbf16>
+ %r2 = 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], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x64xbf16>
+ %packed = vector.shape_cast %rhs : vector<16x64xbf16> to vector<16x32x2xbf16>
+ %even_odd = vector.transpose %packed, [2, 0, 1] : vector<16x32x2xbf16> to vector<2x16x32xbf16>
+ %even = vector.extract %even_odd[0] : vector<16x32xbf16> from vector<2x16x32xbf16>
+ %odd = vector.extract %even_odd[1] : vector<16x32xbf16> from vector<2x16x32xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+ %even3 = vector.shape_cast %even_t : vector<32x16xbf16> to vector<32x16x1xbf16>
+ %odd3 = vector.shape_cast %odd_t : vector<32x16xbf16> to vector<32x16x1xbf16>
+ %join0 = vector.insert_strided_slice %even3, %poison offsets = [0, 0, 0], strides = [1, 1, 1] : vector<32x16x1xbf16> into vector<32x16x2xbf16>
+ %join = vector.insert_strided_slice %odd3, %join0 offsets = [0, 0, 0], strides = [1, 1, 1] : vector<32x16x1xbf16> into vector<32x16x2xbf16>
+ %interleaved = vector.shape_cast %join : vector<32x16x2xbf16> to vector<32x32xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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 %r2 : 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-COUNT-4: x86.amx.tile_muli
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_muli
// NANO: return
// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x4xi8>
// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
@@ -98,11 +386,14 @@ func.func @direct_vnni_int8(%a: memref<?x32x64xi8>,
return %result : vector<32x32xi32>
}
-// fold VNNI unpacking of flat i8 rhs read split into lanes by shape_cast
+// fold VNNI unpacking of flat i8 rhs read reshaped into VNNI layout
// 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-COUNT-4: x86.amx.tile_muli
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_muli
// NANO: return
// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x4xi8>
// CHECK: %[[B_PACKED:.+]] = memref.expand_shape %[[B]] {{.*}} into memref<?x16x32x4xi8>
@@ -134,6 +425,9 @@ func.func @split_vnni_int8(%a: memref<?x32x64xi8>,
// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x?xbf16>
// NANO-LABEL: func.func @dynamic_inner_k(
// NANO-NOT: vector.contract
+// NANO-COUNT-4: x86.amx.tile_mulf
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_mulf
// 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]>>
@@ -169,6 +463,9 @@ func.func @dynamic_inner_k(%a: memref<?x32x?xbf16>,
// CHECK: vector.contract {{.*}} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
// NANO-LABEL: func.func @broadcast_vnni_bf16(
// NANO-NOT: vector.contract
+// NANO-COUNT-4: x86.amx.tile_mulf
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_mulf
// NANO: return
func.func @broadcast_vnni_bf16(%a: memref<?x32x32xbf16>,
%b: memref<?x16x32x2xbf16>,
@@ -265,17 +562,17 @@ func.func @negative_nondivisible_inner_k(%a: memref<?x32x33xbf16>,
return %result : vector<32x32xf32>
}
-// negative test: rhs VNNI lanes are reordered before unpacking -> skip
-// CHECK-LABEL: func.func @negative_interleaved_lanes(
+// negative test: even and odd rhs elements are swapped before unpacking -> skip
+// CHECK-LABEL: func.func @negative_interleaved_elements(
// 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: %[[REORDERED:.+]] = vector.interleave %[[ODD]], %[[EVEN]]
+// CHECK: %[[T:.+]] = vector.transpose %[[REORDERED]], [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 @negative_interleaved_lanes(%a: memref<?x32x64xi8>,
+func.func @negative_interleaved_elements(%a: memref<?x32x64xi8>,
%b: memref<?x16x32x4xi8>,
%init: vector<32x32xi32>, %ub: index) -> vector<32x32xi32> {
%zero = arith.constant 0 : index
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
index 20c9022cd9037d..a2c73666a282e9 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/X86/prepare-contract-loops.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/X86/prepare-contract-loops.mlir
@@ -47,6 +47,62 @@ func.func @vnni_i8(%a: memref<1x2x8xi8>,
return %result : vector<2x4xi8>
}
+// Check that folding deinterleave/transpose/interleave/transpose unpacking
+// and converting the i32 loop to index preserve the matmul result.
+func.func @interleaved_bf16(%a: memref<1x2x8xbf16>,
+ %b: memref<1x4x6xbf16>) -> vector<2x3xbf16> {
+ %zero = arith.constant 0 : index
+ %lb = arith.constant 0 : i32
+ %ub = arith.constant 1 : i32
+ %pad = arith.constant 0.0 : bf16
+ %init = arith.constant dense<0.0> : vector<2x3xbf16>
+ %result = scf.for %iv = %lb to %ub step %ub iter_args(%accum = %init) -> (vector<2x3xbf16>) : i32 {
+ %k = arith.index_cast %iv : i32 to index
+ %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], %pad {in_bounds = [true, true]} : memref<1x4x6xbf16>, vector<4x6xbf16>
+ %even, %odd = vector.deinterleave %right : vector<4x6xbf16> -> vector<4x3xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<4x3xbf16> to vector<3x4xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<4x3xbf16> to vector<3x4xbf16>
+ %interleaved = vector.interleave %even_t, %odd_t : vector<3x4xbf16> -> vector<3x8xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<3x8xbf16> to vector<8x3xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %flat, %accum : vector<2x8xbf16>, vector<8x3xbf16> into vector<2x3xbf16>
+ scf.yield %dot : vector<2x3xbf16>
+ }
+ return %result : vector<2x3xbf16>
+}
+
+// Check the equivalent transpose/extract and slice-insert form, including
+// removal of the IV's index -> i32 -> index round-trip.
+func.func @split_join_bf16(%a: memref<1x2x8xbf16>,
+ %b: memref<1x4x6xbf16>) -> vector<2x3xbf16> {
+ %zero = arith.constant 0 : index
+ %end = arith.constant 1 : index
+ %pad = arith.constant 0.0 : bf16
+ %poison = ub.poison : vector<3x4x2xbf16>
+ %init = arith.constant dense<0.0> : vector<2x3xbf16>
+ %result = scf.for %iv = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x3xbf16>) {
+ %k32 = arith.index_cast %iv : index to i32
+ %k = arith.index_castui %k32 nneg : i32 to index
+ %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], %pad {in_bounds = [true, true]} : memref<1x4x6xbf16>, vector<4x6xbf16>
+ %packed = vector.shape_cast %right : vector<4x6xbf16> to vector<4x3x2xbf16>
+ %even_odd = vector.transpose %packed, [2, 0, 1] : vector<4x3x2xbf16> to vector<2x4x3xbf16>
+ %even = vector.extract %even_odd[0] : vector<4x3xbf16> from vector<2x4x3xbf16>
+ %odd = vector.extract %even_odd[1] : vector<4x3xbf16> from vector<2x4x3xbf16>
+ %even_t = vector.transpose %even, [1, 0] : vector<4x3xbf16> to vector<3x4xbf16>
+ %odd_t = vector.transpose %odd, [1, 0] : vector<4x3xbf16> to vector<3x4xbf16>
+ %even3 = vector.shape_cast %even_t : vector<3x4xbf16> to vector<3x4x1xbf16>
+ %odd3 = vector.shape_cast %odd_t : vector<3x4xbf16> to vector<3x4x1xbf16>
+ %join0 = vector.insert_strided_slice %even3, %poison offsets = [0, 0, 0], strides = [1, 1, 1] : vector<3x4x1xbf16> into vector<3x4x2xbf16>
+ %join = vector.insert_strided_slice %odd3, %join0 offsets = [0, 0, 1], strides = [1, 1, 1] : vector<3x4x1xbf16> into vector<3x4x2xbf16>
+ %interleaved = vector.shape_cast %join : vector<3x4x2xbf16> to vector<3x8xbf16>
+ %flat = vector.transpose %interleaved, [1, 0] : vector<3x8xbf16> to vector<8x3xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %flat, %accum : vector<2x8xbf16>, vector<8x3xbf16> into vector<2x3xbf16>
+ scf.yield %dot : vector<2x3xbf16>
+ }
+ return %result : vector<2x3xbf16>
+}
+
func.func @main() {
%zero = arith.constant 0 : index
%af = memref.alloc() : memref<1x2x8xbf16>
@@ -80,6 +136,25 @@ func.func @main() {
%i8_value = vector.extract %i8_result[0, 0] : i8 from vector<2x4xi8>
vector.print %i8_value : i8
// CHECK-NEXT: 100
+
+ // Unequal even/odd elements and rectangular shapes expose parity swaps and
+ // incorrect reshapes when comparing the original and folded contractions.
+ %ai = memref.alloc() : memref<1x2x8xbf16>
+ %bi = memref.alloc() : memref<1x4x6xbf16>
+ %lhs_i = arith.constant dense<[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
+ [2.0, 0.0, 1.0, 3.0, 0.0, 2.0, 1.0, 3.0]]> : vector<2x8xbf16>
+ %rhs_i = arith.constant dense<[[0.0, 1.0, 2.0, 3.0, 1.0, 0.0],
+ [2.0, 0.0, 3.0, 1.0, 0.0, 2.0],
+ [1.0, 3.0, 0.0, 0.0, 2.0, 1.0],
+ [3.0, 2.0, 1.0, 0.0, 1.0, 3.0]]> : vector<4x6xbf16>
+ vector.transfer_write %lhs_i, %ai[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<2x8xbf16>, memref<1x2x8xbf16>
+ vector.transfer_write %rhs_i, %bi[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<4x6xbf16>, memref<1x4x6xbf16>
+ %interleaved = func.call @interleaved_bf16(%ai, %bi) : (memref<1x2x8xbf16>, memref<1x4x6xbf16>) -> vector<2x3xbf16>
+ vector.print %interleaved : vector<2x3xbf16>
+ // CHECK-NEXT: ( ( 68, 28, 56 ), ( 17, 11, 20 ) )
+ %split_join = func.call @split_join_bf16(%ai, %bi) : (memref<1x2x8xbf16>, memref<1x4x6xbf16>) -> vector<2x3xbf16>
+ vector.print %split_join : vector<2x3xbf16>
+ // CHECK-NEXT: ( ( 68, 28, 56 ), ( 17, 11, 20 ) )
return
}
>From 53c0b088d1dc8bc7290db77b5f3df61f4a23ce59 Mon Sep 17 00:00:00 2001
From: fschlimb <frank.schlimbach at intel.com>
Date: Thu, 29 Oct 2026 03:51:30 -0700
Subject: [PATCH 2/2] [mlir][X86] Fold two-level/multi-way split-and-join VNNI
unpacking for i8
---
.../X86/TransformOps/X86TransformOps.td | 9 +-
.../Transforms/PrepareVectorContractLoops.cpp | 174 ++++--
.../X86/vector-contract-prepare-loops.mlir | 499 ++++++++----------
.../CPU/X86/prepare-contract-loops.mlir | 209 +++++---
4 files changed, 494 insertions(+), 397 deletions(-)
diff --git a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
index 19ee3d71957d54..47b6c7d2a94424 100644
--- a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
+++ b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
@@ -99,11 +99,10 @@ def ApplyPrepareVectorContractLoopsPatternsOp : Op<Transform_Dialect,
let description = [{
Apply patterns that rewrite bf16 and i8 matmul loops to contract packed
VNNI operands directly, avoiding RHS unpacking. Supported unpacking uses
- shape casts and transposes, or bf16 even/odd split-and-join sequences.
-
- Before applying these patterns, remove leading unit dimensions from
- vectors with
- `apply_patterns.vector.cast_away_vector_leading_one_dim`.
+ shape casts and transposes, or split-and-join sequences. Splits may use
+ deinterleave (binary) or transpose/extract (any width dividing the VNNI
+ factor); joins may use interleave (binary) or slice inserts (any width
+ dividing the VNNI factor).
}];
let assemblyFormat = "attr-dict";
diff --git a/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp b/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
index 1097db2fa14555..432d28b92b2f6b 100644
--- a/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/PrepareVectorContractLoops.cpp
@@ -74,10 +74,11 @@ static Value skipReshapes(Value v) {
return v;
}
-// Find the source and parity of an even/odd split: v[i] = src[2 * i + j].
-// Match either pattern below (unit dimensions may also be present):
+// Find the source and VNNI lane j of a split: v[i] = src[factor * i + j].
+// A split step of width s keeps every s-th element. Match either form below
+// (unit dimensions may also be present):
//
-// %even, %odd = vector.deinterleave %src
+// %even, %odd = vector.deinterleave %src // s = 2
// : vector<4x8xbf16> -> vector<4x4xbf16>
//
// or:
@@ -90,11 +91,36 @@ static Value skipReshapes(Value v) {
// : vector<4x4xbf16> from vector<2x4x4xbf16>
// %odd = vector.extract %split[1]
// : vector<4x4xbf16> from vector<2x4x4xbf16>
-static std::optional<std::pair<Value, int64_t>> matchEvenOrOdd(Value v) {
+//
+// If s < factor, the split's input must itself come from splits of src
+// totaling factor / s. For example, for factor = 4, a 4-way transpose/extract
+// split of src matches, as do two levels of binary splits:
+//
+// %even, %odd = vector.deinterleave %src // lanes {0, 2} and {1, 3}
+// %l1, %l3 = vector.deinterleave %odd // %l3[i] = src[4 * i + 3]
+//
+// Splits closer to src give the low lane bits; result j of the split that
+// produces v gives the high bits.
+static std::optional<std::pair<Value, int64_t>> matchSplit(Value v,
+ int64_t factor) {
v = skipReshapes(v);
+ if (factor == 1)
+ return std::pair(v, int64_t(0));
+
+ auto matchSource =
+ [&](Value source, int64_t ways,
+ int64_t j) -> std::optional<std::pair<Value, int64_t>> {
+ if (factor % ways != 0)
+ return std::nullopt;
+ auto lane = matchSplit(source, factor / ways);
+ if (!lane)
+ return std::nullopt;
+ lane->second += factor / ways * j;
+ return lane;
+ };
if (auto deinterleave = v.getDefiningOp<vector::DeinterleaveOp>())
- return std::pair(skipReshapes(deinterleave.getSource()),
- int64_t(cast<OpResult>(v).getResultNumber()));
+ return matchSource(deinterleave.getSource(), 2,
+ cast<OpResult>(v).getResultNumber());
auto extract = v.getDefiningOp<vector::ExtractOp>();
auto transpose =
@@ -109,37 +135,38 @@ static std::optional<std::pair<Value, int64_t>> matchEvenOrOdd(Value v) {
auto isNonUnit = [&](int64_t dim) { return shape[dim] != 1; };
SmallVector<int64_t> order =
llvm::filter_to_vector(llvm::seq<int64_t>(0, shape.size()), isNonUnit);
- if (order.empty() || shape[order.back()] != 2)
+ if (order.empty())
return std::nullopt;
+ int64_t ways = shape[order.back()];
std::rotate(order.begin(), std::prev(order.end()), order.end());
ArrayRef<int64_t> perm = transpose.getPermutation();
if (llvm::filter_to_vector(perm, isNonUnit) != order)
return std::nullopt;
- // The dimension of size 2 is now first among the non-unit dimensions.
- // Keeping exactly half the elements means the extract indexes that
- // dimension; any other dimensions it indexes must be unit dimensions.
- if (2 * type.getNumElements() !=
+ // The split dimension is now first among the non-unit dimensions. Keeping
+ // exactly 1 / ways of the elements means the extract indexes that dimension;
+ // any other dimensions it indexes must be unit dimensions.
+ if (ways * type.getNumElements() !=
transpose.getSourceVectorType().getNumElements())
return std::nullopt;
int64_t j =
extract
.getStaticPosition()[llvm::find(perm, order.front()) - perm.begin()];
- if (j != 0 && j != 1)
+ if (j < 0 || j >= ways)
return std::nullopt;
- return std::pair(skipReshapes(transpose.getVector()), j);
+ return matchSource(transpose.getVector(), ways, j);
}
-// Find the even and odd inputs of either pattern below. The slice inserts may
-// appear in either order.
+// Find the inputs in VNNI-lane order for a join. A join step of width s
+// interleaves s inputs. Match either form below:
//
-// %joined = vector.interleave %even, %odd
+// %joined = vector.interleave %even, %odd // s = 2
// : vector<4x1xbf16> -> vector<4x2xbf16>
//
-// or:
+// or s slice inserts in any order (shown for s = 2):
//
// %poison = ub.poison : vector<4x2xbf16>
// %partial = vector.insert_strided_slice %odd, %poison
@@ -148,37 +175,61 @@ static std::optional<std::pair<Value, int64_t>> matchEvenOrOdd(Value v) {
// %joined = vector.insert_strided_slice %even, %partial
// offsets = [0, 0], strides = [1, 1]
// : vector<4x1xbf16> into vector<4x2xbf16>
-static std::optional<std::array<Value, 2>> matchInterleave(Value v) {
+//
+// If s < factor, each input must itself be a join of width factor / s
+// (possibly built from smaller joins).
+static std::optional<SmallVector<Value, 4>> matchJoin(Value v, int64_t factor) {
v = skipReshapes(v);
- if (auto interleave = v.getDefiningOp<vector::InterleaveOp>())
- return std::array<Value, 2>{interleave.getLhs(), interleave.getRhs()};
-
- // Both slices must be complete and disjoint so the original destination
- // contributes no elements to the result.
- std::array<Value, 2> halves;
- for (int i = 0; i < 2; ++i) {
- auto insert = v.getDefiningOp<vector::InsertStridedSliceOp>();
- if (!insert)
+ if (factor == 1)
+ return SmallVector<Value, 4>{v};
+
+ SmallVector<Value, 4> parts;
+ if (auto interleave = v.getDefiningOp<vector::InterleaveOp>()) {
+ parts = {interleave.getLhs(), interleave.getRhs()};
+ } else if (auto insert = v.getDefiningOp<vector::InsertStridedSliceOp>()) {
+ int64_t ways = insert.getDestVectorType().getShape().back();
+ if (ways < 2 || factor % ways != 0)
return std::nullopt;
+ parts.resize(ways);
+ // All slices must be complete and disjoint so the original destination
+ // contributes no elements to the result.
+ for (int64_t i = 0; i < ways; ++i) {
+ auto insert = v.getDefiningOp<vector::InsertStridedSliceOp>();
+ if (!insert)
+ return std::nullopt;
+ ArrayRef<int64_t> shape = insert.getDestVectorType().getShape();
+ SmallVector<int64_t> slice(shape.drop_back());
+ slice.push_back(1);
+ int64_t j =
+ extractFromIntegerArrayAttr<int64_t>(insert.getOffsets()).back();
+ if (insert.getSourceVectorType().getShape() != ArrayRef(slice) ||
+ parts[j])
+ return std::nullopt;
+ parts[j] = insert.getValueToStore();
+ v = insert.getDest();
+ }
+ } else {
+ return std::nullopt;
+ }
- ArrayRef<int64_t> shape = insert.getDestVectorType().getShape();
- SmallVector<int64_t> slice(shape.drop_back());
- slice.push_back(1);
- auto offsets = extractFromIntegerArrayAttr<int64_t>(insert.getOffsets());
- int64_t j = offsets.back();
- if (shape.back() != 2 ||
- insert.getSourceVectorType().getShape() != ArrayRef(slice) || halves[j])
+ int64_t ways = parts.size();
+ if (factor % ways != 0)
+ return std::nullopt;
+ SmallVector<Value, 4> lanes(factor);
+ for (auto [j, part] : llvm::enumerate(parts)) {
+ auto inner = matchJoin(part, factor / ways);
+ if (!inner)
return std::nullopt;
-
- halves[j] = insert.getValueToStore();
- v = insert.getDest();
+ for (auto [i, lane] : llvm::enumerate(*inner))
+ lanes[ways * i + j] = lane;
}
- return halves;
+ return lanes;
}
// Find the packed RHS so the contraction can use it without unpacking.
-// For a K x N RHS and VNNI factor f, require
-// rhs[f * r + j][c] = src[r][c][j], with src viewed as K/f x N x f.
+// For a K x N RHS and VNNI factor f, require for each VNNI lane j:
+// rhs[f * r + j][c] = src[r][c][j],
+// where 0 <= j < f and src is viewed as K/f x N x f.
// Match transpose/shape_cast unpacking (shown for f = 2):
//
// %packed = vector.shape_cast %src
@@ -188,7 +239,7 @@ static std::optional<std::array<Value, 2>> matchInterleave(Value v) {
// %rhs = vector.shape_cast %transposed
// : vector<4x2x3xbf16> to vector<8x3xbf16>
//
-// or, for f = 2, split/transpose/join/transpose unpacking:
+// or split/transpose/join/transpose unpacking (shown for f = 2):
//
// %even, %odd = vector.deinterleave %src
// : vector<4x6xbf16> -> vector<4x3xbf16>
@@ -200,6 +251,21 @@ static std::optional<std::array<Value, 2>> matchInterleave(Value v) {
// : vector<3x4xbf16> -> vector<3x8xbf16>
// %rhs = vector.transpose %joined, [1, 0]
// : vector<3x8xbf16> to vector<8x3xbf16>
+//
+// For f = 4, splits and joins may each be one 4-way transpose/extract or
+// slice-insert step, or two levels of binary steps. With binary steps, the
+// second-level splits produce VNNI lanes 0/2 and 1/3:
+//
+// %even, %odd = vector.deinterleave %src // lanes {0, 2} and {1, 3}
+// %l0, %l2 = vector.deinterleave %even
+// %l1, %l3 = vector.deinterleave %odd
+// // Transpose each lane %lj to %tj.
+// %even_t = vector.interleave %t0, %t2
+// %odd_t = vector.interleave %t1, %t3
+// %joined = vector.interleave %even_t, %odd_t
+//
+// Any binary split or join may instead use the transpose/extract or
+// slice-insert form.
static Value matchUnpackedRhs(Value rhs, int64_t factor) {
auto type = cast<VectorType>(rhs.getType());
int64_t k = type.getDimSize(0) / factor, n = type.getDimSize(1);
@@ -222,28 +288,28 @@ static Value matchUnpackedRhs(Value rhs, int64_t factor) {
}
const int64_t swap[] = {1, 0};
- if (factor != 2 || perm != ArrayRef(swap) ||
- shape != ArrayRef<int64_t>{n, 2 * k})
+ if (perm != ArrayRef(swap) || shape != ArrayRef<int64_t>{n, factor * k})
return {};
- std::optional<std::array<Value, 2>> halves =
- matchInterleave(transpose.getVector());
- if (!halves)
+ auto lanes = matchJoin(transpose.getVector(), factor);
+ if (!lanes)
return {};
Value src;
- for (auto [j, half] : llvm::enumerate(*halves)) {
- auto transpose = skipReshapes(half).getDefiningOp<vector::TransposeOp>();
- if (!transpose || transpose.getPermutation() != ArrayRef(swap) ||
- transpose.getSourceVectorType().getShape() != ArrayRef<int64_t>{k, n})
+ for (auto [j, lane] : llvm::enumerate(*lanes)) {
+ auto laneTranspose =
+ skipReshapes(lane).getDefiningOp<vector::TransposeOp>();
+ if (!laneTranspose || laneTranspose.getPermutation() != ArrayRef(swap) ||
+ laneTranspose.getSourceVectorType().getShape() !=
+ ArrayRef<int64_t>{k, n})
return {};
- auto evenOrOdd = matchEvenOrOdd(transpose.getVector());
- if (!evenOrOdd || evenOrOdd->second != int64_t(j) ||
- (src && evenOrOdd->first != src))
+ auto source = matchSplit(laneTranspose.getVector(), factor);
+ if (!source || source->second != int64_t(j) ||
+ (src && source->first != src))
return {};
- src = evenOrOdd->first;
+ src = source->first;
}
return src;
}
diff --git a/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir b/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
index 2bcac9a35caa78..82d9cf8616e527 100644
--- a/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-prepare-loops.mlir
@@ -9,165 +9,44 @@
// 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-COUNT-4: x86.amx.tile_mulf
-// NANO-NOT: vector.contract
-// NANO-NOT: x86.amx.tile_mulf
-// 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 flat bf16 rhs read reshaped into VNNI layout
-// 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-COUNT-4: x86.amx.tile_mulf
-// NANO-NOT: vector.contract
-// NANO-NOT: x86.amx.tile_mulf
-// 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 deinterleave/transpose/interleave/transpose unpacking with strided
-// memrefs and an i32 loop index.
+// Fold deinterleave/transpose/interleave/transpose unpacking of non-square
+// operands with strided memrefs, aligned nonzero inner read offsets, and
+// fastmath flags.
// CHECK-LABEL: func.func @interleaved_vnni_bf16(
-// CHECK-SAME: %[[A:[^:]+]]: memref<?x?x32x32xbf16, strided<[?, 1024, 32, 1]>>, %[[B:[^:]+]]: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>
+// CHECK-SAME: %[[A:[^:]+]]: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>, %[[B:[^:]+]]: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>
// NANO-LABEL: func.func @interleaved_vnni_bf16(
-// NANO: vector.contract {{.*}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
+// NANO: vector.contract {{.*}} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
// NANO: return
-// CHECK-DAG: %[[AP:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x?x32x16x2xbf16, strided<[?, 1024, 32, 2, 1]>>
+// CHECK-DAG: %[[ONE:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[AP:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x?x16x32x2xbf16, strided<[?, 1024, 64, 2, 1]>>
// CHECK-DAG: %[[BP:.+]] = memref.expand_shape %[[B]] {{.*}} into memref<?x?x16x32x2xbf16, strided<[?, 1024, 64, 2, 1]>>
-// CHECK: scf.for {{.*}} : i32 {
-// CHECK: %[[K:.+]] = arith.index_cast {{.*}} : i32 to index
-// CHECK: %[[LHS:.+]] = vector.transfer_read %[[AP]][%{{.+}}, %[[K]], {{.*}} vector<32x16x2xbf16>
-// CHECK-NEXT: %[[RHS:.+]] = vector.transfer_read %[[BP]][%{{.+}}, %[[K]], {{.*}} vector<16x32x2xbf16>
-// CHECK-NEXT: %[[DOT:.+]] = 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>
+// CHECK: scf.for %[[K:[^ ]+]] =
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[AP]][%{{.+}}, %[[K]], %{{.+}}, %[[ONE]], %{{.+}}]{{.*}} vector<16x16x2xbf16>
+// CHECK-NEXT: %[[RHS:.+]] = vector.transfer_read %[[BP]][%{{.+}}, %[[K]], %{{.+}}, %[[ONE]], %{{.+}}]{{.*}} vector<16x16x2xbf16>
+// CHECK-NEXT: %[[DOT:.+]] = 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<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
// CHECK-NEXT: scf.yield %[[DOT]]
-func.func @interleaved_vnni_bf16(%a: memref<?x?x32x32xbf16, strided<[?, 1024, 32, 1]>>,
+func.func @interleaved_vnni_bf16(%a: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>,
%b: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>,
- %init: vector<32x32xf32>, %lb: i32, %ub: i32,
- %m: index, %n: index) -> vector<32x32xf32> {
+ %init: vector<16x16xf32>, %lb: index, %ub: index,
+ %m: index, %n: index) -> vector<16x16xf32> {
%zero = arith.constant 0 : index
- %step = arith.constant 1 : i32
- %pad = arith.constant 0.0 : bf16
- %result = scf.for %iv = %lb to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) : i32 {
- %k = arith.index_cast %iv : i32 to index
- %lhs = vector.transfer_read %a[%m, %k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x?x32x32xbf16, strided<[?, 1024, 32, 1]>>, vector<32x32xbf16>
- %rhs = vector.transfer_read %b[%n, %k, %zero, %zero], %pad {in_bounds = [true, true]} : memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>, vector<16x64xbf16>
- %even, %odd = vector.deinterleave %rhs : vector<16x64xbf16> -> vector<16x32xbf16>
- %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
- %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
- %interleaved = vector.interleave %even_t, %odd_t : vector<32x16xbf16> -> vector<32x32xbf16>
- %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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>
-}
-
-// Also handle non-square operands and aligned nonzero inner read offsets.
-// CHECK-LABEL: func.func @interleaved_vnni_rectangular(
-// CHECK: %[[ONE:.+]] = arith.constant 1 : index
-// CHECK: memref.expand_shape {{.*}} into memref<?x16x32x2xbf16>
-// CHECK: memref.expand_shape {{.*}} into memref<?x16x32x2xbf16>
-// CHECK: scf.for %[[K:[^ ]+]] =
-// CHECK: vector.transfer_read %{{.+}}[%[[K]], %{{.+}}, %[[ONE]], %{{.+}}]{{.*}} vector<16x16x2xbf16>
-// CHECK: vector.transfer_read %{{.+}}[%[[K]], %{{.+}}, %[[ONE]], %{{.+}}]{{.*}} vector<16x16x2xbf16>
-// CHECK-NOT: vector.deinterleave
-// CHECK-NOT: vector.interleave
-// CHECK-NOT: vector.transpose
-// CHECK: vector.contract {{.*}} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
-func.func @interleaved_vnni_rectangular(%a: memref<?x16x64xbf16>,
- %b: memref<?x16x64xbf16>,
- %init: vector<16x16xf32>,
- %ub: index) -> vector<16x16xf32> {
- %zero = arith.constant 0 : index
- %step = arith.constant 1 : index
%offset = arith.constant 2 : index
+ %step = arith.constant 1 : index
%pad = arith.constant 0.0 : bf16
- %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<16x16xf32>) {
- %lhs = vector.transfer_read %a[%k, %zero, %offset], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x32xbf16>
- %rhs = vector.transfer_read %b[%k, %zero, %offset], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x32xbf16>
+ %result = scf.for %k = %lb to %ub step %step iter_args(%accum = %init) -> (vector<16x16xf32>) {
+ %lhs = vector.transfer_read %a[%m, %k, %zero, %offset], %pad {in_bounds = [true, true]} : memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>, vector<16x32xbf16>
+ %rhs = vector.transfer_read %b[%n, %k, %zero, %offset], %pad {in_bounds = [true, true]} : memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>, vector<16x32xbf16>
%even, %odd = vector.deinterleave %rhs : vector<16x32xbf16> -> vector<16x16xbf16>
%even_t = vector.transpose %even, [1, 0] : vector<16x16xbf16> to vector<16x16xbf16>
%odd_t = vector.transpose %odd, [1, 0] : vector<16x16xbf16> to vector<16x16xbf16>
%interleaved = vector.interleave %even_t, %odd_t : vector<16x16xbf16> -> vector<16x32xbf16>
%flat = vector.transpose %interleaved, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
- %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %flat, %accum : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
+ %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<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
scf.yield %dot : vector<16x16xf32>
}
return %result : vector<16x16xf32>
}
-// Swapping the even and odd elements changes the packed element order.
-// CHECK-LABEL: func.func @negative_interleaved_swapped(
-// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
-func.func @negative_interleaved_swapped(%a: memref<?x32x32xbf16>, %b: memref<?x16x64xbf16>,
- %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
- %zero = arith.constant 0 : index
- %step = arith.constant 1 : index
- %pad = arith.constant 0.0 : bf16
- %r0 = 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], %pad {in_bounds = [true, true]} : memref<?x16x64xbf16>, vector<16x64xbf16>
- %even, %odd = vector.deinterleave %rhs : vector<16x64xbf16> -> vector<16x32xbf16>
- %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
- %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
- %interleaved = vector.interleave %odd_t, %even_t : vector<32x16xbf16> -> vector<32x32xbf16>
- %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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 %r0 : vector<32x32xf32>
-}
-
// The even and odd elements must come from the same packed source.
// CHECK-LABEL: func.func @negative_interleaved_different_sources(
// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
@@ -193,27 +72,6 @@ func.func @negative_interleaved_different_sources(%a: memref<?x32x32xbf16>, %b:
return %r1 : vector<32x32xf32>
}
-// A loop-invariant RHS does not provide a packed read to fold.
-// CHECK-LABEL: func.func @negative_interleaved_invariant(
-// CHECK: vector.contract {{.*}} : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
-func.func @negative_interleaved_invariant(%a: memref<?x32x32xbf16>, %invariant: vector<16x64xbf16>,
- %init: vector<32x32xf32>, %ub: index) -> vector<32x32xf32> {
- %zero = arith.constant 0 : index
- %step = arith.constant 1 : index
- %pad = arith.constant 0.0 : bf16
- %r2 = 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>
- %even, %odd = vector.deinterleave %invariant : vector<16x64xbf16> -> vector<16x32xbf16>
- %even_t = vector.transpose %even, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
- %odd_t = vector.transpose %odd, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
- %interleaved = vector.interleave %even_t, %odd_t : vector<32x16xbf16> -> vector<32x32xbf16>
- %flat = vector.transpose %interleaved, [1, 0] : vector<32x32xbf16> 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 %r2 : vector<32x32xf32>
-}
-
// Fold transpose/extract and slice-insert unpacking despite interspersed unit
// dimensions and odd-first insertion.
// CHECK-LABEL: func.func @split_join_bf16(
@@ -386,44 +244,12 @@ func.func @direct_vnni_int8(%a: memref<?x32x64xi8>,
return %result : vector<32x32xi32>
}
-// fold VNNI unpacking of flat i8 rhs read reshaped into VNNI layout
-// 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-COUNT-4: x86.amx.tile_muli
-// NANO-NOT: vector.contract
-// NANO-NOT: x86.amx.tile_muli
-// 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>
-}
-
-// 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(
+// Fold transpose/shape_cast unpacking of a packed RHS read through a leading
+// unit-dimension broadcast, and expand the LHS through a static view of its
+// dynamic innermost dimension.
+// CHECK-LABEL: func.func @packed_vnni_bf16(
+// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x?xbf16>, %[[B:[^,]+]]: memref<?x16x32x2xbf16>
+// NANO-LABEL: func.func @packed_vnni_bf16(
// NANO-NOT: vector.contract
// NANO-COUNT-4: x86.amx.tile_mulf
// NANO-NOT: vector.contract
@@ -433,49 +259,19 @@ func.func @split_vnni_int8(%a: memref<?x32x64xi8>,
// 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> {
+// CHECK-NEXT: %[[RHS:.+]] = vector.transfer_read %[[B]][%[[K]], {{.*}} vector<16x32x2xbf16>
+// CHECK-NEXT: 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>
+// CHECK-NEXT: scf.yield
+func.func @packed_vnni_bf16(%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>
-}
-
-// fold a leading unit-dimension broadcast of an already packed RHS
-// CHECK-LABEL: func.func @broadcast_vnni_bf16(
-// CHECK: memref.expand_shape
-// CHECK: %[[LHS:.+]] = vector.transfer_read {{.*}} vector<32x16x2xbf16>
-// CHECK: %[[RHS:.+]] = vector.transfer_read {{.*}} vector<16x32x2xbf16>
-// CHECK-NOT: vector.broadcast
-// CHECK-NOT: vector.transpose
-// CHECK-NOT: vector.shape_cast
-// CHECK: vector.contract {{.*}} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
-// NANO-LABEL: func.func @broadcast_vnni_bf16(
-// NANO-NOT: vector.contract
-// NANO-COUNT-4: x86.amx.tile_mulf
-// NANO-NOT: vector.contract
-// NANO-NOT: x86.amx.tile_mulf
-// NANO: return
-func.func @broadcast_vnni_bf16(%a: memref<?x32x32xbf16>,
- %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<?x32x32xbf16>, vector<32x32xbf16>
- %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x2xbf16>, vector<16x32x2xbf16>
%broadcast = vector.broadcast %rhs : vector<16x32x2xbf16> to vector<1x16x32x2xbf16>
%transpose = vector.transpose %broadcast, [0, 1, 3, 2] : vector<1x16x32x2xbf16> to vector<1x16x2x32xbf16>
%flat = vector.shape_cast %transpose : vector<1x16x2x32xbf16> to vector<32x32xbf16>
@@ -562,35 +358,6 @@ func.func @negative_nondivisible_inner_k(%a: memref<?x32x33xbf16>,
return %result : vector<32x32xf32>
}
-// negative test: even and odd rhs elements are swapped before unpacking -> skip
-// CHECK-LABEL: func.func @negative_interleaved_elements(
-// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x64xi8>
-// CHECK-NOT: memref.expand_shape
-// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A]][{{.*}} vector<32x64xi8>
-// CHECK: %[[EVEN:[^,]+]], %[[ODD:.+]] = vector.deinterleave
-// CHECK: %[[REORDERED:.+]] = vector.interleave %[[ODD]], %[[EVEN]]
-// CHECK: %[[T:.+]] = vector.transpose %[[REORDERED]], [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 @negative_interleaved_elements(%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 @negative_scalable_vnni(
// CHECK-SAME: %[[A:[^,]+]]: memref<?x?x32xbf16>
@@ -643,6 +410,208 @@ func.func @negative_vector_element_lhs(%a: memref<?x32xvector<32xbf16>>,
return %result : vector<32x32xf32>
}
+// Fold i8 unpacking through two levels of binary splits and joins. The first
+// split yields VNNI lanes {0, 2} and {1, 3}; mix deinterleave with
+// transpose/extract, and interleave with out-of-order slice inserts.
+// CHECK-LABEL: func.func @mixed_split_join_int8(
+// CHECK-DAG: %[[AP:.+]] = memref.expand_shape {{.*}} into memref<?x32x16x4xi8>
+// CHECK-DAG: %[[BP:.+]] = memref.expand_shape {{.*}} into memref<?x16x32x4xi8>
+// CHECK: scf.for
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[AP]][{{.*}} vector<32x16x4xi8>
+// CHECK-NEXT: %[[RHS:.+]] = vector.transfer_read %[[BP]][{{.*}} vector<16x32x4xi8>
+// CHECK-NEXT: vector.contract {{.*}} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x4xi8>, vector<16x32x4xi8> into vector<32x32xi32>
+// CHECK-NEXT: scf.yield
+// NANO-LABEL: func.func @mixed_split_join_int8(
+// NANO-NOT: vector.contract
+// NANO-COUNT-4: x86.amx.tile_muli
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_muli
+// NANO: return
+func.func @mixed_split_join_int8(%a: memref<?x32x64xi8>, %b: memref<?x16x128xi8>,
+ %init: vector<32x32xi32>, %ub: index) -> vector<32x32xi32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0 : i8
+ %poison = ub.poison : vector<32x16x2xi8>
+ %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], %pad {in_bounds = [true, true]} : memref<?x16x128xi8>, vector<16x128xi8>
+ %even, %odd = vector.deinterleave %rhs : vector<16x128xi8> -> vector<16x64xi8>
+ %lane0, %lane2 = vector.deinterleave %even : vector<16x64xi8> -> vector<16x32xi8>
+ %odd_packed = vector.shape_cast %odd : vector<16x64xi8> to vector<16x32x2xi8>
+ %odd_split = vector.transpose %odd_packed, [2, 0, 1] : vector<16x32x2xi8> to vector<2x16x32xi8>
+ %lane1 = vector.extract %odd_split[0] : vector<16x32xi8> from vector<2x16x32xi8>
+ %lane3 = vector.extract %odd_split[1] : vector<16x32xi8> from vector<2x16x32xi8>
+ %t0 = vector.transpose %lane0, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t1 = vector.transpose %lane1, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t2 = vector.transpose %lane2, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t3 = vector.transpose %lane3, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %even_t = vector.interleave %t0, %t2 : vector<32x16xi8> -> vector<32x32xi8>
+ %s1 = vector.shape_cast %t1 : vector<32x16xi8> to vector<32x16x1xi8>
+ %s3 = vector.shape_cast %t3 : vector<32x16xi8> to vector<32x16x1xi8>
+ %join0 = vector.insert_strided_slice %s3, %poison offsets = [0, 0, 1], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x2xi8>
+ %join1 = vector.insert_strided_slice %s1, %join0 offsets = [0, 0, 0], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x2xi8>
+ %odd_t = vector.shape_cast %join1 : vector<32x16x2xi8> to vector<32x32xi8>
+ %joined = vector.interleave %even_t, %odd_t : vector<32x32xi8> -> vector<32x64xi8>
+ %flat = vector.transpose %joined, [1, 0] : vector<32x64xi8> 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>
+}
+
+// Joining adjacent VNNI lanes at the inner level swaps VNNI lanes 1 and 2.
+// CHECK-LABEL: func.func @negative_int8_lane_order(
+// CHECK-NOT: memref.expand_shape
+// CHECK: vector.contract {{.*}} : vector<32x64xi8>, vector<64x32xi8> into vector<32x32xi32>
+func.func @negative_int8_lane_order(%a: memref<?x32x64xi8>, %b: memref<?x16x128xi8>,
+ %init: vector<32x32xi32>, %ub: index) -> vector<32x32xi32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : 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], %pad {in_bounds = [true, true]} : memref<?x16x128xi8>, vector<16x128xi8>
+ %even, %odd = vector.deinterleave %rhs : vector<16x128xi8> -> vector<16x64xi8>
+ %lane0, %lane2 = vector.deinterleave %even : vector<16x64xi8> -> vector<16x32xi8>
+ %lane1, %lane3 = vector.deinterleave %odd : vector<16x64xi8> -> vector<16x32xi8>
+ %t0 = vector.transpose %lane0, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t1 = vector.transpose %lane1, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t2 = vector.transpose %lane2, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t3 = vector.transpose %lane3, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %join0 = vector.interleave %t0, %t1 : vector<32x16xi8> -> vector<32x32xi8>
+ %join1 = vector.interleave %t2, %t3 : vector<32x16xi8> -> vector<32x32xi8>
+ %joined = vector.interleave %join0, %join1 : vector<32x32xi8> -> vector<32x64xi8>
+ %flat = vector.transpose %joined, [1, 0] : vector<32x64xi8> 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 i8 unpacking through a single 4-way transpose/extract split and a
+// 4-way slice-insert join with out-of-order inserts.
+// CHECK-LABEL: func.func @fourway_split_join_int8(
+// CHECK-DAG: %[[AP:.+]] = memref.expand_shape {{.*}} into memref<?x32x16x4xi8>
+// CHECK-DAG: %[[BP:.+]] = memref.expand_shape {{.*}} into memref<?x16x32x4xi8>
+// CHECK: scf.for
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[AP]][{{.*}} vector<32x16x4xi8>
+// CHECK-NEXT: %[[RHS:.+]] = vector.transfer_read %[[BP]][{{.*}} vector<16x32x4xi8>
+// CHECK-NEXT: vector.contract {{.*}} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x4xi8>, vector<16x32x4xi8> into vector<32x32xi32>
+// CHECK-NEXT: scf.yield
+// NANO-LABEL: func.func @fourway_split_join_int8(
+// NANO-NOT: vector.contract
+// NANO-COUNT-4: x86.amx.tile_muli
+// NANO-NOT: vector.contract
+// NANO-NOT: x86.amx.tile_muli
+// NANO: return
+func.func @fourway_split_join_int8(%a: memref<?x32x64xi8>, %b: memref<?x16x128xi8>,
+ %init: vector<32x32xi32>, %ub: index) -> vector<32x32xi32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0 : i8
+ %poison = ub.poison : vector<32x16x4xi8>
+ %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], %pad {in_bounds = [true, true]} : memref<?x16x128xi8>, vector<16x128xi8>
+ %packed = vector.shape_cast %rhs : vector<16x128xi8> to vector<16x32x4xi8>
+ %split = vector.transpose %packed, [2, 0, 1] : vector<16x32x4xi8> to vector<4x16x32xi8>
+ %lane0 = vector.extract %split[0] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane1 = vector.extract %split[1] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane2 = vector.extract %split[2] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane3 = vector.extract %split[3] : vector<16x32xi8> from vector<4x16x32xi8>
+ %t0 = vector.transpose %lane0, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t1 = vector.transpose %lane1, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t2 = vector.transpose %lane2, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t3 = vector.transpose %lane3, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %s0 = vector.shape_cast %t0 : vector<32x16xi8> to vector<32x16x1xi8>
+ %s1 = vector.shape_cast %t1 : vector<32x16xi8> to vector<32x16x1xi8>
+ %s2 = vector.shape_cast %t2 : vector<32x16xi8> to vector<32x16x1xi8>
+ %s3 = vector.shape_cast %t3 : vector<32x16xi8> to vector<32x16x1xi8>
+ %join0 = vector.insert_strided_slice %s2, %poison offsets = [0, 0, 2], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x4xi8>
+ %join1 = vector.insert_strided_slice %s0, %join0 offsets = [0, 0, 0], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x4xi8>
+ %join2 = vector.insert_strided_slice %s3, %join1 offsets = [0, 0, 3], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x4xi8>
+ %join3 = vector.insert_strided_slice %s1, %join2 offsets = [0, 0, 1], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x4xi8>
+ %joined = vector.shape_cast %join3 : vector<32x16x4xi8> to vector<32x64xi8>
+ %flat = vector.transpose %joined, [1, 0] : vector<32x64xi8> 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 a 4-way split combined with two levels of binary joins.
+// CHECK-LABEL: func.func @fourway_split_binary_join_int8(
+// CHECK-DAG: memref.expand_shape {{.*}} into memref<?x32x16x4xi8>
+// CHECK-DAG: memref.expand_shape {{.*}} into memref<?x16x32x4xi8>
+// CHECK: vector.contract {{.*}} : vector<32x16x4xi8>, vector<16x32x4xi8> into vector<32x32xi32>
+func.func @fourway_split_binary_join_int8(%a: memref<?x32x64xi8>, %b: memref<?x16x128xi8>,
+ %init: vector<32x32xi32>, %ub: index) -> vector<32x32xi32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : 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], %pad {in_bounds = [true, true]} : memref<?x16x128xi8>, vector<16x128xi8>
+ %packed = vector.shape_cast %rhs : vector<16x128xi8> to vector<16x32x4xi8>
+ %split = vector.transpose %packed, [2, 0, 1] : vector<16x32x4xi8> to vector<4x16x32xi8>
+ %lane0 = vector.extract %split[0] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane1 = vector.extract %split[1] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane2 = vector.extract %split[2] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane3 = vector.extract %split[3] : vector<16x32xi8> from vector<4x16x32xi8>
+ %t0 = vector.transpose %lane0, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t1 = vector.transpose %lane1, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t2 = vector.transpose %lane2, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t3 = vector.transpose %lane3, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %even_t = vector.interleave %t0, %t2 : vector<32x16xi8> -> vector<32x32xi8>
+ %odd_t = vector.interleave %t1, %t3 : vector<32x16xi8> -> vector<32x32xi8>
+ %joined = vector.interleave %even_t, %odd_t : vector<32x32xi8> -> vector<32x64xi8>
+ %flat = vector.transpose %joined, [1, 0] : vector<32x64xi8> 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>
+}
+
+// Inserting lanes 1 and 2 at each other's offsets swaps VNNI lanes.
+// CHECK-LABEL: func.func @negative_fourway_join_lane_order(
+// CHECK-NOT: memref.expand_shape
+// CHECK: vector.contract {{.*}} : vector<32x64xi8>, vector<64x32xi8> into vector<32x32xi32>
+func.func @negative_fourway_join_lane_order(%a: memref<?x32x64xi8>, %b: memref<?x16x128xi8>,
+ %init: vector<32x32xi32>, %ub: index) -> vector<32x32xi32> {
+ %zero = arith.constant 0 : index
+ %step = arith.constant 1 : index
+ %pad = arith.constant 0 : i8
+ %poison = ub.poison : vector<32x16x4xi8>
+ %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], %pad {in_bounds = [true, true]} : memref<?x16x128xi8>, vector<16x128xi8>
+ %packed = vector.shape_cast %rhs : vector<16x128xi8> to vector<16x32x4xi8>
+ %split = vector.transpose %packed, [2, 0, 1] : vector<16x32x4xi8> to vector<4x16x32xi8>
+ %lane0 = vector.extract %split[0] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane1 = vector.extract %split[1] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane2 = vector.extract %split[2] : vector<16x32xi8> from vector<4x16x32xi8>
+ %lane3 = vector.extract %split[3] : vector<16x32xi8> from vector<4x16x32xi8>
+ %t0 = vector.transpose %lane0, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t1 = vector.transpose %lane1, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t2 = vector.transpose %lane2, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %t3 = vector.transpose %lane3, [1, 0] : vector<16x32xi8> to vector<32x16xi8>
+ %s0 = vector.shape_cast %t0 : vector<32x16xi8> to vector<32x16x1xi8>
+ %s1 = vector.shape_cast %t1 : vector<32x16xi8> to vector<32x16x1xi8>
+ %s2 = vector.shape_cast %t2 : vector<32x16xi8> to vector<32x16x1xi8>
+ %s3 = vector.shape_cast %t3 : vector<32x16xi8> to vector<32x16x1xi8>
+ %join0 = vector.insert_strided_slice %s0, %poison offsets = [0, 0, 0], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x4xi8>
+ %join1 = vector.insert_strided_slice %s2, %join0 offsets = [0, 0, 1], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x4xi8>
+ %join2 = vector.insert_strided_slice %s1, %join1 offsets = [0, 0, 2], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x4xi8>
+ %join3 = vector.insert_strided_slice %s3, %join2 offsets = [0, 0, 3], strides = [1, 1, 1] : vector<32x16x1xi8> into vector<32x16x4xi8>
+ %joined = vector.shape_cast %join3 : vector<32x16x4xi8> to vector<32x64xi8>
+ %flat = vector.transpose %joined, [1, 0] : vector<32x64xi8> 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>
+}
+
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
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
index a2c73666a282e9..ad6d0a951d0307 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/X86/prepare-contract-loops.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/X86/prepare-contract-loops.mlir
@@ -14,50 +14,48 @@
// Generic CPU lowering executes VNNI contractions only without mixed precision.
func.func @vnni_bf16(%a: memref<1x2x8xbf16>,
- %b: memref<1x4x4x2xbf16>) -> vector<2x4xbf16> {
+ %b: memref<1x4x3x2xbf16>) -> vector<2x3xbf16> {
%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>) {
+ %init = arith.constant dense<0.0> : vector<2x3xbf16>
+ %result = scf.for %k = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x3xbf16>) {
%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>
+ %right = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<1x4x3x2xbf16>, vector<4x3x2xbf16>
+ %transpose = vector.transpose %right, [0, 2, 1] : vector<4x3x2xbf16> to vector<4x2x3xbf16>
+ %flat = vector.shape_cast %transpose : vector<4x2x3xbf16> to vector<8x3xbf16>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %flat, %accum : vector<2x8xbf16>, vector<8x3xbf16> into vector<2x3xbf16>
+ scf.yield %dot : vector<2x3xbf16>
}
- return %result : vector<2x4xbf16>
+ return %result : vector<2x3xbf16>
}
func.func @vnni_i8(%a: memref<1x2x8xi8>,
- %b: memref<1x2x4x4xi8>) -> vector<2x4xi8> {
+ %b: memref<1x2x3x4xi8>) -> vector<2x3xi8> {
%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>) {
+ %init = arith.constant dense<0> : vector<2x3xi8>
+ %result = scf.for %k = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x3xi8>) {
%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>
+ %right = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<1x2x3x4xi8>, vector<2x3x4xi8>
+ %transpose = vector.transpose %right, [0, 2, 1] : vector<2x3x4xi8> to vector<2x4x3xi8>
+ %flat = vector.shape_cast %transpose : vector<2x4x3xi8> to vector<8x3xi8>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %flat, %accum : vector<2x8xi8>, vector<8x3xi8> into vector<2x3xi8>
+ scf.yield %dot : vector<2x3xi8>
}
- return %result : vector<2x4xi8>
+ return %result : vector<2x3xi8>
}
// Check that folding deinterleave/transpose/interleave/transpose unpacking
-// and converting the i32 loop to index preserve the matmul result.
+// preserves the matmul result.
func.func @interleaved_bf16(%a: memref<1x2x8xbf16>,
%b: memref<1x4x6xbf16>) -> vector<2x3xbf16> {
%zero = arith.constant 0 : index
- %lb = arith.constant 0 : i32
- %ub = arith.constant 1 : i32
+ %end = arith.constant 1 : index
%pad = arith.constant 0.0 : bf16
%init = arith.constant dense<0.0> : vector<2x3xbf16>
- %result = scf.for %iv = %lb to %ub step %ub iter_args(%accum = %init) -> (vector<2x3xbf16>) : i32 {
- %k = arith.index_cast %iv : i32 to index
+ %result = scf.for %k = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x3xbf16>) {
%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], %pad {in_bounds = [true, true]} : memref<1x4x6xbf16>, vector<4x6xbf16>
%even, %odd = vector.deinterleave %right : vector<4x6xbf16> -> vector<4x3xbf16>
@@ -71,8 +69,7 @@ func.func @interleaved_bf16(%a: memref<1x2x8xbf16>,
return %result : vector<2x3xbf16>
}
-// Check the equivalent transpose/extract and slice-insert form, including
-// removal of the IV's index -> i32 -> index round-trip.
+// Check the equivalent transpose/extract and slice-insert form.
func.func @split_join_bf16(%a: memref<1x2x8xbf16>,
%b: memref<1x4x6xbf16>) -> vector<2x3xbf16> {
%zero = arith.constant 0 : index
@@ -80,9 +77,7 @@ func.func @split_join_bf16(%a: memref<1x2x8xbf16>,
%pad = arith.constant 0.0 : bf16
%poison = ub.poison : vector<3x4x2xbf16>
%init = arith.constant dense<0.0> : vector<2x3xbf16>
- %result = scf.for %iv = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x3xbf16>) {
- %k32 = arith.index_cast %iv : index to i32
- %k = arith.index_castui %k32 nneg : i32 to index
+ %result = scf.for %k = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x3xbf16>) {
%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], %pad {in_bounds = [true, true]} : memref<1x4x6xbf16>, vector<4x6xbf16>
%packed = vector.shape_cast %right : vector<4x6xbf16> to vector<4x3x2xbf16>
@@ -103,58 +98,126 @@ func.func @split_join_bf16(%a: memref<1x2x8xbf16>,
return %result : vector<2x3xbf16>
}
-func.func @main() {
+// Check i8 unpacking through two levels of binary splits and joins. The first
+// split yields VNNI lanes {0, 2} and {1, 3}; mix deinterleave with
+// transpose/extract, and interleave with out-of-order slice inserts.
+func.func @split_join_i8(%a: memref<1x2x8xi8>,
+ %b: memref<1x2x12xi8>) -> vector<2x3xi8> {
%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
+ %end = arith.constant 1 : index
+ %pad = arith.constant 0 : i8
+ %poison = ub.poison : vector<3x2x2xi8>
+ %init = arith.constant dense<0> : vector<2x3xi8>
+ %result = scf.for %k = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x3xi8>) {
+ %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], %pad {in_bounds = [true, true]} : memref<1x2x12xi8>, vector<2x12xi8>
+ %even, %odd = vector.deinterleave %right : vector<2x12xi8> -> vector<2x6xi8>
+ %lane0, %lane2 = vector.deinterleave %even : vector<2x6xi8> -> vector<2x3xi8>
+ %odd_packed = vector.shape_cast %odd : vector<2x6xi8> to vector<2x3x2xi8>
+ %odd_split = vector.transpose %odd_packed, [2, 0, 1] : vector<2x3x2xi8> to vector<2x2x3xi8>
+ %lane1 = vector.extract %odd_split[0] : vector<2x3xi8> from vector<2x2x3xi8>
+ %lane3 = vector.extract %odd_split[1] : vector<2x3xi8> from vector<2x2x3xi8>
+ %t0 = vector.transpose %lane0, [1, 0] : vector<2x3xi8> to vector<3x2xi8>
+ %t1 = vector.transpose %lane1, [1, 0] : vector<2x3xi8> to vector<3x2xi8>
+ %t2 = vector.transpose %lane2, [1, 0] : vector<2x3xi8> to vector<3x2xi8>
+ %t3 = vector.transpose %lane3, [1, 0] : vector<2x3xi8> to vector<3x2xi8>
+ %even_t = vector.interleave %t0, %t2 : vector<3x2xi8> -> vector<3x4xi8>
+ %s1 = vector.shape_cast %t1 : vector<3x2xi8> to vector<3x2x1xi8>
+ %s3 = vector.shape_cast %t3 : vector<3x2xi8> to vector<3x2x1xi8>
+ %join0 = vector.insert_strided_slice %s3, %poison offsets = [0, 0, 1], strides = [1, 1, 1] : vector<3x2x1xi8> into vector<3x2x2xi8>
+ %join1 = vector.insert_strided_slice %s1, %join0 offsets = [0, 0, 0], strides = [1, 1, 1] : vector<3x2x1xi8> into vector<3x2x2xi8>
+ %odd_t = vector.shape_cast %join1 : vector<3x2x2xi8> to vector<3x4xi8>
+ %joined = vector.interleave %even_t, %odd_t : vector<3x4xi8> -> vector<3x8xi8>
+ %flat = vector.transpose %joined, [1, 0] : vector<3x8xi8> to vector<8x3xi8>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %flat, %accum : vector<2x8xi8>, vector<8x3xi8> into vector<2x3xi8>
+ scf.yield %dot : vector<2x3xi8>
+ }
+ return %result : vector<2x3xi8>
+}
+
+// Check i8 unpacking through a single 4-way transpose/extract split and a
+// 4-way slice-insert join with out-of-order inserts.
+func.func @fourway_split_join_i8(%a: memref<1x2x8xi8>,
+ %b: memref<1x2x12xi8>) -> vector<2x3xi8> {
+ %zero = arith.constant 0 : index
+ %end = arith.constant 1 : index
+ %pad = arith.constant 0 : i8
+ %poison = ub.poison : vector<3x2x4xi8>
+ %init = arith.constant dense<0> : vector<2x3xi8>
+ %result = scf.for %k = %zero to %end step %end iter_args(%accum = %init) -> (vector<2x3xi8>) {
+ %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], %pad {in_bounds = [true, true]} : memref<1x2x12xi8>, vector<2x12xi8>
+ %packed = vector.shape_cast %right : vector<2x12xi8> to vector<2x3x4xi8>
+ %split = vector.transpose %packed, [2, 0, 1] : vector<2x3x4xi8> to vector<4x2x3xi8>
+ %lane0 = vector.extract %split[0] : vector<2x3xi8> from vector<4x2x3xi8>
+ %lane1 = vector.extract %split[1] : vector<2x3xi8> from vector<4x2x3xi8>
+ %lane2 = vector.extract %split[2] : vector<2x3xi8> from vector<4x2x3xi8>
+ %lane3 = vector.extract %split[3] : vector<2x3xi8> from vector<4x2x3xi8>
+ %t0 = vector.transpose %lane0, [1, 0] : vector<2x3xi8> to vector<3x2xi8>
+ %t1 = vector.transpose %lane1, [1, 0] : vector<2x3xi8> to vector<3x2xi8>
+ %t2 = vector.transpose %lane2, [1, 0] : vector<2x3xi8> to vector<3x2xi8>
+ %t3 = vector.transpose %lane3, [1, 0] : vector<2x3xi8> to vector<3x2xi8>
+ %s0 = vector.shape_cast %t0 : vector<3x2xi8> to vector<3x2x1xi8>
+ %s1 = vector.shape_cast %t1 : vector<3x2xi8> to vector<3x2x1xi8>
+ %s2 = vector.shape_cast %t2 : vector<3x2xi8> to vector<3x2x1xi8>
+ %s3 = vector.shape_cast %t3 : vector<3x2xi8> to vector<3x2x1xi8>
+ %join0 = vector.insert_strided_slice %s2, %poison offsets = [0, 0, 2], strides = [1, 1, 1] : vector<3x2x1xi8> into vector<3x2x4xi8>
+ %join1 = vector.insert_strided_slice %s0, %join0 offsets = [0, 0, 0], strides = [1, 1, 1] : vector<3x2x1xi8> into vector<3x2x4xi8>
+ %join2 = vector.insert_strided_slice %s3, %join1 offsets = [0, 0, 3], strides = [1, 1, 1] : vector<3x2x1xi8> into vector<3x2x4xi8>
+ %join3 = vector.insert_strided_slice %s1, %join2 offsets = [0, 0, 1], strides = [1, 1, 1] : vector<3x2x1xi8> into vector<3x2x4xi8>
+ %joined = vector.shape_cast %join3 : vector<3x2x4xi8> to vector<3x8xi8>
+ %flat = vector.transpose %joined, [1, 0] : vector<3x8xi8> to vector<8x3xi8>
+ %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %left, %flat, %accum : vector<2x8xi8>, vector<8x3xi8> into vector<2x3xi8>
+ scf.yield %dot : vector<2x3xi8>
+ }
+ return %result : vector<2x3xi8>
+}
- %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
+func.func @main() {
+ %zero = arith.constant 0 : index
- // Unequal even/odd elements and rectangular shapes expose parity swaps and
- // incorrect reshapes when comparing the original and folded contractions.
- %ai = memref.alloc() : memref<1x2x8xbf16>
- %bi = memref.alloc() : memref<1x4x6xbf16>
- %lhs_i = arith.constant dense<[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
- [2.0, 0.0, 1.0, 3.0, 0.0, 2.0, 1.0, 3.0]]> : vector<2x8xbf16>
- %rhs_i = arith.constant dense<[[0.0, 1.0, 2.0, 3.0, 1.0, 0.0],
- [2.0, 0.0, 3.0, 1.0, 0.0, 2.0],
- [1.0, 3.0, 0.0, 0.0, 2.0, 1.0],
- [3.0, 2.0, 1.0, 0.0, 1.0, 3.0]]> : vector<4x6xbf16>
- vector.transfer_write %lhs_i, %ai[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<2x8xbf16>, memref<1x2x8xbf16>
- vector.transfer_write %rhs_i, %bi[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<4x6xbf16>, memref<1x4x6xbf16>
- %interleaved = func.call @interleaved_bf16(%ai, %bi) : (memref<1x2x8xbf16>, memref<1x4x6xbf16>) -> vector<2x3xbf16>
+ // All functions unpack the same logical 8x3 RHS from either 2-way (bf16) or
+ // 4-way (i8) VNNI layout and multiply it by the same 2x8 LHS. Unequal VNNI
+ // lanes and rectangular shapes expose lane permutations and wrong reshapes.
+ %af = memref.alloc() : memref<1x2x8xbf16>
+ %bf = memref.alloc() : memref<1x4x6xbf16>
+ %bf_packed = memref.expand_shape %bf [[0], [1], [2, 3]] output_shape [1, 4, 3, 2] : memref<1x4x6xbf16> into memref<1x4x3x2xbf16>
+ %lhs_bf = arith.constant dense<[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
+ [2.0, 0.0, 1.0, 3.0, 0.0, 2.0, 1.0, 3.0]]> : vector<2x8xbf16>
+ %rhs_bf = arith.constant dense<[[0.0, 1.0, 2.0, 3.0, 1.0, 0.0],
+ [2.0, 0.0, 3.0, 1.0, 0.0, 2.0],
+ [1.0, 3.0, 0.0, 0.0, 2.0, 1.0],
+ [3.0, 2.0, 1.0, 0.0, 1.0, 3.0]]> : vector<4x6xbf16>
+ vector.transfer_write %lhs_bf, %af[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<2x8xbf16>, memref<1x2x8xbf16>
+ vector.transfer_write %rhs_bf, %bf[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<4x6xbf16>, memref<1x4x6xbf16>
+ %vnni_bf16 = func.call @vnni_bf16(%af, %bf_packed) : (memref<1x2x8xbf16>, memref<1x4x3x2xbf16>) -> vector<2x3xbf16>
+ vector.print %vnni_bf16 : vector<2x3xbf16>
+ // CHECK: ( ( 68, 28, 56 ), ( 17, 11, 20 ) )
+ %interleaved = func.call @interleaved_bf16(%af, %bf) : (memref<1x2x8xbf16>, memref<1x4x6xbf16>) -> vector<2x3xbf16>
vector.print %interleaved : vector<2x3xbf16>
// CHECK-NEXT: ( ( 68, 28, 56 ), ( 17, 11, 20 ) )
- %split_join = func.call @split_join_bf16(%ai, %bi) : (memref<1x2x8xbf16>, memref<1x4x6xbf16>) -> vector<2x3xbf16>
+ %split_join = func.call @split_join_bf16(%af, %bf) : (memref<1x2x8xbf16>, memref<1x4x6xbf16>) -> vector<2x3xbf16>
vector.print %split_join : vector<2x3xbf16>
// CHECK-NEXT: ( ( 68, 28, 56 ), ( 17, 11, 20 ) )
+
+ %ai = memref.alloc() : memref<1x2x8xi8>
+ %bi = memref.alloc() : memref<1x2x12xi8>
+ %bi_packed = memref.expand_shape %bi [[0], [1], [2, 3]] output_shape [1, 2, 3, 4] : memref<1x2x12xi8> into memref<1x2x3x4xi8>
+ %lhs_i = arith.constant dense<[[1, 2, 3, 4, 5, 6, 7, 8],
+ [2, 0, 1, 3, 0, 2, 1, 3]]> : vector<2x8xi8>
+ %rhs_i = arith.constant dense<[[0, 1, 2, 0, 2, 3, 3, 1, 1, 0, 0, 2],
+ [1, 3, 3, 2, 0, 0, 1, 0, 2, 1, 1, 3]]> : vector<2x12xi8>
+ vector.transfer_write %lhs_i, %ai[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<2x8xi8>, memref<1x2x8xi8>
+ vector.transfer_write %rhs_i, %bi[%zero, %zero, %zero] {in_bounds = [true, true]} : vector<2x12xi8>, memref<1x2x12xi8>
+ %vnni_i8 = func.call @vnni_i8(%ai, %bi_packed) : (memref<1x2x8xi8>, memref<1x2x3x4xi8>) -> vector<2x3xi8>
+ vector.print %vnni_i8 : vector<2x3xi8>
+ // CHECK-NEXT: ( ( 68, 28, 56 ), ( 17, 11, 20 ) )
+ %split_join_i8 = func.call @split_join_i8(%ai, %bi) : (memref<1x2x8xi8>, memref<1x2x12xi8>) -> vector<2x3xi8>
+ vector.print %split_join_i8 : vector<2x3xi8>
+ // CHECK-NEXT: ( ( 68, 28, 56 ), ( 17, 11, 20 ) )
+ %fourway_i8 = func.call @fourway_split_join_i8(%ai, %bi) : (memref<1x2x8xi8>, memref<1x2x12xi8>) -> vector<2x3xi8>
+ vector.print %fourway_i8 : vector<2x3xi8>
+ // CHECK-NEXT: ( ( 68, 28, 56 ), ( 17, 11, 20 ) )
return
}
More information about the Mlir-commits
mailing list