Google Git Google Git
Sign in
llvm/llvm-project/792bf8ef030d7c2aa12abbab7662ffbcbe2d8b21/./mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
blob: 6b0c09b3b48f4b4df9ce9c4ec8c6e58fd1f9b94d [file]
//===- VectorToXeGPU.cpp - Convert vector to XeGPU dialect ------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements lowering of vector operations to XeGPU dialect ops.
//
//===----------------------------------------------------------------------===//
#include "mlir/Conversion/VectorToXeGPU/VectorToXeGPU.h"
#include "mlir/Conversion/VectorToGPU/VectorToGPU.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/TypeSwitch.h"
#include <algorithm>
#include <optional>
namespace mlir {
#define GEN_PASS_DEF_CONVERTVECTORTOXEGPU
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
using namespace mlir;
namespace {
// Return true if value represents a zero constant.
static bool isZeroConstant(Value val) {
auto constant = val.getDefiningOp<arith::ConstantOp>();
if (!constant)
return false;
return TypeSwitch<Attribute, bool>(constant.getValue())
.Case([](FloatAttr floatAttr) { return floatAttr.getValue().isZero(); })
.Case([](IntegerAttr intAttr) { return intAttr.getValue().isZero(); })
.Default(false);
}
// Return true if the transfer padding value is compatible with the implicit
// padding of an nd block load. LoadNdOp fills out-of-bounds elements with zero,
// so a zero constant matches its semantics exactly. A poison padding means the
// out-of-bounds elements are "don't care", so any implicit padding (including
// zero) is also acceptable.
static bool isZeroOrPoisonPadding(Value val) {
return isZeroConstant(val) || val.getDefiningOp<ub::PoisonOp>();
}
// Return true if the permutation map keeps every dimension in place except the
// innermost two, which are swapped, e.g.:
// (d0, d1) -> (d1, d0)
// (d0, d1, d2) -> (d2, d1)
// (d0, d1, d2, d3) -> (d0, d1, d3, d2)
// This is the only non-identity permutation an nd block load can realize (by
// loading the untransposed block and applying a trailing vector.transpose).
static bool isInnermostTwoDimsTransposed(AffineMap map) {
unsigned numResults = map.getNumResults();
if (numResults < 2)
return false;
MLIRContext *ctx = map.getContext();
unsigned numInputs = map.getNumInputs();
// All but the innermost two results must match the minor-identity map.
for (unsigned i = 0; i + 2 < numResults; ++i)
if (map.getResult(i) != getAffineDimExpr(numInputs - numResults + i, ctx))
return false;
// The innermost two results must be the last two input dims, swapped.
return map.getResult(numResults - 2) ==
getAffineDimExpr(numInputs - 1, ctx) &&
map.getResult(numResults - 1) == getAffineDimExpr(numInputs - 2, ctx);
}
// Return true if `uArch` can transfer a `shape`-shaped tile of `elemTy`
// elements with the subgroup 2D block instruction `instKind`.
//
// A 2D block instruction accesses the two innermost dimensions of the tile;
// leading dimensions are unrolled into a sequence of 2D accesses. The
// innermost 2D tile does not have to match a hardware block exactly, because
// the later XeGPU passes (layout propagation and blocking) split it into
// hardware-sized blocks - but that split only exists when each of the two
// innermost extents is a multiple of a supported block extent. Transfers
// without such a split cannot be expressed with block instructions at all and
// must use a different lowering.
static bool isSupportedBlockShape(const xegpu::uArch::uArch *uArch,
xegpu::uArch::InstructionKind instKind,
ArrayRef<int64_t> shape, Type elemTy,
bool hasTranspose = false) {
if (shape.size() < 2)
return false;
if (!uArch || !uArch->isSupportedInstruction(instKind))
return false;
const auto *blockInst = dyn_cast<xegpu::uArch::BlockIOInstructionInterface>(
uArch->getInstruction(instKind));
if (!blockInst)
return false;
int width = static_cast<int>(shape.back());
int height = static_cast<int>(shape[shape.size() - 2]);
// A tile is accessible when both of its extents are multiples of a supported
// block extent, so that the later XeGPU passes can split it into hardware
// blocks. A missing entry means this variant does not support the element
// type.
auto fitsBlockShapes = [&](bool hasTransform) {
std::optional<xegpu::uArch::BlockIOInstructionInterface::BlockShapes>
blockShapes = blockInst->getBlockWidthHeightCount(elemTy, hasTransform,
hasTranspose);
if (!blockShapes)
return false;
auto [widths, heights, counts] = *blockShapes;
return xegpu::getLargestDivisor(width, widths) != -1 &&
xegpu::getLargestDivisor(height, heights) != -1;
};
// Whether a tile is eventually loaded with the transformed (VNNI) variant is
// only decided later, when layout propagation knows whether it feeds the B
// operand of a matrix operation, so accept a tile that either variant can
// access.
return fitsBlockShapes(/*hasTransform=*/false) ||
fitsBlockShapes(/*hasTransform=*/true);
}
static LogicalResult transferPreconditions(PatternRewriter &rewriter,
VectorTransferOpInterface xferOp) {
if (xferOp.getMask())
return rewriter.notifyMatchFailure(xferOp,
"Masked transfer is not supported");
auto srcTy = dyn_cast<MemRefType>(xferOp.getShapedType());
if (!srcTy)
return rewriter.notifyMatchFailure(xferOp, "Expects memref source");
// Validate further transfer op semantics.
SmallVector<int64_t> strides;
int64_t offset;
if (failed(srcTy.getStridesAndOffset(strides, offset)))
return rewriter.notifyMatchFailure(xferOp,
"The memref strides cannot be inferred");
if (strides.empty())
return rewriter.notifyMatchFailure(xferOp, "0D memref is not supported");
if (strides.back() != 1)
return rewriter.notifyMatchFailure(
xferOp, "Buffer must be contiguous in the innermost dimension");
VectorType vecTy = xferOp.getVectorType();
unsigned vecRank = vecTy.getRank();
if (vecRank == 0)
return rewriter.notifyMatchFailure(xferOp, "0D vectors are not supported");
AffineMap map = xferOp.getPermutationMap();
if (!map.isProjectedPermutation(/*allowZeroInResults=*/false))
return rewriter.notifyMatchFailure(xferOp, "Unsupported permutation map");
unsigned numInputDims = map.getNumInputs();
for (AffineExpr expr : map.getResults().take_back(vecRank)) {
auto dim = dyn_cast<AffineDimExpr>(expr);
if (dim.getPosition() < (numInputDims - vecRank))
return rewriter.notifyMatchFailure(
xferOp, "Only the innermost dimensions can be accessed");
}
return success();
}
// Adjusts the strides of a memref according to a given permutation map for
// vector operations.
//
// This function updates the innermost strides in the `strides` array to
// reflect the permutation specified by `permMap`. The permutation is computed
// using the inverse and broadcasting-aware version of the permutation map,
// and is applied to the relevant strides. This ensures that memory accesses
// are consistent with the logical permutation of vector elements.
//
// Example:
// Suppose we have a memref of rank 4 with strides `[s0, s1, s2, s3]`.
// If the permutation map swaps the last two dimensions (e.g., [0, 1] -> [1,
// 0]), then after calling this function, the last two strides will be
// swapped:
// Original strides: [s0, s1, s2, s3]
// After permutation: [s0, s1, s3, s2]
//
static void adjustStridesForPermutation(AffineMap permMap,
SmallVectorImpl<Value> &strides) {
AffineMap invMap = inverseAndBroadcastProjectedPermutation(permMap);
SmallVector<unsigned> perms;
invMap.isPermutationOfMinorIdentityWithBroadcasting(perms);
SmallVector<int64_t> perms64(perms.begin(), perms.end());
strides = applyPermutation(strides, perms64);
}
// Computes memory strides and a memref offset for vector transfer operations,
// handling both static and dynamic memrefs while applying permutation
// transformations for XeGPU lowering.
template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
std::decay_t<OpType>, vector::TransferReadOp, vector::TransferWriteOp,
vector::GatherOp, vector::ScatterOp>::value>>
static std::pair<SmallVector<Value>, Value>
computeMemrefMeta(OpType xferOp, PatternRewriter &rewriter) {
SmallVector<Value> strides;
Value baseMemref = xferOp.getBase();
MemRefType memrefType = dyn_cast<MemRefType>(baseMemref.getType());
Location loc = xferOp.getLoc();
Value offsetVal = nullptr;
if (memrefType.hasStaticShape()) {
int64_t offset;
SmallVector<int64_t> intStrides;
if (failed(memrefType.getStridesAndOffset(intStrides, offset)))
return {{}, offsetVal};
bool hasDynamicStrides = llvm::any_of(intStrides, [](int64_t strideVal) {
return ShapedType::isDynamic(strideVal);
});
if (!hasDynamicStrides)
for (int64_t s : intStrides)
strides.push_back(arith::ConstantIndexOp::create(rewriter, loc, s));
if (!ShapedType::isDynamic(offset))
offsetVal = arith::ConstantIndexOp::create(rewriter, loc, offset);
}
if (strides.empty() || !offsetVal) {
// For dynamic shape memref, use memref.extract_strided_metadata to get
// stride values
unsigned rank = memrefType.getRank();
Type indexType = rewriter.getIndexType();
// Result types: [base_memref, offset, stride0, stride1, ..., strideN-1,
// size0, size1, ..., sizeN-1]
SmallVector<Type> resultTypes;
resultTypes.push_back(MemRefType::get(
{}, memrefType.getElementType())); // base memref (unranked)
resultTypes.push_back(indexType); // offset
for (unsigned i = 0; i < rank; ++i)
resultTypes.push_back(indexType); // strides
for (unsigned i = 0; i < rank; ++i)
resultTypes.push_back(indexType); // sizes
auto meta = memref::ExtractStridedMetadataOp::create(
rewriter, loc, resultTypes, baseMemref);
if (strides.empty())
strides.append(meta.getStrides().begin(), meta.getStrides().end());
if (!offsetVal)
offsetVal = meta.getOffset();
}
// Strides are returned in original memref order; permutation is applied in
// computeOffsets only where offsets are indexed in vector order.
return {strides, offsetVal};
}
// Adds the transfer's indices to `baseOffset`, the memref's own element offset:
// the element at `indices` sits at `baseOffset + sum(indices[d] * strides[d])`.
static Value computeBaseOffset(VectorTransferOpInterface xferOp,
PatternRewriter &rewriter,
ArrayRef<Value> strides, Value baseOffset) {
Location loc = xferOp.getLoc();
for (auto [index, stride] : llvm::zip_equal(xferOp.getIndices(), strides)) {
Value contrib = arith::MulIOp::create(rewriter, loc, index, stride);
baseOffset = arith::AddIOp::create(rewriter, loc, baseOffset, contrib);
}
return baseOffset;
}
// This function compute the vectors of localOffsets for scattered load/stores.
// It is used in the lowering of vector.transfer_read/write to
// load_gather/store_scatter Example:
// %0 = vector.transfer_read %expand_shape[%block_id_y, %c0, %c0, %c0, %c0],
// %cst {in_bounds = [true, true, true, true]}>} :
// memref<8x4x2x6x32xbf16>, vector<4x2x6x32xbf16>
//
// %6 = vector.step: vector<4xindex>
// %7 = vector.step: vector<2xindex>
// %8 = vector.step: vector<6xindex>
// %9 = vector.step: vector<32xindex>
// %10 = arith.mul %6, 384
// %11 = arith.mul %7, 192
// %12 = arith.mul %8, 32
// %13 = arith.mul %9, 1
// %14 = vector.shape_cast %10: vector<4xindex> -> vector<4x1x1x1xbf16>
// %15 = vector.shape_cast %11: vector<2xindex> -> vector<1x2x1x1xbf16>
// %16 = vector.shape_cast %12: vector<6xindex> -> vector<1x1x6x1xbf16>
// %17 = vector.shape_cast %13: vector<32xindex> -> vector<1x1x1x32xbf16>
// %18 = vector.broadcast %14: vector<4x1x1x1xbf16> -> vector<4x2x6x32xindex>
// %19 = vector.broadcast %15: vector<1x2x1x1xbf16> -> vector<4x2x6x32xindex>
// %20 = vector.broadcast %16: vector<1x1x6x1xbf16> -> vector<4x2x6x32xindex>
// %21 = vector.broadcast %17: vector<1x1x1x32xbf16> -> vector<4x2x6x32xindex>
// %22 = arith.add %18, %19
// %23 = arith.add %20, %21
// %local_offsets = arith.add %22, %23
// %orig_offset = %block_id_y * 4x2x6x32 // consider using affine map
// %offsets = memref_offset + orig_offset + local_offsets
static Value computeOffsets(VectorTransferOpInterface xferOp,
PatternRewriter &rewriter, ArrayRef<Value> strides,
Value baseOffset) {
Location loc = xferOp.getLoc();
VectorType vectorType = xferOp.getVectorType();
ArrayRef<int64_t> vectorShape = vectorType.getShape();
// Create vector.step operations for each dimension
SmallVector<Value> stepVectors;
llvm::map_to_vector(vectorShape, [&](int64_t dim) {
auto stepType = VectorType::get({dim}, rewriter.getIndexType());
auto stepOp = vector::StepOp::create(rewriter, loc, stepType);
stepVectors.push_back(stepOp);
return stepOp;
});
// Local offsets are indexed in vector order, so permute strides; the base
// offset below uses the original memref-order strides.
SmallVector<Value> permutedStrides(strides.begin(), strides.end());
adjustStridesForPermutation(xferOp.getPermutationMap(), permutedStrides);
// Multiply step vectors by corresponding strides
size_t memrefRank = permutedStrides.size();
size_t vectorRank = vectorShape.size();
SmallVector<Value> strideMultiplied;
for (size_t i = 0; i < vectorRank; ++i) {
size_t memrefDim = memrefRank - vectorRank + i;
Value strideValue = permutedStrides[memrefDim];
auto mulType = dyn_cast<VectorType>(stepVectors[i].getType());
auto bcastOp =
vector::BroadcastOp::create(rewriter, loc, mulType, strideValue);
auto mulOp = arith::MulIOp::create(rewriter, loc, stepVectors[i], bcastOp);
strideMultiplied.push_back(mulOp);
}
// Shape cast each multiplied vector to add singleton dimensions
SmallVector<Value> shapeCasted;
for (size_t i = 0; i < vectorRank; ++i) {
SmallVector<int64_t> newShape(vectorRank, 1);
newShape[i] = vectorShape[i];
auto newType = VectorType::get(newShape, rewriter.getIndexType());
auto castOp = vector::ShapeCastOp::create(rewriter, loc, newType,
strideMultiplied[i]);
shapeCasted.push_back(castOp);
}
// Broadcast each shape-casted vector to full vector shape
SmallVector<Value> broadcasted;
auto fullIndexVectorType =
VectorType::get(vectorShape, rewriter.getIndexType());
for (Value shapeCastVal : shapeCasted) {
auto broadcastOp = vector::BroadcastOp::create(
rewriter, loc, fullIndexVectorType, shapeCastVal);
broadcasted.push_back(broadcastOp);
}
// Add all broadcasted vectors together to compute local offsets
Value localOffsets = broadcasted[0];
for (size_t i = 1; i < broadcasted.size(); ++i)
localOffsets =
arith::AddIOp::create(rewriter, loc, localOffsets, broadcasted[i]);
// Broadcast base offset to match vector shape
baseOffset = computeBaseOffset(xferOp, rewriter, strides, baseOffset);
Value bcastBase = vector::BroadcastOp::create(
rewriter, loc, fullIndexVectorType, baseOffset);
localOffsets = arith::AddIOp::create(rewriter, loc, bcastBase, localOffsets);
return localOffsets;
}
// Returns the size of memref dimension `dim` as a value.
static Value getMemrefDimSize(VectorTransferOpInterface xferOp, unsigned dim,
PatternRewriter &rewriter) {
Location loc = xferOp.getLoc();
auto memrefTy = cast<MemRefType>(xferOp.getShapedType());
if (memrefTy.isDynamicDim(dim))
return memref::DimOp::create(rewriter, loc, xferOp.getBase(), dim)
.getResult();
return arith::ConstantIndexOp::create(rewriter, loc, memrefTy.getDimSize(dim))
.getResult();
}
// Builds the mask a scattered access needs to stay within the source bounds.
//
// The element at vector position `i` along vector dimension `v` accesses memref
// dimension `d` (the one `v` maps to) at `indices[d] + i`, so it is in bounds
// iff `i < dim(d) - indices[d]`. Dimensions the transfer declares in-bounds are
// skipped; non-vector dimensions are always in bounds by the op's contract. The
// per-dimension masks are spread over the full vector shape and combined, the
// same way `computeOffsets` spreads the per-dimension offsets.
//
// Example, for a `vector<8xf32>` read of a `memref<?xf32>` at `%off`:
// %dim = memref.dim %src, %c0
// %limit = arith.subi %dim, %off
// %step = vector.step : vector<8xindex>
// %bcast = vector.broadcast %limit : index to vector<8xindex>
// %mask = arith.cmpi slt, %step, %bcast : vector<8xindex>
static Value computeInBoundsMask(VectorTransferOpInterface xferOp,
PatternRewriter &rewriter) {
Location loc = xferOp.getLoc();
ArrayRef<int64_t> vectorShape = xferOp.getVectorType().getShape();
auto maskType = VectorType::get(vectorShape, rewriter.getI1Type());
AffineMap map = xferOp.getPermutationMap();
OperandRange indices = xferOp.getIndices();
Value mask;
for (unsigned v = 0, e = vectorShape.size(); v < e; ++v) {
if (xferOp.isDimInBounds(v))
continue;
unsigned d = cast<AffineDimExpr>(map.getResult(v)).getPosition();
Value bound = getMemrefDimSize(xferOp, d, rewriter);
// A negative limit masks the dimension off entirely, which is what a
// starting index past the end of the dimension should do.
Value limit = arith::SubIOp::create(rewriter, loc, bound, indices[d]);
auto stepType = VectorType::get({vectorShape[v]}, rewriter.getIndexType());
Value step = vector::StepOp::create(rewriter, loc, stepType);
Value limitVec =
vector::BroadcastOp::create(rewriter, loc, stepType, limit);
Value dimMask = arith::CmpIOp::create(
rewriter, loc, arith::CmpIPredicate::slt, step, limitVec);
if (vectorShape.size() > 1) {
SmallVector<int64_t> expandedShape(vectorShape.size(), 1);
expandedShape[v] = vectorShape[v];
dimMask = vector::ShapeCastOp::create(
rewriter, loc, VectorType::get(expandedShape, rewriter.getI1Type()),
dimMask);
dimMask = vector::BroadcastOp::create(rewriter, loc, maskType, dimMask);
}
mask = mask
? arith::AndIOp::create(rewriter, loc, mask, dimMask).getResult()
: dimMask;
}
if (mask)
return mask;
return vector::ConstantMaskOp::create(rewriter, loc, maskType, vectorShape);
}
// Builds that mask as a scalar `i1`, for a transfer of a single element. The
// element sits at `indices`, so `i < dim(d) - indices[d]` at `i == 0` reduces
// to `indices[d] < dim(d)`, which needs neither the limit subtraction nor the
// step vector. The dimensions checked are the same.
//
// Example, for a `vector<1xf32>` read of a `memref<?xf32>` at `%off`:
// %dim = memref.dim %src, %c0
// %mask = arith.cmpi slt, %off, %dim : index
static Value computeUnitInBoundsMask(VectorTransferOpInterface xferOp,
PatternRewriter &rewriter) {
Location loc = xferOp.getLoc();
AffineMap map = xferOp.getPermutationMap();
OperandRange indices = xferOp.getIndices();
Value mask;
for (unsigned v = 0, e = xferOp.getVectorType().getRank(); v < e; ++v) {
if (xferOp.isDimInBounds(v))
continue;
unsigned d = cast<AffineDimExpr>(map.getResult(v)).getPosition();
Value bound = getMemrefDimSize(xferOp, d, rewriter);
Value dimMask = arith::CmpIOp::create(
rewriter, loc, arith::CmpIPredicate::slt, indices[d], bound);
mask = mask
? arith::AndIOp::create(rewriter, loc, mask, dimMask).getResult()
: dimMask;
}
if (mask)
return mask;
return arith::ConstantOp::create(rewriter, loc, rewriter.getBoolAttr(true));
}
// Compute the element-wise offsets for vector.gather or vector.scatter ops.
//
// This function linearizes the base offsets of the gather/scatter operation
// and combines them with the per-element indices to produce a final vector of
// memory offsets.
template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
std::decay_t<OpType>, vector::GatherOp, vector::ScatterOp>::value>>
static Value computeOffsets(PatternRewriter &rewriter, OpType gatScatOp,
ArrayRef<Value> strides, Value baseOffset) {
Location loc = gatScatOp.getLoc();
SmallVector<Value> offsets = gatScatOp.getOffsets();
for (size_t i = 0; i < offsets.size(); ++i) {
Value offsetContrib =
arith::MulIOp::create(rewriter, loc, offsets[i], strides[i]);
baseOffset =
arith::AddIOp::create(rewriter, loc, baseOffset, offsetContrib);
}
Value indices = gatScatOp.getIndices();
VectorType vecType = cast<VectorType>(indices.getType());
Value strideVector =
vector::BroadcastOp::create(rewriter, loc, vecType, strides.back())
.getResult();
Value stridedIndices =
arith::MulIOp::create(rewriter, loc, strideVector, indices).getResult();
Value baseVector =
vector::BroadcastOp::create(
rewriter, loc,
VectorType::get(vecType.getShape(), rewriter.getIndexType()),
baseOffset)
.getResult();
return arith::AddIOp::create(rewriter, loc, baseVector, stridedIndices)
.getResult();
}
// Collapses shapes of a nD memref to the target rank while applying offsets for
// the collapsed dimensions. Returns the new memref value and the remaining
// offsets for the last targetRank dimensions. For example:
// input: %memref = memref<2x4x8x32xf32>, offsets=[%i0, %i1, %i2, %i3],
// output: %memref[%i0, %i1, 0, 0] -> memref<8x32xf32>, offsets: [%i2, %i3]
static std::pair<Value, SmallVector<OpFoldResult>>
convertMemrefAndOffsetsToTargetRank(PatternRewriter &rewriter, Location loc,
Value memref,
SmallVector<OpFoldResult> offsets,
int64_t targetRank) {
auto memrefType = cast<MemRefType>(memref.getType());
unsigned rank = memrefType.getRank();
if (rank <= targetRank)
return {memref, offsets};
int64_t numCombinedDims = rank - targetRank;
SmallVector<OpFoldResult> subviewOffsets;
SmallVector<OpFoldResult> subviewSizes;
SmallVector<OpFoldResult> subviewStrides;
// For the combined dimensions: use the provided offsets, size=1, stride=1
for (unsigned i = 0; i < numCombinedDims; ++i) {
subviewOffsets.push_back(offsets[i]);
subviewSizes.push_back(rewriter.getI64IntegerAttr(1));
subviewStrides.push_back(rewriter.getI64IntegerAttr(1));
}
// For the last targetRank dimensions: offset=0, use full size, stride=1
SmallVector<int64_t> resultShape;
auto originalShape = memrefType.getShape();
auto meta = memref::ExtractStridedMetadataOp::create(rewriter, loc, memref);
for (unsigned i = numCombinedDims; i < rank; ++i) {
subviewOffsets.push_back(rewriter.getI64IntegerAttr(0));
if (ShapedType::isDynamic(originalShape[i])) {
subviewSizes.push_back(meta.getSizes()[i]);
resultShape.push_back(ShapedType::kDynamic);
} else {
subviewSizes.push_back(rewriter.getI64IntegerAttr(originalShape[i]));
resultShape.push_back(originalShape[i]);
}
subviewStrides.push_back(rewriter.getI64IntegerAttr(1));
}
auto resultType = memref::SubViewOp::inferRankReducedResultType(
resultShape, memrefType, subviewOffsets, subviewSizes, subviewStrides);
auto subviewOp =
memref::SubViewOp::create(rewriter, loc, resultType, memref,
subviewOffsets, subviewSizes, subviewStrides);
// Return the remaining offsets for the last targetRank dimensions
SmallVector<OpFoldResult> newOffsets(offsets.begin() + numCombinedDims,
offsets.end());
return {subviewOp.getResult(), newOffsets};
}
template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
std::decay_t<OpType>, vector::TransferReadOp, vector::TransferWriteOp,
vector::GatherOp, vector::ScatterOp>::value>>
// Convert memref to i64 base pointer
static Value memrefToIndexPtr(OpType xferOp, PatternRewriter &rewriter) {
Location loc = xferOp.getLoc();
auto indexPtr = memref::ExtractAlignedPointerAsIndexOp::create(
rewriter, loc, xferOp.getBase())
.getResult();
return arith::IndexCastOp::create(rewriter, loc, rewriter.getI64Type(),
indexPtr)
.getResult();
}
// Returns true if every use of `vec` extracts a scalar element from it. An
// unused value does not qualify.
static bool isUsedAsScalar(Value vec) {
if (vec.use_empty())
return false;
return llvm::all_of(vec.getUsers(), [](Operation *user) {
auto extractOp = dyn_cast<vector::ExtractOp>(user);
return extractOp && !isa<VectorType>(extractOp.getResult().getType());
});
}
// Lowers a transfer of a single element to a scalar `xegpu.load`.
//
// The transfer touches one location, which its own indices name, so the load
// takes a scalar offset and a scalar mask - no descriptor, no lane vectors -
// and the result is broadcast back to the transfer's unit-size vector type.
//
// %off = arith.addi %base, %contrib : index
// %inb = arith.cmpi slt, %idx, %dim : index
// %val = xegpu.load %src[%off], %inb : i64, index, i1 -> f32
// %vec = vector.broadcast %val : f32 to vector<1xf32>
static LogicalResult lowerToScalarLoadOp(vector::TransferReadOp readOp,
PatternRewriter &rewriter) {
Location loc = readOp.getLoc();
VectorType vectorType = readOp.getVectorType();
if (!isa<MemRefType>(readOp.getShapedType()))
return rewriter.notifyMatchFailure(readOp, "Expected memref source");
auto meta = computeMemrefMeta(readOp, rewriter);
if (meta.first.empty())
return rewriter.notifyMatchFailure(readOp, "Failed to compute strides");
Value offset = computeBaseOffset(readOp, rewriter, meta.first, meta.second);
Value flatMemref = memrefToIndexPtr(readOp, rewriter);
Value mask = computeUnitInBoundsMask(readOp, rewriter);
auto loadOp = xegpu::LoadGatherOp::create(
rewriter, loc, vectorType.getElementType(), flatMemref, offset, mask,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
/*layout=*/nullptr, /*contiguity=*/nullptr);
// A masked-off xegpu.load is unspecified, so the padding has to be applied
// explicitly, as in lowerToScatteredLoadOp. A poison padding is "don't care".
Value scalar = loadOp.getResult();
if (readOp.hasOutOfBoundsDim() &&
!readOp.getPadding().getDefiningOp<ub::PoisonOp>())
scalar = arith::SelectOp::create(rewriter, loc, mask, scalar,
readOp.getPadding());
rewriter.replaceOpWithNewOp<vector::BroadcastOp>(readOp, vectorType, scalar);
return success();
}
static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
PatternRewriter &rewriter) {
Location loc = readOp.getLoc();
VectorType vectorType = readOp.getVectorType();
auto memrefType = dyn_cast<MemRefType>(readOp.getShapedType());
if (!memrefType)
return rewriter.notifyMatchFailure(readOp, "Expected memref source");
auto meta = computeMemrefMeta(readOp, rewriter);
if (meta.first.empty())
return rewriter.notifyMatchFailure(readOp, "Failed to compute strides");
Value localOffsets =
computeOffsets(readOp, rewriter, meta.first, meta.second);
Value flatMemref = memrefToIndexPtr(readOp, rewriter);
Value mask = computeInBoundsMask(readOp, rewriter);
auto gatherOp = xegpu::LoadGatherOp::create(
rewriter, loc, vectorType, flatMemref, localOffsets, mask,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
/*layout=*/nullptr, /*contiguity=*/nullptr);
// The masked-off lanes of an xegpu.load are unspecified, so the transfer's
// padding has to be applied explicitly. A poison padding leaves them "don't
// care", so the select can be skipped.
Value result = gatherOp.getResult();
if (readOp.hasOutOfBoundsDim() &&
!readOp.getPadding().getDefiningOp<ub::PoisonOp>()) {
Value padding = vector::BroadcastOp::create(rewriter, loc, vectorType,
readOp.getPadding());
result = arith::SelectOp::create(rewriter, loc, mask, result, padding);
}
rewriter.replaceOp(readOp, result);
return success();
}
static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
PatternRewriter &rewriter) {
Location loc = writeOp.getLoc();
auto memrefType = dyn_cast<MemRefType>(writeOp.getShapedType());
if (!memrefType)
return rewriter.notifyMatchFailure(writeOp, "Expected memref source");
auto meta = computeMemrefMeta(writeOp, rewriter);
if (meta.first.empty())
return rewriter.notifyMatchFailure(writeOp, "Failed to compute strides");
Value localOffsets =
computeOffsets(writeOp, rewriter, meta.first, meta.second);
Value flatMemref = memrefToIndexPtr(writeOp, rewriter);
// Out-of-bounds elements are simply not stored, so the mask is all this
// needs - no counterpart to the read's padding.
Value mask = computeInBoundsMask(writeOp, rewriter);
xegpu::StoreScatterOp::create(rewriter, loc, writeOp.getVector(), flatMemref,
localOffsets, mask,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
/*layout=*/nullptr, /*contiguity=*/nullptr);
rewriter.eraseOp(writeOp);
return success();
}
struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
using Base::Base;
LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
PatternRewriter &rewriter) const override {
Location loc = readOp.getLoc();
if (failed(transferPreconditions(rewriter, readOp)))
return failure();
auto readMemTy = cast<MemRefType>(readOp.getShapedType());
VectorType loadedVecTy = readOp.getVectorType();
bool isOutOfBounds = readOp.hasOutOfBoundsDim();
// Check if the memref has address space 3 (shared local memory)
bool isSharedMemory = xegpu::XeGPUDialect::isSharedMemory(readMemTy);
// Handle the SLM case.
if (isSharedMemory) {
// load_matrix supports 1D and 2D loads from SLM.
if (loadedVecTy.getRank() != 1 && loadedVecTy.getRank() != 2)
return rewriter.notifyMatchFailure(
readOp, "Only 1D and 2D vector loads are supported for SLM");
AffineMap readMap = readOp.getPermutationMap();
if (!readMap.isMinorIdentity())
return rewriter.notifyMatchFailure(
readOp,
"Non identity transposition is not supported for SLM loads.");
// Out of bounds case is not supported for SLM loads.
if (isOutOfBounds)
return rewriter.notifyMatchFailure(
readOp, "Out-of-bounds access is not supported for SLM loads");
// Create mem_desc for SLM
auto memDescType =
xegpu::MemDescType::get(rewriter.getContext(), readMemTy.getShape(),
readMemTy.getElementType(),
/*mem_layout=*/nullptr);
auto createMemDescOp = xegpu::CreateMemDescOp::create(
rewriter, loc, memDescType, readOp.getBase());
// Convert indices to OpFoldResult for LoadMatrixOp
SmallVector<OpFoldResult> indices =
getAsOpFoldResult(readOp.getIndices());
auto loadMatrixOp = xegpu::LoadMatrixOp::create(
rewriter, loc, loadedVecTy, createMemDescOp.getResult(), indices,
/*layout=*/nullptr);
rewriter.replaceOp(readOp, loadMatrixOp.getResult());
return success();
}
// A transfer of a single element that is only ever extracted to a scalar
// becomes a scalar load, at any rank: the vector is a wrapper its consumers
// undo, so neither an nd descriptor nor a lane-vector gather buys anything.
// A unit-size vector genuinely used as a vector keeps the paths below.
if (loadedVecTy.getNumElements() == 1 && isUsedAsScalar(readOp.getResult()))
return lowerToScalarLoadOp(readOp, rewriter);
const xegpu::uArch::uArch *uArch =
xegpu::uArch::getUArch(xegpu::getChipStr(readOp).value_or(""));
// An nd block load can realize a minor-identity map directly, or an
// innermost-two-dims transpose via a trailing vector.transpose. Any other
// permutation (e.g. a mid-vector transpose of a high-dim load) is left to
// the scattered path, which permutes strides explicitly.
AffineMap readMap = readOp.getPermutationMap();
bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
// A block load transfers the tile as it is laid out in memory, so a
// transposing read describes a tile whose innermost two dims are swapped.
Type elementType = loadedVecTy.getElementType();
SmallVector<int64_t> descShape(loadedVecTy.getShape());
if (isTransposeLoad) {
size_t rank = descShape.size();
assert(rank >= 2 && "Transpose requires at least 2 dimensions");
std::swap(descShape[rank - 1], descShape[rank - 2]);
}
// Prefer an nd block load. It requires a vector of rank >= 2 backed by a
// scalar-element memref, a map the block load can realize, and a tile shape
// the target's 2D block load can access. 1D vectors use the scattered
// xegpu.load path instead, which has a richer interface (e.g. layout
// capabilities). Out-of-bounds reads are allowed as long as the padding
// matches load_nd's implicit zero padding.
bool canLowerToLoadNd =
loadedVecTy.getRank() > 1 &&
(readMap.isMinorIdentity() || isTransposeLoad) &&
readMemTy.getElementType().isIntOrFloat() &&
(!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding())) &&
isSupportedBlockShape(
uArch, xegpu::uArch::InstructionKind::Subgroup2DBlockLoad,
descShape, elementType, /*hasTranspose=*/isTransposeLoad);
if (canLowerToLoadNd) {
// The load produces the memory-ordered tile; the transpose below restores
// the shape the transfer_read asked for.
if (isTransposeLoad)
loadedVecTy = VectorType::get(descShape, elementType);
auto descType = xegpu::TensorDescType::get(
descShape, elementType, /*array_length=*/1,
/*boundary_check=*/isOutOfBounds, xegpu::MemorySpace::Global);
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
rewriter, loc, readOp.getBase(),
getAsOpFoldResult(readOp.getIndices()), loadedVecTy.getRank());
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
Operation *loadedOp =
xegpu::LoadNdOp::create(rewriter, loc, loadedVecTy, ndDesc, indices,
/*packed=*/nullptr, /*transpose=*/nullptr,
/*l1_hint=*/hint,
/*l2_hint=*/hint, /*l3_hint=*/hint,
/*layout=*/nullptr);
if (isTransposeLoad) {
// Undo the innermost-two-dims swap with a trailing vector.transpose:
// keep the leading dimensions in place and interchange only the last
// two.
int64_t rank = loadedVecTy.getRank();
SmallVector<int64_t> perm(llvm::to_vector(llvm::seq<int64_t>(0, rank)));
std::swap(perm[rank - 1], perm[rank - 2]);
loadedOp = vector::TransposeOp::create(rewriter, loc,
loadedOp->getResult(0), perm);
}
rewriter.replaceOp(readOp, loadedOp);
return success();
}
// Fall back to a scattered load. It supports arbitrary permutations and any
// rank, and masks off the out-of-bounds elements.
return lowerToScatteredLoadOp(readOp, rewriter);
}
};
struct TransferWriteLowering
: public OpRewritePattern<vector::TransferWriteOp> {
using Base::Base;
LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
PatternRewriter &rewriter) const override {
Location loc = writeOp.getLoc();
if (failed(transferPreconditions(rewriter, writeOp)))
return failure();
// Perform common data transfer checks.
VectorType vecTy = writeOp.getVectorType();
auto writeMemTy = cast<MemRefType>(writeOp.getShapedType());
// Check if the memref has address space 3 (shared local memory)
bool isSharedMemory = xegpu::XeGPUDialect::isSharedMemory(writeMemTy);
// For shared local memory (address space 3), use create_mem_desc +
// store_matrix
if (isSharedMemory) {
// store_matrix supports 1D and 2D stores to SLM.
if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
return rewriter.notifyMatchFailure(
writeOp, "Only 1D and 2D vector stores are supported for SLM");
// Out of bounds case is not supported for SLM stores.
if (writeOp.hasOutOfBoundsDim())
return rewriter.notifyMatchFailure(
writeOp, "Out-of-bounds access is not supported for SLM stores");
// Create mem_desc for SLM
auto memDescType =
xegpu::MemDescType::get(rewriter.getContext(), writeMemTy.getShape(),
writeMemTy.getElementType(),
/*mem_layout=*/nullptr);
auto createMemDescOp = xegpu::CreateMemDescOp::create(
rewriter, loc, memDescType, writeOp.getBase());
// Convert indices to OpFoldResult for StoreMatrixOp
SmallVector<OpFoldResult> indices =
getAsOpFoldResult(writeOp.getIndices());
xegpu::StoreMatrixOp::create(rewriter, loc, writeOp.getVector(),
createMemDescOp.getResult(), indices,
/*layout=*/nullptr);
rewriter.eraseOp(writeOp);
return success();
}
const xegpu::uArch::uArch *uArch =
xegpu::uArch::getUArch(xegpu::getChipStr(writeOp).value_or(""));
// Prefer an nd block store. It requires a vector of rank >= 2 backed by a
// scalar-element memref, a minor-identity map (block stores have no
// transpose support), and a tile shape the target's 2D block store can
// access. 1D vectors use the scattered xegpu.store path instead, which has
// a richer interface. Out-of-bounds writes are handled by the descriptor's
// boundary check.
AffineMap map = writeOp.getPermutationMap();
bool canLowerToStoreNd =
vecTy.getRank() > 1 && map.isMinorIdentity() &&
writeMemTy.getElementType().isIntOrFloat() &&
isSupportedBlockShape(
uArch, xegpu::uArch::InstructionKind::Subgroup2DBlockStore,
vecTy.getShape(), vecTy.getElementType());
if (canLowerToStoreNd) {
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
rewriter, loc, writeOp.getBase(),
getAsOpFoldResult(writeOp.getIndices()), vecTy.getRank());
auto descType = xegpu::TensorDescType::get(
vecTy.getShape(), vecTy.getElementType(),
/*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
xegpu::MemorySpace::Global);
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
auto storeOp = xegpu::StoreNdOp::create(
rewriter, loc, writeOp.getVector(), ndDesc, indices,
/*l1_hint=*/hint,
/*l2_hint=*/hint, /*l3_hint=*/hint,
/*layout=*/nullptr);
rewriter.replaceOp(writeOp, storeOp);
return success();
}
// Fall back to a scattered store. It supports arbitrary permutations and
// any rank, and masks off the out-of-bounds elements.
return lowerToScatteredStoreOp(writeOp, rewriter);
}
};
struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
using Base::Base;
LogicalResult matchAndRewrite(vector::GatherOp gatherOp,
PatternRewriter &rewriter) const override {
auto srcTy = dyn_cast<MemRefType>(gatherOp.getBase().getType());
if (!srcTy)
return rewriter.notifyMatchFailure(gatherOp, "Expects memref source");
Location loc = gatherOp.getLoc();
VectorType vectorType = gatherOp.getVectorType();
auto meta = computeMemrefMeta(gatherOp, rewriter);
if (meta.first.empty())
return rewriter.notifyMatchFailure(gatherOp, "Failed to compute strides");
Value localOffsets =
computeOffsets(rewriter, gatherOp, meta.first, meta.second);
Value flatMemref = memrefToIndexPtr(gatherOp, rewriter);
auto xeGatherOp = xegpu::LoadGatherOp::create(
rewriter, loc, vectorType, flatMemref, localOffsets, gatherOp.getMask(),
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
/*layout=*/nullptr, /*contiguity=*/nullptr);
auto selectOp =
arith::SelectOp::create(rewriter, loc, gatherOp.getMask(),
xeGatherOp.getResult(), gatherOp.getPassThru());
rewriter.replaceOp(gatherOp, selectOp.getResult());
return success();
}
};
struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
using Base::Base;
LogicalResult matchAndRewrite(vector::ScatterOp scatterOp,
PatternRewriter &rewriter) const override {
auto srcTy = dyn_cast<MemRefType>(scatterOp.getBase().getType());
if (!srcTy)
return rewriter.notifyMatchFailure(scatterOp, "Expects memref source");
Location loc = scatterOp.getLoc();
auto meta = computeMemrefMeta(scatterOp, rewriter);
if (meta.first.empty())
return rewriter.notifyMatchFailure(scatterOp,
"Failed to compute strides");
Value localOffsets =
computeOffsets(rewriter, scatterOp, meta.first, meta.second);
Value flatMemref = memrefToIndexPtr(scatterOp, rewriter);
xegpu::StoreScatterOp::create(rewriter, loc, scatterOp.getValueToStore(),
flatMemref, localOffsets, scatterOp.getMask(),
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
/*layout=*/nullptr,
/*contiguity=*/nullptr);
rewriter.eraseOp(scatterOp);
return success();
}
};
struct LoadLowering : public OpRewritePattern<vector::LoadOp> {
using Base::Base;
LogicalResult matchAndRewrite(vector::LoadOp loadOp,
PatternRewriter &rewriter) const override {
Location loc = loadOp.getLoc();
VectorType vecTy = loadOp.getResult().getType();
MemRefType memTy = loadOp.getBase().getType();
// The plain vector.load lowering only supports 1D/2D block loads.
if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
return rewriter.notifyMatchFailure(loadOp, "Expects 1D or 2D vector");
if (!memTy.getElementType().isIntOrFloat())
return rewriter.notifyMatchFailure(
loadOp, "Unsupported memref element type: expected integer or float");
// Boundary check is available only for block instructions.
bool boundaryCheck = vecTy.getRank() > 1;
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
rewriter, loc, loadOp.getBase(), getAsOpFoldResult(loadOp.getIndices()),
vecTy.getRank());
auto descType = xegpu::TensorDescType::get(
vecTy.getShape(), vecTy.getElementType(), /*array_length=*/1,
boundaryCheck, xegpu::MemorySpace::Global);
xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
auto loadNdOp =
xegpu::LoadNdOp::create(rewriter, loc, vecTy, ndDesc, indices,
/*packed=*/nullptr, /*transpose=*/nullptr,
/*l1_hint=*/hint,
/*l2_hint=*/hint, /*l3_hint=*/hint,
/*layout=*/nullptr);
rewriter.replaceOp(loadOp, loadNdOp);
return success();
}
};
struct StoreLowering : public OpRewritePattern<vector::StoreOp> {
using Base::Base;
LogicalResult matchAndRewrite(vector::StoreOp storeOp,
PatternRewriter &rewriter) const override {
Location loc = storeOp.getLoc();
TypedValue<VectorType> vector = storeOp.getValueToStore();
VectorType vecTy = vector.getType();
MemRefType memTy = storeOp.getBase().getType();
// The plain vector.store lowering only supports 1D/2D block stores.
if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
return rewriter.notifyMatchFailure(storeOp, "Expects 1D or 2D vector");
if (!memTy.getElementType().isIntOrFloat())
return rewriter.notifyMatchFailure(
storeOp,
"Unsupported memref element type: expected integer or float");
// Boundary check is available only for block instructions.
bool boundaryCheck = vecTy.getRank() > 1;
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
rewriter, loc, storeOp.getBase(),
getAsOpFoldResult(storeOp.getIndices()), vecTy.getRank());
auto descType = xegpu::TensorDescType::get(
vecTy.getShape(), vecTy.getElementType(),
/*array_length=*/1, boundaryCheck, xegpu::MemorySpace::Global);
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
auto storeNdOp =
xegpu::StoreNdOp::create(rewriter, loc, vector, ndDesc, indices,
/*l1_hint=*/hint,
/*l2_hint=*/hint, /*l3_hint=*/hint,
/*layout=*/nullptr);
rewriter.replaceOp(storeOp, storeNdOp);
return success();
}
};
// If `indexingMaps` describe a (batched) row-major matmul
// lhs[b..., m, k], rhs[b..., k, n], acc[b..., m, n]
// return the number of leading batch dims (0 for a plain 2D matmul);
// otherwise return std::nullopt.
static std::optional<int64_t>
getRowMajorMatmulBatchRank(ArrayAttr indexingMaps) {
if (indexingMaps.size() != 3)
return std::nullopt;
AffineMap mapA = cast<AffineMapAttr>(indexingMaps[0]).getValue();
AffineMap mapB = cast<AffineMapAttr>(indexingMaps[1]).getValue();
AffineMap mapC = cast<AffineMapAttr>(indexingMaps[2]).getValue();
// The result map exposes the batch dims followed by the core (m, n) dims.
if (mapC.getNumResults() < 2)
return std::nullopt;
int64_t batchRank = mapC.getNumResults() - 2;
// A single `k` reduction gives batchRank + 3 iteration dims; each operand
// map exposes batchRank + 2 dims (batch dims + 2 core dims).
unsigned numDims = static_cast<unsigned>(batchRank) + 3;
unsigned numOperandResults = static_cast<unsigned>(batchRank) + 2;
if (mapA.getNumInputs() != numDims || mapB.getNumInputs() != numDims ||
mapC.getNumInputs() != numDims)
return std::nullopt;
if (mapA.getNumResults() != numOperandResults ||
mapB.getNumResults() != numOperandResults)
return std::nullopt;
// Reconstruct the canonical maps from the batch/m/n dims of the result and
// the k dim of lhs, then compare against the actual maps.
MLIRContext *context = indexingMaps.getContext();
ArrayRef<AffineExpr> batchDims = mapC.getResults().take_front(batchRank);
AffineExpr m = mapC.getResult(batchRank);
AffineExpr n = mapC.getResult(batchRank + 1);
AffineExpr k = mapA.getResult(batchRank + 1);
SmallVector<AffineExpr> aDims = llvm::to_vector(batchDims);
aDims.push_back(m);
aDims.push_back(k);
SmallVector<AffineExpr> bDims = llvm::to_vector(batchDims);
bDims.push_back(k);
bDims.push_back(n);
SmallVector<AffineExpr> cDims = llvm::to_vector(batchDims);
cDims.push_back(m);
cDims.push_back(n);
auto expected = ArrayAttr::get(
context,
{AffineMapAttr::get(AffineMap::get(numDims, 0, aDims, context)),
AffineMapAttr::get(AffineMap::get(numDims, 0, bDims, context)),
AffineMapAttr::get(AffineMap::get(numDims, 0, cDims, context))});
if (indexingMaps != expected)
return std::nullopt;
return batchRank;
}
struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
using Base::Base;
LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
PatternRewriter &rewriter) const override {
Location loc = contractOp.getLoc();
if (contractOp.getKind() != vector::CombiningKind::ADD)
return rewriter.notifyMatchFailure(contractOp,
"Expects add combining kind");
TypedValue<VectorType> lhs = contractOp.getLhs();
TypedValue<VectorType> rhs = contractOp.getRhs();
TypedValue<Type> acc = contractOp.getAcc();
VectorType accType = dyn_cast<VectorType>(acc.getType());
if (!accType)
return rewriter.notifyMatchFailure(contractOp, "Expects vector acc");
std::optional<int64_t> batchRank =
getRowMajorMatmulBatchRank(contractOp.getIndexingMapsAttr());
if (!batchRank)
return rewriter.notifyMatchFailure(
contractOp,
"Expects a (batched) row-major matmul: leading dims must "
"be batch dims shared by lhs, rhs, and acc; innermost two "
"dims must be (M, K), (K, N), and (M, N)");
// xegpu.dpas operands are limited to 2 batch + 2 core dims.
if (*batchRank > 2)
return rewriter.notifyMatchFailure(contractOp,
"Expects operands of rank 4 or less");
auto dpasOp = xegpu::DpasOp::create(
rewriter, loc, contractOp.getResultType(), lhs, rhs, acc,
/*layout_a=*/nullptr, /*layout_b=*/nullptr, /*layout_cd=*/nullptr);
rewriter.replaceOp(contractOp, dpasOp);
return success();
}
};
// Returns the `vector.shape_cast` that flattened a value of type `ndType` into
// `flat`, if that is how `flat` was produced.
static vector::ShapeCastOp getFlattenCast(Value flat, VectorType ndType) {
auto shapeCast = flat.getDefiningOp<vector::ShapeCastOp>();
if (shapeCast && shapeCast.getSourceVectorType() == ndType)
return shapeCast;
return nullptr;
}
static DenseElementsAttr getDenseConstant(Value flat) {
DenseElementsAttr elements;
if (matchPattern(flat, m_Constant(&elements)))
return elements;
return nullptr;
}
static vector::BroadcastOp getSplatBroadcast(Value flat) {
auto broadcast = flat.getDefiningOp<vector::BroadcastOp>();
if (broadcast && !isa<VectorType>(broadcast.getSourceType()))
return broadcast;
return nullptr;
}
// Returns true if the cast `unflatten` creates will fold away, i.e. if `flat`
// is a flattening cast, a constant or a splat.
static bool canUnflatten(Value flat, VectorType ndType) {
return getFlattenCast(flat, ndType) || getDenseConstant(flat) ||
getSplatBroadcast(flat);
}
// Reshape `flat` to `ndType`; `canUnflatten` must hold so the cast folds away.
static Value unflatten(PatternRewriter &rewriter, Value flat,
VectorType ndType) {
assert(canUnflatten(flat, ndType) && "expected the cast to fold away");
return vector::ShapeCastOp::create(rewriter, flat.getLoc(), ndType, flat);
}
// Restore the N-D form of a flattened `vector.gather` / `vector.scatter`.
//
// XeGPU layouts are expressed in terms of the N-D shape of the accessed data,
// so a flattened gather/scatter forces layout propagation to reason through the
// surrounding `vector.shape_cast` ops. That adds complexity and tends to yield
// layouts that lower to unoptimized code.
//
// Before:
// %cst = arith.constant dense<0.0> : vector<8192xbf16>
// %fi = vector.shape_cast %idx : vector<128x64xindex> to vector<8192xindex>
// %fm = vector.shape_cast %mask : vector<128x64xi1> to vector<8192xi1>
// %fr = vector.gather %src[%c0] [%fi], %fm, %cst : memref<?xbf16>,
// vector<8192xindex>, vector<8192xi1>, vector<8192xbf16>
// into vector<8192xbf16>
// %res = vector.shape_cast %fr : vector<8192xbf16> to vector<128x64xbf16>
//
// After:
// %cst = arith.constant dense<0.0> : vector<128x64xbf16>
// %res = vector.gather %src[%c0] [%idx], %mask, %cst : memref<?xbf16>,
// vector<128x64xindex>, vector<128x64xi1>, vector<128x64xbf16>
// into vector<128x64xbf16>
template <typename OpTy>
struct UnflattenGatherScatter : public OpRewritePattern<OpTy> {
using OpRewritePattern<OpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(OpTy op,
PatternRewriter &rewriter) const override {
constexpr bool isGather = std::is_same_v<OpTy, vector::GatherOp>;
if (!isa<MemRefType>(op.getBase().getType()))
return rewriter.notifyMatchFailure(op, "expects a memref source");
if (op.getIndexVectorType().getRank() != 1)
return rewriter.notifyMatchFailure(op, "index vector is not 1-D");
// The N-D shape comes from the index operand's producer: this only undoes
// a flattening that already happened, it never invents a shape.
auto indexCast =
op.getIndices().template getDefiningOp<vector::ShapeCastOp>();
if (!indexCast || indexCast.getSourceVectorType().getRank() < 2)
return rewriter.notifyMatchFailure(
op, "index vector is not a shape_cast of an N-D vector");
VectorType ndIndexType = indexCast.getSourceVectorType();
VectorType ndMaskType =
ndIndexType.cloneWith(std::nullopt, rewriter.getI1Type());
VectorType ndType = ndIndexType.cloneWith(
std::nullopt, op.getVectorType().getElementType());
// Check everything before creating any IR: a partially applied rewrite
// would leave dead ops behind.
if (!canUnflatten(op.getMask(), ndMaskType))
return rewriter.notifyMatchFailure(op, "cannot un-flatten the mask");
if constexpr (isGather) {
if (!canUnflatten(op.getPassThru(), ndType))
return rewriter.notifyMatchFailure(op,
"cannot un-flatten the pass-thru");
Value mask = unflatten(rewriter, op.getMask(), ndMaskType);
Value passThru = unflatten(rewriter, op.getPassThru(), ndType);
auto ndGather = vector::GatherOp::create(
rewriter, op.getLoc(), ndType, op.getBase(), op.getOffsets(),
indexCast.getSource(), mask, passThru, op.getAlignmentAttr());
ndGather->setDiscardableAttrs(op->getDiscardableAttrDictionary());
rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, op.getVectorType(),
ndGather);
} else {
if (!canUnflatten(op.getValueToStore(), ndType))
return rewriter.notifyMatchFailure(
op, "cannot un-flatten the stored value");
Value mask = unflatten(rewriter, op.getMask(), ndMaskType);
Value valueToStore = unflatten(rewriter, op.getValueToStore(), ndType);
// Only operand types change, and this keeps the optional tensor result
// untouched.
rewriter.modifyOpInPlace(op, [&] {
op.getIndicesMutable().assign(indexCast.getSource());
op.getMaskMutable().assign(mask);
op.getValueToStoreMutable().assign(valueToStore);
});
}
return success();
}
};
// Un-flatten every gather/scatter that was flattened to 1-D operands, so that
// the conversion patterns and the XeGPU layouts downstream of them see the N-D
// shape of the accessed data.
static LogicalResult unflattenGatherScatter(Operation *root) {
MLIRContext *ctx = root->getContext();
RewritePatternSet patterns(ctx);
patterns.add<UnflattenGatherScatter<vector::GatherOp>,
UnflattenGatherScatter<vector::ScatterOp>>(ctx);
vector::ShapeCastOp::getCanonicalizationPatterns(patterns, ctx);
return applyPatternsGreedily(root, std::move(patterns));
}
// Returns `memrefTy` with its memory space replaced by `newMemSpace`.
static MemRefType withMemorySpace(MemRefType memrefTy, Attribute newMemSpace) {
return MemRefType::get(memrefTy.getShape(), memrefTy.getElementType(),
memrefTy.getLayout(), newMemSpace);
}
// Rewrite every `memref.alloca` not already in shared local memory (SLM) to
// be in SLM (address space 3), and propagate the new memory space through
// memref-producing aliasing users (e.g. memref.cast, memref.subview,
// memref.expand_shape, ...). Consumers that take a memref operand but
// produce a non-memref result (e.g. vector.transfer_read, vector.load) are
// left untouched: their operand type simply reflects the new memory space.
//
// This makes `xegpu.load_matrix`/`xegpu.store_matrix` lowering work end-to-end
// for IR coming from bufferization, which by default assigns memory space 0/1
// to allocations.
static void promoteAllocasToSLM(Operation *root) {
MLIRContext *ctx = root->getContext();
Attribute slmAttr = IntegerAttr::get(IntegerType::get(ctx, 64), 3);
// A user is treated as a memref-producing alias (e.g. memref.cast,
// memref.subview, memref.expand_shape, ...) if it is side-effect free and
// produces at least one memref result. This excludes ops like memref.copy
// that have memory effects.
auto isMemrefResultOp = [](Operation *op) {
if (!isMemoryEffectFree(op))
return false;
return llvm::any_of(op->getResultTypes(),
[](Type t) { return isa<MemRefType>(t); });
};
// Update `v`'s type to have SLM memory space, then walk forward through
// memref-producing users and update their result types accordingly.
std::function<void(Value)> propagate = [&](Value v) {
auto memrefTy = dyn_cast<MemRefType>(v.getType());
if (!memrefTy || xegpu::XeGPUDialect::isSharedMemory(memrefTy))
return;
v.setType(withMemorySpace(memrefTy, slmAttr));
for (Operation *user : v.getUsers()) {
if (!isMemrefResultOp(user))
continue;
for (Value result : user->getResults())
propagate(result);
}
};
SmallVector<memref::AllocaOp> allocas;
root->walk([&](memref::AllocaOp op) {
auto memrefTy = dyn_cast<MemRefType>(op.getResult().getType());
if (!memrefTy || xegpu::XeGPUDialect::isSharedMemory(memrefTy))
return;
allocas.push_back(op);
});
for (memref::AllocaOp alloca : allocas) {
OpBuilder builder(alloca);
auto memrefTy = cast<MemRefType>(alloca.getResult().getType());
auto newTy = withMemorySpace(memrefTy, slmAttr);
auto newOp = memref::AllocaOp::create(
builder, alloca.getLoc(), newTy, alloca.getDynamicSizes(),
alloca.getSymbolOperands(), alloca.getAlignmentAttr());
alloca.getResult().replaceAllUsesWith(newOp.getResult());
alloca.erase();
// Propagate the new memory space through memref-producing consumers.
for (Operation *user : newOp.getResult().getUsers()) {
if (!isMemrefResultOp(user))
continue;
for (Value result : user->getResults())
propagate(result);
}
}
}
struct ConvertVectorToXeGPUPass
: public impl::ConvertVectorToXeGPUBase<ConvertVectorToXeGPUPass> {
void runOnOperation() override {
// Promote local allocations to SLM (address space 3) so that
// load_matrix/store_matrix lowerings have well-typed memref operands.
promoteAllocasToSLM(getOperation());
// Undo any flattening of gather/scatter operands, so that the conversion
// below sees the N-D shape the XeGPU layouts are expressed in.
if (failed(unflattenGatherScatter(getOperation())))
return signalPassFailure();
RewritePatternSet patterns(&getContext());
populateVectorToXeGPUConversionPatterns(patterns);
populatePrepareVectorToMMAPatterns(patterns);
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
return signalPassFailure();
}
};
} // namespace
void mlir::populateVectorToXeGPUConversionPatterns(
RewritePatternSet &patterns) {
patterns
.add<TransferReadLowering, TransferWriteLowering, LoadLowering,
ScatterLowering, GatherLowering, StoreLowering, ContractionLowering>(
patterns.getContext());
}