[Mlir-commits] [mlir] [mlir][x86] Add pattern to tile and unroll vector.contract (PR #229076)
Julian Oppermann
llvmlistbot at llvm.org
Fri Oct 9 01:30:25 PDT 2026
================
@@ -0,0 +1,817 @@
+//===- VectorContractMultiLevelUnroll.cpp ---------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/Dialect/X86/Transforms.h"
+#include "mlir/Dialect/X86/Utils/X86Utils.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
+
+#define DEBUG_TYPE "x86-vector-contract-multi-level-unroll"
+
+using namespace mlir;
+
+using llvm::all_of;
+using llvm::count;
+
+// Discardable attribute to control vector unroll patterns.
+static constexpr auto nativeShapeAttrName = "x86_vcmlu_native_shape";
+
+namespace {
+// Shared state for the pattern match and IR transformation.
+struct MLUCandidate {
+ // The CPU feature flag to target with the transformation, in the spelling
+ // used by llvm/include/llvm/TargetParser/X86TargetParser.def, e.g.
+ // "amx-bf16". Determines native shapes and supported data types for operands
+ // and the accumulator.
+ StringRef target;
+
+ // Anchor operation.
+ vector::ContractionOp contract;
+ bool isInVnniLayout = false;
+
+ // Shape of the original contraction.
+ VectorType accType;
+ VectorType lhsType;
+ VectorType rhsType;
+
+ // Shape of the register-tiled contraction. We'll insert a loop nest to cover
+ // the original shape.
+ VectorType accRegTileType;
+ VectorType lhsRegTileType;
+ VectorType rhsRegTileType;
+
+ // Shape of the native contraction (for the X86 nanokernel patterns). We'll
+ // unroll the register-tiled contraction to this shape.
+ VectorType accNativeType;
+ VectorType lhsNativeType;
+ VectorType rhsNativeType;
+ SmallVector<int64_t> contractNativeShape;
+
+ // Bookkeeping for the accumulation loop and operands.
+ scf::ForOp accLoop;
+ Value accInitVal;
+
+ bool accIsZeroInit = false;
+ vector::TransferReadOp lhsRead;
+ vector::TransferReadOp rhsRead;
+
+ // Memory backing the accumulator.
+ Value accMemIn;
+ Value accMemOut;
+
+ // Reconstructed accumulation loop with native-shaped iter_args.
+ scf::ForOp slicedAccLoop;
+
+ // Keep track of operations that should be deleted after transformation
+ // (depending on which ops could be integrated into the new loop nest).
+ SmallVector<Operation *> opsToDelete;
+};
+} // anonymous namespace
+
+//----------------------------------------------------------------------------//
+// Utilities //
+//----------------------------------------------------------------------------//
+
+// Return true if the argument is an all-zero vector.
+static bool isZeroVector(Value v) {
+ auto constOp = v.getDefiningOp<arith::ConstantOp>();
+ if (!constOp)
+ return false;
+ auto denseAttr = dyn_cast<DenseElementsAttr>(constOp.getValue());
+ if (!denseAttr || !denseAttr.isSplat())
+ return false;
+
+ return isZeroIntegerOrFloat(denseAttr.getSplatValue<Attribute>());
+}
+
+// Return true if given vector.transfer_read/write has a default permutation
+// map, no mask and no out-of-bounds check.
+static bool isSimpleTransferOp(VectorTransferOpInterface op) {
+ return op.getPermutationMap().isMinorIdentity() && !op.getMask() &&
+ !op.hasOutOfBoundsDim();
+}
+
+// Create a rank-reducing subview of a memref that matches the target vector
+// type, e.g. with targetType = vector<64x32xbf16>:
+//
+// memref.subview %A[%block_m, %block_k, 0, 0] [1, 1, 64, 32] [1, 1, 1, 1]
+// : memref<?x?x64x32xbf16> to
+// memref<64x32xbf16, strided<[32, 1], offset: ?>>
+static Value makeRankReducingSubview(OpBuilder &b, Location loc, Value base,
+ ValueRange indices,
+ VectorType targetType) {
+ auto *ctx = b.getContext();
+
+ auto baseType = dyn_cast<MemRefType>(base.getType());
+ assert(baseType && "base must be a MemRefType");
+
+ int64_t baseRank = baseType.getRank();
+ int64_t targetRank = targetType.getRank();
+ assert(targetRank <= baseRank && "target shape's rank > memref's rank");
+
+ ArrayRef<int64_t> targetShape = targetType.getShape();
+
+ SmallVector<int64_t> sizesInt(baseType.getRank(), 1);
+ copy(targetShape, sizesInt.begin() + baseRank - targetRank);
+ SmallVector<int64_t> stridesInt(baseRank, 1);
+
+ auto offsets = getAsOpFoldResult(indices);
+ auto sizes = getAsIndexOpFoldResult(ctx, sizesInt);
+ auto strides = getAsIndexOpFoldResult(ctx, stridesInt);
+
+ MemRefType rankReducedType = memref::SubViewOp::inferRankReducedResultType(
+ targetShape, baseType, offsets, sizes, strides);
+ return memref::SubViewOp::create(b, loc, rankReducedType, base, offsets,
+ sizes, strides);
+}
+
+// Slice a 2D vector into smaller tiles of the given native type.
+static SmallVector<Value> slice(OpBuilder &b, Location loc, Value value,
+ VectorType nativeType) {
+ VectorType regTileType = cast<VectorType>(value.getType());
+ ArrayRef<int64_t> regTileShape = regTileType.getShape();
+ ArrayRef<int64_t> nativeShape = nativeType.getShape();
+
+ assert(regTileShape.size() == 2 && nativeShape.size() == 2 &&
+ "Only 2D shapes are supported");
+
+ int64_t regTileM = regTileShape[0];
+ int64_t regTileN = regTileShape[1];
+ int64_t nativeM = nativeShape[0];
+ int64_t nativeN = nativeShape[1];
+
+ SmallVector<Value> result;
+ result.reserve((regTileM / nativeM) * (regTileN / nativeN));
+
+ for (int64_t m = 0; m < regTileM; m += nativeM)
+ for (int64_t n = 0; n < regTileN; n += nativeN)
+ result.push_back(vector::ExtractStridedSliceOp::create(
+ b, loc, value, ArrayRef<int64_t>{m, n}, nativeShape,
+ ArrayRef<int64_t>{1, 1}));
+
+ return result;
+}
+
+// Splice smaller tiles back into a 2D vector of the given register tile type.
+static Value splice(OpBuilder &b, Location loc, ValueRange values,
+ VectorType regTileType) {
+ auto types = values.getTypes();
+ assert(!types.empty() &&
+ count(types, types.front()) == static_cast<int64_t>(types.size()) &&
+ "All values must have the same type");
+
+ VectorType nativeType = cast<VectorType>(types.front());
+
+ ArrayRef<int64_t> nativeShape = nativeType.getShape();
+ ArrayRef<int64_t> regTileShape = regTileType.getShape();
+
+ assert(nativeShape.size() == 2 && regTileShape.size() == 2 &&
+ "Only 2D shapes are supported");
+
+ int64_t nativeM = nativeShape[0];
+ int64_t nativeN = nativeShape[1];
+ int64_t regTileM = regTileShape[0];
+ int64_t regTileN = regTileShape[1];
+
+ assert((regTileM / nativeM) * (regTileN / nativeN) ==
+ static_cast<int64_t>(values.size()) &&
+ "The number of values must match the number of slices");
+
+ Value result = ub::PoisonOp::create(b, loc, regTileType);
+
+ auto valueIt = values.begin();
+ for (int64_t m = 0; m < regTileM; m += nativeM)
+ for (int64_t n = 0; n < regTileN; n += nativeN) {
+ result = vector::InsertStridedSliceOp::create(b, loc, *valueIt, result,
+ ArrayRef<int64_t>{m, n},
+ ArrayRef<int64_t>{1, 1});
+ ++valueIt;
+ }
+
+ return result;
+}
+
+static LogicalResult hasNativeShapeAttr(Operation *op) {
+ return success(op->hasAttr(nativeShapeAttrName));
+}
+
+static std::optional<SmallVector<int64_t>>
+getFromNativeShapeAttr(Operation *op) {
+ auto attr = op->getAttrOfType<DenseI64ArrayAttr>(nativeShapeAttrName);
+ if (!attr)
+ return std::nullopt;
+ return SmallVector<int64_t>(attr.asArrayRef());
+}
+
+static void setNativeShapeAttr(VectorTransferOpInterface op,
+ VectorType nativeType) {
+ if (op.getVectorType() != nativeType)
+ op->setAttr(
+ nativeShapeAttrName,
+ DenseI64ArrayAttr::get(nativeType.getContext(), nativeType.getShape()));
+}
+
+//----------------------------------------------------------------------------//
+// Pattern match and target shape computation //
+//----------------------------------------------------------------------------//
+
+// Establish whether the candidate contraction is a canonical matrix
+// multiplication on flat or VNNI-packed vectors. If successful, the discovered
+// operand and accumulator types are recorded in the candidate struct.
+//
+// Example (flat):
+// vector.contract {
+// kind = #vector.kind<add>, indexing_maps = [
+// affine_map<(dm, dn, dk) -> (dm, dk)>,
+// affine_map<(dm, dn, dk) -> (dk, dn)>,
+// affine_map<(dm, dn, dk) -> (dm, dn)>],
+// iterator_types = ["parallel", "parallel", "reduction"],
+// } %lhs, %rhs, %acc : !lhsType, !rhsType into !accType
+//
+// Example (VNNI-packed):
+// vector.contract {
+// kind = #vector.kind<add>, indexing_maps = [
+// affine_map<(dm, dn, dk, dvnni) -> (dm, dk, dvnni)>,
+// affine_map<(dm, dn, dk, dvnni) -> (dk, dn, dvnni)>,
+// affine_map<(dm, dn, dk, dvnni) -> (dm, dn)>],
+// iterator_types = ["parallel", "parallel", "reduction", "reduction"],
+// } %lhs, %rhs, %acc : !lhsType, !rhsType into !accType
+static LogicalResult matchCanonicalMatmul(MLUCandidate &candidate,
+ PatternRewriter &rewriter) {
+ vector::ContractionOp contract = candidate.contract;
+ if (contract.getKind() != vector::CombiningKind::ADD)
+ return rewriter.notifyMatchFailure(contract,
+ "not using ADD combining kind");
+
+ VectorType accType = dyn_cast<VectorType>(contract.getAccType());
+ if (!accType || accType.getRank() != 2)
+ return rewriter.notifyMatchFailure(contract,
+ "accumulator is not a 2D vector");
+
+ VectorType lhsType = contract.getLhsType();
+ VectorType rhsType = contract.getRhsType();
+ if (lhsType.getElementType() != rhsType.getElementType())
+ return rewriter.notifyMatchFailure(
+ contract, "LHS and RHS element types do not match");
+
+ auto iteratorTypes = contract.getIteratorTypesArray();
+ auto indexingMaps = contract.getIndexingMapsArray();
+ assert(indexingMaps.size() == 3 && "expected 3 indexing maps");
+ bool isInVnniLayout = x86::isInVnniLayout(contract, indexingMaps);
+
+ auto par = vector::IteratorType::parallel;
+ auto red = vector::IteratorType::reduction;
+ auto m = rewriter.getAffineDimExpr(0);
+ auto n = rewriter.getAffineDimExpr(1);
+ auto k = rewriter.getAffineDimExpr(2);
+ auto vnni = rewriter.getAffineDimExpr(3);
+
+ if (!isInVnniLayout) {
+ if (!equal(indexingMaps[0].getResults(), ArrayRef{m, k}) ||
+ !equal(indexingMaps[1].getResults(), ArrayRef{k, n}) ||
+ !equal(indexingMaps[2].getResults(), ArrayRef{m, n}))
+ return rewriter.notifyMatchFailure(contract,
+ "indexing maps are not canonical");
+
+ if (!equal(iteratorTypes, ArrayRef{par, par, red}))
+ return rewriter.notifyMatchFailure(contract,
+ "iterator types are not canonical");
+ } else {
+ if (!equal(indexingMaps[0].getResults(), ArrayRef{m, k, vnni}) ||
+ !equal(indexingMaps[1].getResults(), ArrayRef{k, n, vnni}) ||
+ !equal(indexingMaps[2].getResults(), ArrayRef{m, n}))
+ return rewriter.notifyMatchFailure(
+ contract, "indexing maps are not canonical (VNNI)");
+
+ if (!equal(iteratorTypes, ArrayRef{par, par, red, red}))
+ return rewriter.notifyMatchFailure(
+ contract, "iterator types are not canonical (VNNI)");
+ }
+
+ candidate.isInVnniLayout = isInVnniLayout;
+ candidate.accType = accType;
+ candidate.lhsType = lhsType;
+ candidate.rhsType = rhsType;
+
+ return success();
+}
+
+// Check compatibility of datatypes with the given target ISA extension, and try
+// to determine a suitable tiling strategy for it. If successful, the strategy
+// is encoded in the *RegTileType and *NativeType fields of the candidate
+// struct.
+static LogicalResult matchShapesAndTypes(MLUCandidate &candidate,
+ PatternRewriter &rewriter) {
+ vector::ContractionOp contract = candidate.contract;
+ StringRef target = candidate.target;
+
+ Type accElemType = candidate.accType.getElementType();
+ Type inpElemType = candidate.lhsType.getElementType();
+
+ ArrayRef<int64_t> accShape = candidate.accType.getShape();
+ ArrayRef<int64_t> lhsShape = candidate.lhsType.getShape();
+
+ int64_t origM = accShape[0];
+ int64_t origN = accShape[1];
+ int64_t origK = lhsShape[1];
+ int64_t origVnni = candidate.isInVnniLayout ? lhsShape[2] : -1;
+
+ int64_t regTileM, regTileN, regTileK;
+ int64_t nativeM, nativeN, nativeK, vnni;
+
+ if (target.starts_with("amx")) {
+ if (!((target == "amx-bf16" && inpElemType.isBF16() &&
+ accElemType.isF32()) ||
+ (target == "amx-int8" && inpElemType.isInteger(8) &&
+ accElemType.isInteger(32))))
+ return rewriter.notifyMatchFailure(
+ contract, "unsupported combination of input and accumulator types");
+
+ nativeM = nativeN = 16;
+ vnni = 32 / inpElemType.getIntOrFloatBitWidth();
+
+ if (candidate.isInVnniLayout) {
+ nativeK = 16;
+ if (origM % nativeM != 0 || origN % nativeN != 0)
+ return rewriter.notifyMatchFailure(
+ contract, "vector shape cannot be cleanly unrolled");
+ if (origK != nativeK || origVnni != vnni)
+ return rewriter.notifyMatchFailure(
+ contract, "K dimension or VNNI factor mismatch");
+ if (origM == nativeM && origN == nativeN)
+ return rewriter.notifyMatchFailure(contract,
+ "already in native tile size");
+ // Prefer 2x2 tiling if shapes allow.
+ regTileM = origM % (2 * nativeM) == 0 ? 2 * nativeM : nativeM;
+ regTileN = origN % (2 * nativeN) == 0 ? 2 * nativeN : nativeN;
+ regTileK = nativeK;
+ } else {
+ // Online packing requires a 2x2 tile register layout.
+ nativeK = 16 * vnni;
+ regTileM = 2 * nativeM;
----------------
jopperm wrote:
Good catch, added with a new test.
https://github.com/llvm/llvm-project/pull/229076
More information about the Mlir-commits
mailing list