Google Git Google Git
Sign in
llvm/llvm-project/refs/heads/users/void/array-bounds-accessed/./mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
blob: c0b3adf8402e03b613121be29781c03a2fc9fa2d [file] [edit]
//===-- XeGPUToXeVM.cpp - XeGPU to XeVM dialect conversion ------*- C++ -*-===//
//
// This file is licensed 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/Conversion/XeGPUToXeVM/XeGPUToXeVM.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.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/uArchCommon.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Types.h"
#include "llvm/ADT/TypeSwitch.h"
#include <string>
namespace mlir {
#define GEN_PASS_DEF_CONVERTXEGPUTOXEVMPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
using namespace mlir;
namespace {
// TODO: Below are uArch dependent values, should move away from hardcoding
static constexpr int32_t systolicDepth{8};
static constexpr int32_t executionSize{16};
// Offsets to individual fields of the 8xi32 layout nd tensor descriptor.
enum class NdTdescOffset : uint32_t {
BasePtr = 0, // Base pointer (i64)
BaseShapeW = 2, // Base shape width (i32)
BaseShapeH = 3, // Base shape height (i32)
BasePitch = 4, // Base pitch/stride of dim rank-2 (i32)
LeadingStride0 = 5, // Row strides of the leading (batch) dims of a >2D
LeadingStride1 = 6, // descriptor (i32); added into offset_h by the load/store
LeadingStride2 = 7, // lowering. Left at 0 for 2D descriptors.
};
// Spare payload slots above, and the resulting max lowerable descriptor rank.
static constexpr int64_t maxNdTdescLeadingDims{3};
static constexpr int64_t maxNdTdescRank{2 + maxNdTdescLeadingDims};
static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
switch (xeGpuMemspace) {
case xegpu::MemorySpace::Global:
return static_cast<int>(xevm::AddrSpace::GLOBAL);
case xegpu::MemorySpace::SLM:
return static_cast<int>(xevm::AddrSpace::SHARED);
}
llvm_unreachable("Unknown XeGPU memory space");
}
/// Translates a memref memory space attribute into XeVM's numeric address
/// space, which follows the OpenCL/SPIR-V convention (0 = private, 1 =
/// global, 2 = constant, 3 = shared/local, 4 = generic). A null attribute,
/// meaning the memory space was left unspecified, maps to the default space
/// 0. Returns failure if `memSpace` is a representation this pass does not
/// know how to translate (e.g. a SPIR-V storage class or an arbitrary string
/// attribute), rather than assuming it is an `IntegerAttr` and asserting.
static FailureOr<unsigned> getNumericMemorySpace(Attribute memSpace) {
if (!memSpace)
return 0u;
if (auto intAttr = llvm::dyn_cast<IntegerAttr>(memSpace))
return static_cast<unsigned>(intAttr.getInt());
if (auto xevmSpace = llvm::dyn_cast<xevm::AddrSpaceAttr>(memSpace))
return static_cast<unsigned>(xevmSpace.getValue());
if (auto gpuSpace = llvm::dyn_cast<gpu::AddressSpaceAttr>(memSpace)) {
switch (gpuSpace.getValue()) {
case gpu::AddressSpace::Global:
return static_cast<unsigned>(xevm::AddrSpace::GLOBAL);
case gpu::AddressSpace::Workgroup:
return static_cast<unsigned>(xevm::AddrSpace::SHARED);
case gpu::AddressSpace::Private:
return static_cast<unsigned>(xevm::AddrSpace::PRIVATE);
case gpu::AddressSpace::Constant:
return static_cast<unsigned>(xevm::AddrSpace::CONSTANT);
}
llvm_unreachable("Unknown GPU address space");
}
return failure();
}
/// Checks if the given MemRefType refers to shared memory.
static bool isSharedMemRef(const MemRefType &memrefTy) {
FailureOr<unsigned> addrSpace =
getNumericMemorySpace(memrefTy.getMemorySpace());
return succeeded(addrSpace) &&
*addrSpace == static_cast<unsigned>(xevm::AddrSpace::SHARED);
}
// Get same bitwidth flat vector type of new element type.
static VectorType encodeVectorTypeTo(VectorType currentVecType,
Type toElemType) {
auto elemType = currentVecType.getElementType();
auto currentBitWidth = elemType.getIntOrFloatBitWidth();
auto newBitWidth = toElemType.getIntOrFloatBitWidth();
const int size =
currentVecType.getNumElements() * currentBitWidth / newBitWidth;
return VectorType::get(size, toElemType);
}
static xevm::LoadCacheControl
translateLoadXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
std::optional<xegpu::CachePolicy> L3hint) {
// If no hints are provided, use the default cache control.
if (!L1hint && !L3hint)
return xevm::LoadCacheControl::USE_DEFAULT;
// If only one of the hints is provided, use the default for the other level.
auto L1hintVal = L1hint.value_or(xegpu::CachePolicy::CACHED);
auto L3hintVal = L3hint.value_or(xegpu::CachePolicy::CACHED);
switch (L1hintVal) {
case xegpu::CachePolicy::CACHED:
if (L3hintVal == xegpu::CachePolicy::CACHED)
return xevm::LoadCacheControl::L1C_L2UC_L3C;
else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
return xevm::LoadCacheControl::L1C_L2UC_L3UC;
else
llvm_unreachable("Unsupported cache control.");
case xegpu::CachePolicy::UNCACHED:
if (L3hintVal == xegpu::CachePolicy::CACHED)
return xevm::LoadCacheControl::L1UC_L2UC_L3C;
else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
return xevm::LoadCacheControl::L1UC_L2UC_L3UC;
else
llvm_unreachable("Unsupported cache control.");
case xegpu::CachePolicy::STREAMING:
if (L3hintVal == xegpu::CachePolicy::CACHED)
return xevm::LoadCacheControl::L1S_L2UC_L3C;
else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
return xevm::LoadCacheControl::L1S_L2UC_L3UC;
else
llvm_unreachable("Unsupported cache control.");
case xegpu::CachePolicy::READ_INVALIDATE:
return xevm::LoadCacheControl::INVALIDATE_READ;
default:
llvm_unreachable("Unsupported cache control.");
}
}
static xevm::StoreCacheControl
translateStoreXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
std::optional<xegpu::CachePolicy> L3hint) {
// If no hints are provided, use the default cache control.
if (!L1hint && !L3hint)
return xevm::StoreCacheControl::USE_DEFAULT;
// If only one of the hints is provided, use the default for the other level.
auto L1hintVal = L1hint.value_or(xegpu::CachePolicy::UNCACHED);
auto L3hintVal = L3hint.value_or(xegpu::CachePolicy::WRITE_BACK);
switch (L1hintVal) {
case xegpu::CachePolicy::UNCACHED:
if (L3hintVal == xegpu::CachePolicy::UNCACHED)
return xevm::StoreCacheControl::L1UC_L2UC_L3UC;
else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
return xevm::StoreCacheControl::L1UC_L2UC_L3WB;
else
llvm_unreachable("Unsupported cache control.");
case xegpu::CachePolicy::STREAMING:
if (L3hintVal == xegpu::CachePolicy::UNCACHED)
return xevm::StoreCacheControl::L1S_L2UC_L3UC;
else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
return xevm::StoreCacheControl::L1S_L2UC_L3WB;
else
llvm_unreachable("Unsupported cache control.");
case xegpu::CachePolicy::WRITE_BACK:
if (L3hintVal == xegpu::CachePolicy::UNCACHED)
return xevm::StoreCacheControl::L1WB_L2UC_L3UC;
else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
return xevm::StoreCacheControl::L1WB_L2UC_L3WB;
else
llvm_unreachable("Unsupported cache control.");
case xegpu::CachePolicy::WRITE_THROUGH:
if (L3hintVal == xegpu::CachePolicy::UNCACHED)
return xevm::StoreCacheControl::L1WT_L2UC_L3UC;
else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
return xevm::StoreCacheControl::L1WT_L2UC_L3WB;
else
llvm_unreachable("Unsupported cache control.");
default:
llvm_unreachable("Unsupported cache control.");
}
}
//
// Note:
// Block operations for tile of sub byte element types are handled by
// emulating with larger element types.
// Tensor descriptor are keep intact and only ops consuming them are
// emulated
//
//
// High-D (>2D) nd descriptors are lowered by viewing the source as a single
// flattened 2D plane, so the 2D-block surface covers every leading (batch)
// plane at once and a batch position becomes a row offset into it. The surface
// height is the row extent the leading dims reach:
// base_height = size[R-2] + sum_d (size[d] - 1) * (stride[d] / stride[R-2])
//
// Leaving `base_ptr` at the true base means an out-of-range batch index lands
// past `base_height`, where the HW boundary check handles it, instead of aiming
// the surface at unmapped memory. Encoding the leading strides as row counts
// (`stride[d] / stride[R-2]`) also makes them dimensionless, so they survive
// element-type repacking (e.g. the f16 -> i32 transpose repack) without a unit
// conversion.
//
// Limitations of the flattened-plane view:
// 1. Each leading stride must be a whole number of rows. A source with gaps
// between planes (`stride[d] % stride[R-2] != 0`) is not lowered. This can
// only be checked when the strides are static; for dynamic strides the
// divisibility is assumed.
// 2. `base_height` grows with the leading extent, so a source with a large
// batch x head x sequence extent can exceed the HW 2D-block surface height.
// 3. Plane boundaries are invisible to the boundary check, which only knows
// the surface: a tile whose rows run past `size[R-2]` spills into the
// following rows of the surface -- the next plane, or the padding between
// planes -- rather than being clipped there as it would be on a per-plane
// surface. A load reads those rows instead of returning zeros, and a store
// overwrites them, so it can corrupt the next plane. This only matters
// when `size[R-2]` is not a multiple of the tile height.
//
class CreateNdDescToXeVMPattern
: public OpConversionPattern<xegpu::CreateNdDescOp> {
using OpConversionPattern::OpConversionPattern;
// Checks the surface width and pitch against the target's 2D block memory
// restrictions. Only these two are enforced: both are backed by the source's
// static shape and strides, whereas the base address alignment is a property
// of the allocation that is not visible here. The width and the pitch are
// checked independently, each only when it is static; a value known only at
// runtime is left alone, as guarding it inside the kernel would cost far more
// than it is worth.
//
// A violation is a dead end rather than a pattern that does not apply: the op
// has no other legalization path, so it is reported with an error that names
// the limit it misses.
LogicalResult checkMemoryRestrictions(xegpu::CreateNdDescOp op,
int64_t rank) const {
const auto *uArch =
xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
// Without a known target chip there is no uArch to check against, so the
// descriptor must be accepted as-is.
if (!uArch)
return success();
const xegpu::uArch::BlockIOMemoryRestrictions &limits =
uArch->getBlockIOMemoryRestrictions();
// Recover the width and the pitch as a number of elements. A memref source
// carries both in its type, and either one may be dynamic independently of
// the other, so they are queried separately instead of being dropped
// together when the source is not fully static. `getMixedSizes` /
// `getMixedStrides` are only consulted for an integer source, where the op
// itself holds the extents.
std::optional<int64_t> widthElems;
std::optional<int64_t> pitchElems;
if (auto memrefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
if (!ShapedType::isDynamic(memrefTy.getDimSize(rank - 1)))
widthElems = memrefTy.getDimSize(rank - 1);
// The source is known to be strided; the caller checked that already.
SmallVector<int64_t> staticStrides;
int64_t staticOffset;
if (succeeded(
memrefTy.getStridesAndOffset(staticStrides, staticOffset)) &&
!ShapedType::isDynamic(staticStrides[rank - 2]))
pitchElems = staticStrides[rank - 2];
} else {
widthElems = getConstantIntValue(op.getMixedSizes()[rank - 1]);
pitchElems = getConstantIntValue(op.getMixedStrides()[rank - 2]);
}
unsigned elemBitSize =
op.getType().getElementType().getIntOrFloatBitWidth();
// The restrictions are expressed in bytes, but a sub-byte element type can
// leave the width or the pitch on a fraction of a byte, so the comparisons
// are done in bits. Such a size meets no byte alignment at all and is
// reported as a violation rather than waved through.
auto describe = [](int64_t bits) -> std::string {
if (bits % 8 == 0)
return llvm::formatv("{0} bytes", bits / 8).str();
return llvm::formatv("{0} bits", bits).str();
};
auto check = [&](std::optional<int64_t> elems, int64_t minBytes,
int64_t alignBytes,
llvm::StringRef what) -> LogicalResult {
if (!elems)
return success();
int64_t bits = *elems * elemBitSize;
if (bits < minBytes * 8)
return op.emitError(llvm::formatv(
"2D block surface {0} must be at least {1} bytes on this "
"target, got {2}.",
what, minBytes, describe(bits)));
if (alignBytes > 0 && bits % (alignBytes * 8) != 0)
return op.emitError(llvm::formatv(
"2D block surface {0} must be a multiple of {1} bytes on this "
"target, got {2}.",
what, alignBytes, describe(bits)));
return success();
};
if (failed(check(widthElems, limits.minBaseWidthBytes,
limits.baseWidthAlignmentBytes, "width")))
return failure();
// The pitch is the row stride of the innermost 2D tile.
if (failed(check(pitchElems, limits.minBasePitchBytes,
limits.basePitchAlignmentBytes, "pitch")))
return failure();
// Consecutive rows may not overlap, so the pitch has to span the full
// width.
if (widthElems && pitchElems && *pitchElems < *widthElems)
return op.emitError(
llvm::formatv("2D block surface pitch must be at least the surface "
"width, got a pitch of {0} for a surface {1} wide.",
describe(*pitchElems * elemBitSize),
describe(*widthElems * elemBitSize)));
return success();
}
LogicalResult
matchAndRewrite(xegpu::CreateNdDescOp op,
xegpu::CreateNdDescOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto source = op.getSource();
// Check all failure conditions before generating any IR, so nothing has to
// be rolled back.
int64_t rank = op.getType().getRank();
int64_t sourceRank;
auto memrefTy = dyn_cast<MemRefType>(source.getType());
if (memrefTy) {
if (!memrefTy.isStrided())
return rewriter.notifyMatchFailure(op, "Expected strided Memref.");
sourceRank = memrefTy.getRank();
} else if (isa<IntegerType>(source.getType())) {
sourceRank = op.getMixedSizes().size();
} else {
return rewriter.notifyMatchFailure(op,
"Expected ranked Memref or integer.");
}
if (sourceRank != rank)
return rewriter.notifyMatchFailure(
op, "Expected descriptor rank to match source rank; subview the "
"source down to the descriptor rank.");
if (rank > maxNdTdescRank)
return rewriter.notifyMatchFailure(
op, "Batched nd descriptor supports at most " +
std::to_string(maxNdTdescLeadingDims) +
" leading dims (rank <= " + std::to_string(maxNdTdescRank) +
").");
// Limitation 1 above; dynamic strides are assumed to divide evenly.
if (rank > 2) {
SmallVector<std::optional<int64_t>> constStrides(rank, std::nullopt);
if (memrefTy) {
SmallVector<int64_t> staticStrides;
int64_t staticOffset;
if (succeeded(
memrefTy.getStridesAndOffset(staticStrides, staticOffset)))
for (int64_t d = 0; d < rank; ++d)
if (!ShapedType::isDynamic(staticStrides[d]))
constStrides[d] = staticStrides[d];
} else {
SmallVector<OpFoldResult> mixed = op.getMixedStrides();
for (int64_t d = 0; d < rank; ++d)
constStrides[d] = getConstantIntValue(mixed[d]);
}
if (std::optional<int64_t> pitch = constStrides[rank - 2]) {
for (int64_t d = 0; d < rank - 2; ++d) {
std::optional<int64_t> leading = constStrides[d];
if (leading && (*pitch == 0 || *leading % *pitch != 0))
return rewriter.notifyMatchFailure(
op, "Expected each leading (batch) stride to be a multiple of "
"the row stride; the source has gaps between planes.");
}
}
}
// Enforce the target's 2D block memory restrictions. A rank-1 descriptor is
// just a base address and is not consumed by a 2D block op, so it is
// exempt.
if (rank >= 2 && failed(checkMemoryRestrictions(op, rank)))
return failure();
Type payloadElemTy = rewriter.getI32Type();
Type i64Ty = rewriter.getI64Type();
// Access the adaptor only after the failure checks, so a bail-out leaves no
// materialization cast behind.
Value baseAddr = adaptor.getSource();
if (isa<IntegerType>(source.getType()) && baseAddr.getType() != i64Ty) {
// Pointer type may be i32. Cast to i64 if needed.
baseAddr = arith::ExtUIOp::create(rewriter, loc, i64Ty, baseAddr);
}
// 1D tensor descriptor is just the base address.
if (rank == 1) {
rewriter.replaceOp(op, baseAddr);
return success();
}
SmallVector<OpFoldResult> mixedSizes;
SmallVector<OpFoldResult> mixedStrides;
if (memrefTy && !xegpu::hasStaticShapeAndStrides(memrefTy)) {
auto meta =
memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
mixedSizes = meta.getConstifiedMixedSizes();
mixedStrides = meta.getConstifiedMixedStrides();
} else {
mixedSizes = op.getMixedSizes();
mixedStrides = op.getMixedStrides();
}
// Op is lowered to a code sequence that populates payload.
// Payload is a 8xi32 vector. Offset to individual fields are defined in
// NdTdescOffset enum.
VectorType payloadTy = VectorType::get(8, payloadElemTy);
// 4xi64 view is used for inserting the base pointer.
VectorType payloadI64Ty = VectorType::get(4, i64Ty);
// Initialize payload to zero.
Value payload = arith::ConstantOp::create(
rewriter, loc,
DenseElementsAttr::get(payloadTy, IntegerAttr::get(payloadElemTy, 0)));
// Utility for creating offset values from op fold result.
auto createOffset = [&](SmallVector<OpFoldResult> &ofrVec,
unsigned idx) -> Value {
Value val = getValueOrCreateConstantIntOp(rewriter, loc, ofrVec[idx]);
val = getValueOrCreateCastToIndexLike(rewriter, loc, payloadElemTy, val);
return val;
};
// The descriptor's innermost 2 dims are the 2D tile (H, W).
Value baseShapeW = createOffset(mixedSizes, rank - 1);
// Pitch is the stride of dim rank-2 (the row stride of the 2D tile).
Value basePitch = createOffset(mixedStrides, rank - 2);
// Leading (batch) strides as a number of rows (`stride[d] / pitch`).
SmallVector<Value> leadingRowStrides;
for (int64_t d = 0; d < rank - 2; ++d) {
std::optional<int64_t> leading = getConstantIntValue(mixedStrides[d]);
std::optional<int64_t> pitch =
getConstantIntValue(mixedStrides[rank - 2]);
if (leading && pitch && *pitch != 0)
leadingRowStrides.push_back(arith::ConstantIntOp::create(
rewriter, loc, payloadElemTy, *leading / *pitch));
else
leadingRowStrides.push_back(arith::DivUIOp::create(
rewriter, loc, createOffset(mixedStrides, d), basePitch));
}
// Height of the flattened plane: every leading (batch) plane is stacked
// into the surface, so the boundary check covers an out-of-range batch.
// base_height = size[R-2] + sum_d (size[d] - 1) * leadingRowStride[d]
// For rank 2 this is just size[0].
Value baseShapeH = createOffset(mixedSizes, rank - 2);
if (rank > 2) {
Value one = arith::ConstantIntOp::create(rewriter, loc, payloadElemTy, 1);
for (int64_t d = 0; d < rank - 2; ++d) {
Value planesBelow = rewriter.createOrFold<arith::SubIOp>(
loc, createOffset(mixedSizes, d), one);
Value rows = rewriter.createOrFold<arith::MulIOp>(loc, planesBelow,
leadingRowStrides[d]);
baseShapeH =
rewriter.createOrFold<arith::AddIOp>(loc, baseShapeH, rows);
}
}
// Populate payload.
Value payLoadAsI64 =
vector::BitCastOp::create(rewriter, loc, payloadI64Ty, payload);
payLoadAsI64 =
vector::InsertOp::create(rewriter, loc, baseAddr, payLoadAsI64,
static_cast<int>(NdTdescOffset::BasePtr));
payload = vector::BitCastOp::create(rewriter, loc, payloadTy, payLoadAsI64);
payload =
vector::InsertOp::create(rewriter, loc, baseShapeW, payload,
static_cast<int>(NdTdescOffset::BaseShapeW));
payload =
vector::InsertOp::create(rewriter, loc, baseShapeH, payload,
static_cast<int>(NdTdescOffset::BaseShapeH));
payload =
vector::InsertOp::create(rewriter, loc, basePitch, payload,
static_cast<int>(NdTdescOffset::BasePitch));
// The leading (batch) row strides go into the spare payload slots; the
// load/store/prefetch lowering turns the batch offsets into a row offset
// with them.
for (int64_t d = 0; d < rank - 2; ++d)
payload = vector::InsertOp::create(
rewriter, loc, leadingRowStrides[d], payload,
static_cast<int>(NdTdescOffset::LeadingStride0) + d);
rewriter.replaceOp(op, payload);
return success();
}
};
template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
OpType, xegpu::LoadNdOp, xegpu::StoreNdOp, xegpu::PrefetchNdOp>::value>>
class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
using OpConversionPattern<OpType>::OpConversionPattern;
LogicalResult
matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto mixedOffsets = op.getMixedOffsets();
int64_t opOffsetsSize = mixedOffsets.size();
auto loc = op.getLoc();
auto ctxt = rewriter.getContext();
auto tdesc = adaptor.getTensorDesc();
auto tdescTy = op.getTensorDescType();
auto tileRank = tdescTy.getRank();
if (opOffsetsSize != tileRank)
return rewriter.notifyMatchFailure(
op, "Expected offset rank to match descriptor rank.");
if (tileRank > 2 && llvm::any_of(tdescTy.getShape().drop_back(2),
[](int64_t d) { return d != 1; }))
return rewriter.notifyMatchFailure(
op, "Expected leading (batch) descriptor dims to be unit.");
if (tileRank > maxNdTdescRank)
return rewriter.notifyMatchFailure(
op, "Expected descriptor rank <= " + std::to_string(maxNdTdescRank) +
".");
auto elemType = tdescTy.getElementType();
auto elemBitSize = elemType.getIntOrFloatBitWidth();
bool isSubByte = elemBitSize < 8;
uint64_t wScaleFactor = 1;
if (!isSubByte && (elemBitSize % 8 != 0))
return rewriter.notifyMatchFailure(
op, "Expected element type bit width to be multiple of 8.");
auto tileW = tdescTy.getDimSize(tileRank - 1);
// For sub byte types, only 4bits are currently supported.
if (isSubByte) {
if (elemBitSize != 4)
return rewriter.notifyMatchFailure(
op, "Only sub byte types of 4bits are supported.");
if (tileRank != 2)
return rewriter.notifyMatchFailure(
op, "Sub byte types are only supported for 2D tensor descriptors.");
auto subByteFactor = 8 / elemBitSize;
auto tileH = tdescTy.getDimSize(0);
// Handle special case for packed load.
if constexpr (std::is_same_v<OpType, xegpu::LoadNdOp>) {
if (op.getPacked().value_or(false)) {
// packed load is implemented as packed loads of 8bit elements.
if (tileH == systolicDepth * 4 &&
tileW == executionSize * subByteFactor) {
// Usage case for loading as Matrix B with pack request.
// source is assumed to pre-packed into 8bit elements
// Emulate with 8bit loads with pack request.
// scaled_tileW = executionSize
elemType = rewriter.getIntegerType(8);
tileW = executionSize;
wScaleFactor = subByteFactor;
}
}
}
// If not handled by packed load case above, handle other cases.
if (wScaleFactor == 1) {
auto sub16BitFactor = subByteFactor * 2;
if (tileW == executionSize * sub16BitFactor) {
// Usage case for loading as Matrix A operand
// Emulate with 16bit loads/stores.
// scaled_tileW = executionSize
elemType = rewriter.getIntegerType(16);
tileW = executionSize;
wScaleFactor = sub16BitFactor;
} else {
return rewriter.notifyMatchFailure(
op, "Unsupported tile shape for sub byte types.");
}
}
// recompute element bit size for emulation.
elemBitSize = elemType.getIntOrFloatBitWidth();
}
// Get address space from tensor descriptor memory space.
auto ptrTypeLLVM = LLVM::LLVMPointerType::get(
ctxt, getNumericXeVMAddrSpace(tdescTy.getMemorySpace()));
if (tileRank >= 2) {
// Compute element byte size.
Value elemByteSize = arith::ConstantIntOp::create(
rewriter, loc, rewriter.getI32Type(), elemBitSize / 8);
VectorType payloadI64Ty = VectorType::get(4, rewriter.getI64Type());
Value payLoadAsI64 =
vector::BitCastOp::create(rewriter, loc, payloadI64Ty, tdesc);
Value basePtr =
vector::ExtractOp::create(rewriter, loc, payLoadAsI64,
static_cast<int>(NdTdescOffset::BasePtr));
Value baseShapeW = vector::ExtractOp::create(
rewriter, loc, tdesc, static_cast<int>(NdTdescOffset::BaseShapeW));
Value baseShapeH = vector::ExtractOp::create(
rewriter, loc, tdesc, static_cast<int>(NdTdescOffset::BaseShapeH));
Value basePitch = vector::ExtractOp::create(
rewriter, loc, tdesc, static_cast<int>(NdTdescOffset::BasePitch));
Value offsetW = getValueOrCreateConstantIntOp(rewriter, loc,
mixedOffsets[tileRank - 1]);
offsetW = getValueOrCreateCastToIndexLike(rewriter, loc,
rewriter.getI32Type(), offsetW);
Value offsetH = getValueOrCreateConstantIntOp(rewriter, loc,
mixedOffsets[tileRank - 2]);
offsetH = getValueOrCreateCastToIndexLike(rewriter, loc,
rewriter.getI32Type(), offsetH);
// Turn the leading (batch) offsets into a row offset into the flattened
// plane, using the row-unit batch strides encoded at create time:
// offsetH += sum_d offset[d] * leadingRowStride[d]
// The base pointer stays at the true base, so an out-of-range batch index
// is caught by the HW boundary check instead of moving the surface to
// unmapped memory.
for (int64_t d = 0; d < tileRank - 2; ++d) {
Value off =
getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[d]);
off = getValueOrCreateCastToIndexLike(rewriter, loc,
rewriter.getI32Type(), off);
Value rowStride = vector::ExtractOp::create(
rewriter, loc, tdesc,
static_cast<int>(NdTdescOffset::LeadingStride0) + d);
Value term = arith::MulIOp::create(rewriter, loc, off, rowStride);
offsetH = arith::AddIOp::create(rewriter, loc, offsetH, term);
}
// Convert base pointer (i64) to LLVM pointer type.
Value basePtrLLVM =
LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtr);
// FIXME: width or pitch is not the same as baseShapeW it should be the
// stride of the second to last dimension in row major layout.
// Compute width in bytes.
Value baseShapeWInBytes =
arith::MulIOp::create(rewriter, loc, baseShapeW, elemByteSize);
// Compute pitch in bytes.
Value basePitchBytes =
arith::MulIOp::create(rewriter, loc, basePitch, elemByteSize);
if (wScaleFactor > 1) {
// Scale offsetW, baseShapeWInBytes for sub byte emulation.
// Note: tileW is already scaled above.
Value wScaleFactorValLog2 = arith::ConstantIntOp::create(
rewriter, loc, rewriter.getI32Type(), llvm::Log2_64(wScaleFactor));
baseShapeWInBytes = arith::ShRSIOp::create(
rewriter, loc, baseShapeWInBytes, wScaleFactorValLog2);
basePitchBytes = arith::ShRSIOp::create(rewriter, loc, basePitchBytes,
wScaleFactorValLog2);
offsetW =
arith::ShRSIOp::create(rewriter, loc, offsetW, wScaleFactorValLog2);
}
// Get tile height from the tensor descriptor type (second-to-last dim).
auto tileH = tdescTy.getDimSize(tileRank - 2);
// Get vblocks from the tensor descriptor type.
int32_t vblocks = tdescTy.getArrayLength();
if constexpr (std::is_same_v<OpType, xegpu::StoreNdOp>) {
Value src = adaptor.getValue();
// If store value is a scalar, get value from op instead of adaptor.
// Adaptor might have optimized away single element vector
if (src.getType().isIntOrFloat()) {
src = op.getValue();
}
VectorType srcVecTy = dyn_cast<VectorType>(src.getType());
if (!srcVecTy)
return rewriter.notifyMatchFailure(
op, "Expected store value to be a vector type.");
// Get flat vector type of integer type with matching element bit size.
VectorType newSrcVecTy =
encodeVectorTypeTo(srcVecTy, rewriter.getIntegerType(elemBitSize));
if (srcVecTy != newSrcVecTy)
src = vector::BitCastOp::create(rewriter, loc, newSrcVecTy, src);
auto storeCacheControl =
translateStoreXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
xevm::BlockStore2dOp::create(
rewriter, loc, basePtrLLVM, baseShapeWInBytes, baseShapeH,
basePitchBytes, offsetW, offsetH, elemBitSize, tileW, tileH, src,
xevm::StoreCacheControlAttr::get(ctxt, storeCacheControl));
rewriter.eraseOp(op);
} else {
auto loadCacheControl =
translateLoadXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
if constexpr (std::is_same_v<OpType, xegpu::PrefetchNdOp>) {
xevm::BlockPrefetch2dOp::create(
rewriter, loc, basePtrLLVM, baseShapeWInBytes, baseShapeH,
basePitchBytes, offsetW, offsetH, elemBitSize, tileW, tileH,
vblocks, xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
rewriter.eraseOp(op);
} else {
VectorType dstVecTy = cast<VectorType>(op.getValue().getType());
bool vnni = op.getPacked().value_or(false);
auto transposeValue = op.getTranspose();
bool transpose =
transposeValue.has_value() && transposeValue.value()[0] == 1;
// Handle special case of 32x16 and 8bit element load
// with no vnni, no transpose, no vblocks.
// For this special case, vnni and non vnni yields the same output
// and only the vnni variant is supported by HW.
// Check and set vnni of the special case.
if (elemBitSize == 8 && tileW == 16 && tileH == 32 && !vnni &&
!transpose) {
vnni = true;
}
// Handle tranpose request on small element size
// Transpose needs to be requested on 32bit element type.
// offsetW and tileW needs to be adjusted to account for element type
// change.
if (transpose && elemBitSize < 32) {
int32_t scale = 32 / elemBitSize;
Value scaleLog2 = arith::ConstantIntOp::create(
rewriter, loc, rewriter.getI32Type(), llvm::Log2_64(scale));
offsetW = arith::ShRSIOp::create(rewriter, loc, offsetW, scaleLog2);
tileW = tileW * elemBitSize / 32;
elemBitSize = 32;
}
VectorType loadedTy = encodeVectorTypeTo(
dstVecTy, vnni ? rewriter.getI32Type()
: rewriter.getIntegerType(elemBitSize));
Value resultFlatVec = xevm::BlockLoad2dOp::create(
rewriter, loc, loadedTy, basePtrLLVM, baseShapeWInBytes,
baseShapeH, basePitchBytes, offsetW, offsetH, elemBitSize, tileW,
tileH, vblocks, transpose, vnni,
xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
resultFlatVec = vector::BitCastOp::create(
rewriter, loc,
encodeVectorTypeTo(loadedTy, dstVecTy.getElementType()),
resultFlatVec);
rewriter.replaceOp(op, resultFlatVec);
}
}
} else {
// 1D tensor descriptor.
// `tdesc` represents base address as i64
// Offset in number of elements, need to multiply by element byte size.
// Compute byte offset.
// byteOffset = offset * elementByteSize
Value offset =
getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[0]);
offset = getValueOrCreateCastToIndexLike(rewriter, loc,
rewriter.getI64Type(), offset);
// Compute element byte size.
Value elemByteSize = arith::ConstantIntOp::create(
rewriter, loc, rewriter.getI64Type(), elemBitSize / 8);
Value byteOffset =
rewriter.createOrFold<arith::MulIOp>(loc, offset, elemByteSize);
// Final address = basePtr + byteOffset
Value finalAddrI64 = rewriter.createOrFold<arith::AddIOp>(
loc, tdesc,
getValueOrCreateCastToIndexLike(rewriter, loc, rewriter.getI64Type(),
byteOffset));
// Convert base pointer (i64) to LLVM pointer type.
Value finalPtrLLVM =
LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, finalAddrI64);
if constexpr (std::is_same_v<OpType, xegpu::StoreNdOp>) {
Value src = adaptor.getValue();
// If store value is a scalar, get value from op instead of adaptor.
// Adaptor might have optimized away single element vector
if (src.getType().isIntOrFloat()) {
src = op.getValue();
}
VectorType srcVecTy = dyn_cast<VectorType>(src.getType());
if (!srcVecTy)
return rewriter.notifyMatchFailure(
op, "Expected store value to be a vector type.");
// Get flat vector type of integer type with matching element bit size.
VectorType newSrcVecTy =
encodeVectorTypeTo(srcVecTy, rewriter.getIntegerType(elemBitSize));
if (srcVecTy != newSrcVecTy)
src = vector::BitCastOp::create(rewriter, loc, newSrcVecTy, src);
auto storeCacheControl =
translateStoreXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
rewriter.replaceOpWithNewOp<xevm::BlockStoreOp>(
op, finalPtrLLVM, src,
xevm::StoreCacheControlAttr::get(ctxt, storeCacheControl));
} else if constexpr (std::is_same_v<OpType, xegpu::LoadNdOp>) {
auto loadCacheControl =
translateLoadXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
VectorType resTy = cast<VectorType>(op.getValue().getType());
VectorType loadedTy =
encodeVectorTypeTo(resTy, rewriter.getIntegerType(elemBitSize));
Value load = xevm::BlockLoadOp::create(
rewriter, loc, loadedTy, finalPtrLLVM,
xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
if (loadedTy != resTy)
load = vector::BitCastOp::create(rewriter, loc, resTy, load);
rewriter.replaceOp(op, load);
} else {
return rewriter.notifyMatchFailure(
op, "Unsupported operation: xegpu.prefetch_nd with tensor "
"descriptor rank == 1");
}
}
return success();
}
};
// Add a builder that creates
// offset * elemByteSize + baseAddr
static Value addOffsetToBaseAddr(ConversionPatternRewriter &rewriter,
Location loc, Value baseAddr, Value offset,
int64_t elemByteSize) {
Value byteSize = arith::ConstantIntOp::create(
rewriter, loc, baseAddr.getType(), elemByteSize);
Value byteOffset = arith::MulIOp::create(rewriter, loc, offset, byteSize);
Value newAddr = arith::AddIOp::create(rewriter, loc, baseAddr, byteOffset);
return newAddr;
}
// Returns true when every element of `mask` carries the same bit, so gating a
// whole contiguous block on element 0 is equivalent. Splat constants, an
// all-ones or all-zeros `vector.constant_mask`, broadcasts of a scalar, and a
// `vector.from_elements` of one repeated value qualify.
//
// TODO: this recognizes a fixed list of producers. A general uniformity query
// on the vector dialect would cover more forms and would not need updating
// every time a new producer shows up here.
static bool isUniformMask(Value mask) {
if (!isa<VectorType>(mask.getType()))
return true;
DenseElementsAttr splat;
if (matchPattern(mask, m_Constant(&splat)) && splat.isSplat())
return true;
if (auto constantMask = mask.getDefiningOp<vector::ConstantMaskOp>()) {
// All-ones and all-zeros are uniform. A zero dim size is only legal when
// every dim is zero, so those are the only two uniform cases.
return constantMask.isAllOnesMask() ||
llvm::all_of(constantMask.getMaskDimSizes(),
[](int64_t size) { return size == 0; });
}
if (auto broadcast = mask.getDefiningOp<vector::BroadcastOp>())
return !isa<VectorType>(broadcast.getSource().getType());
if (auto fromElements = mask.getDefiningOp<vector::FromElementsOp>())
return llvm::all_equal(fromElements.getElements());
// Flattening a rank > 1 operand to rank 1 inserts a shape cast. It keeps the
// same elements, so it keeps uniformity.
if (auto shapeCast = mask.getDefiningOp<vector::ShapeCastOp>())
return isUniformMask(shapeCast.getSource());
return false;
}
template <typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
OpType, xegpu::LoadGatherOp, xegpu::StoreScatterOp>::value>>
class LoadStoreToXeVMPattern : public OpConversionPattern<OpType> {
using OpConversionPattern<OpType>::OpConversionPattern;
LogicalResult
matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
Value offset = adaptor.getOffsets();
if (!offset)
return rewriter.notifyMatchFailure(op, "Expected offset to be provided.");
auto loc = op.getLoc();
auto ctxt = rewriter.getContext();
Value basePtrI64;
// Load result or Store valye Type can be vector or scalar.
Type valOrResTy;
if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>)
valOrResTy =
this->getTypeConverter()->convertType(op.getResult().getType());
else
valOrResTy = adaptor.getValue().getType();
VectorType valOrResVecTy = dyn_cast<VectorType>(valOrResTy);
bool hasScalarVal = !valOrResVecTy;
int64_t elemBitWidth =
hasScalarVal ? valOrResTy.getIntOrFloatBitWidth()
: valOrResVecTy.getElementType().getIntOrFloatBitWidth();
// Element type must be multiple of 8 bits.
if (elemBitWidth % 8 != 0)
return rewriter.notifyMatchFailure(
op, "Expected element type bit width to be multiple of 8.");
int64_t elemByteSize = elemBitWidth / 8;
// Default memory space is global.
LLVM::LLVMPointerType ptrTypeLLVM = LLVM::LLVMPointerType::get(
ctxt, getNumericXeVMAddrSpace(xegpu::MemorySpace::Global));
// Base pointer can come from source (load) or dest (store).
// If they are memrefs, we use their memory space.
if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>) {
basePtrI64 = adaptor.getSource();
if (auto memRefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
FailureOr<unsigned> addrSpace =
getNumericMemorySpace(memRefTy.getMemorySpace());
if (failed(addrSpace))
return rewriter.notifyMatchFailure(
op, "Unsupported memref memory space attribute.");
if (*addrSpace != 0)
ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, *addrSpace);
}
} else {
basePtrI64 = adaptor.getDest();
if (auto memRefTy = dyn_cast<MemRefType>(op.getDest().getType())) {
FailureOr<unsigned> addrSpace =
getNumericMemorySpace(memRefTy.getMemorySpace());
if (failed(addrSpace))
return rewriter.notifyMatchFailure(
op, "Unsupported memref memory space attribute.");
if (*addrSpace != 0)
ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, *addrSpace);
}
}
// Base pointer is passed as i32 or i64 by adaptor, cast to i64 if needed.
if (basePtrI64.getType() != rewriter.getI64Type()) {
basePtrI64 = arith::ExtUIOp::create(rewriter, loc, rewriter.getI64Type(),
basePtrI64);
}
Value mask = adaptor.getMask();
// Coalesce a lane's multi-element access into one block access.
//
// Distribution gives every element its own offset and mask bit, so a lane
// that takes D neighbouring elements arrives with `vector<D>` offsets and
// mask. One block access needs one base offset and one mask bit, so take
// both from element 0. A non-uniform mask cannot be reduced to one bit, so
// it stays a vector and fails to match below.
auto origOffsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
if (isa<VectorType>(offset.getType()) && origOffsetsTy && valOrResVecTy &&
origOffsetsTy.getNumElements() == valOrResVecTy.getNumElements() &&
isUniformMask(op.getMask())) {
offset = vector::ExtractOp::create(rewriter, loc, offset,
ArrayRef<int64_t>{0});
if (isa<VectorType>(mask.getType()))
mask = vector::ExtractOp::create(rewriter, loc, mask,
ArrayRef<int64_t>{0});
}
if (dyn_cast<VectorType>(offset.getType())) {
// Offset needs be scalar. Single element vector is converted to scalar
// by type converter.
return rewriter.notifyMatchFailure(op, "Expected offset to be a scalar.");
} else {
// If offset is provided, we add them to the base pointer.
// Offset is in number of elements, we need to multiply by
// element byte size.
basePtrI64 =
addOffsetToBaseAddr(rewriter, loc, basePtrI64, offset, elemByteSize);
}
// Convert base pointer (i64) to LLVM pointer type.
Value basePtrLLVM =
LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);
Value maskForLane;
VectorType maskVecTy = dyn_cast<VectorType>(mask.getType());
if (maskVecTy) {
// Mask needs be scalar. Single element vector is converted to scalar by
// type converter.
return rewriter.notifyMatchFailure(op, "Expected mask to be a scalar.");
} else
maskForLane = mask;
if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>) {
scf::IfOp ifOp = scf::IfOp::create(rewriter, loc, {valOrResTy},
maskForLane, true, true);
// If mask is true,- then clause - load from memory and yield.
rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
if (!hasScalarVal)
valOrResTy = VectorType::get({valOrResVecTy.getNumElements()},
valOrResVecTy.getElementType());
Value loaded =
LLVM::LoadOp::create(rewriter, loc, valOrResTy, basePtrLLVM);
// Set cache control attribute on the load operation.
loaded.getDefiningOp()->setDiscardableAttr(
"cache_control", xevm::LoadCacheControlAttr::get(
ctxt, translateLoadXeGPUCacheHint(
op.getL1Hint(), op.getL3Hint())));
scf::YieldOp::create(rewriter, loc, ValueRange{loaded});
rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
// If mask is false - else clause -yield a vector of zeros.
auto eTy = hasScalarVal ? valOrResTy : valOrResVecTy.getElementType();
TypedAttr eVal;
if (eTy.isFloat())
eVal = FloatAttr::get(eTy, 0.0);
else
eVal = IntegerAttr::get(eTy, 0);
if (hasScalarVal)
loaded = arith::ConstantOp::create(rewriter, loc, eVal);
else
loaded = arith::ConstantOp::create(
rewriter, loc, DenseElementsAttr::get(valOrResVecTy, eVal));
scf::YieldOp::create(rewriter, loc, ValueRange{loaded});
rewriter.replaceOp(op, ifOp.getResult(0));
} else {
// If mask is true, perform the store.
scf::IfOp ifOp = scf::IfOp::create(rewriter, loc, maskForLane, false);
auto body = ifOp.getBody();
rewriter.setInsertionPointToStart(body);
auto storeOp =
LLVM::StoreOp::create(rewriter, loc, adaptor.getValue(), basePtrLLVM);
// Set cache control attribute on the store operation.
storeOp.getOperation()->setDiscardableAttr(
"cache_control", xevm::StoreCacheControlAttr::get(
ctxt, translateStoreXeGPUCacheHint(
op.getL1Hint(), op.getL3Hint())));
rewriter.eraseOp(op);
}
return success();
}
};
class CreateMemDescOpPattern final
: public OpConversionPattern<xegpu::CreateMemDescOp> {
public:
using OpConversionPattern<xegpu::CreateMemDescOp>::OpConversionPattern;
LogicalResult
matchAndRewrite(xegpu::CreateMemDescOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
rewriter.replaceOp(op, adaptor.getSource());
return success();
}
};
template <typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
OpType, xegpu::LoadMatrixOp, xegpu::StoreMatrixOp>::value>>
class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
using OpConversionPattern<OpType>::OpConversionPattern;
LogicalResult
matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
SmallVector<OpFoldResult> offsets = op.getMixedOffsets();
if (offsets.empty())
return rewriter.notifyMatchFailure(op, "Expected offset to be provided.");
auto loc = op.getLoc();
auto ctxt = rewriter.getContext();
Value baseAddr32 = adaptor.getMemDesc();
Value mdescVal = op.getMemDesc();
// Load result or Store value Type can be vector or scalar.
Type dataTy;
if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
Type resType = op.getResult().getType();
// Some transforms may leave unit dimension in the 2D vector, adaptors do
// not catch it for results.
if (auto vecType = dyn_cast<VectorType>(resType)) {
assert(llvm::count_if(vecType.getShape(),
[](int64_t d) { return d != 1; }) <= 1 &&
"Expected either 1D vector or nD with unit dimensions");
resType = VectorType::get({vecType.getNumElements()},
vecType.getElementType());
}
dataTy = resType;
} else
dataTy = adaptor.getData().getType();
VectorType valOrResVecTy = dyn_cast<VectorType>(dataTy);
if (!valOrResVecTy)
valOrResVecTy = VectorType::get(1, dataTy);
int64_t elemBitWidth =
valOrResVecTy.getElementType().getIntOrFloatBitWidth();
// Element type must be multiple of 8 bits.
if (elemBitWidth % 8 != 0)
return rewriter.notifyMatchFailure(
op, "Expected element type bit width to be multiple of 8.");
int64_t elemByteSize = elemBitWidth / 8;
// Default memory space is SLM.
LLVM::LLVMPointerType ptrTypeLLVM = LLVM::LLVMPointerType::get(
ctxt, getNumericXeVMAddrSpace(xegpu::MemorySpace::SLM));
auto mdescTy = cast<xegpu::MemDescType>(mdescVal.getType());
Value linearOffset = mdescTy.getLinearOffsets(rewriter, loc, offsets);
linearOffset = arith::IndexCastUIOp::create(
rewriter, loc, rewriter.getI32Type(), linearOffset);
Value basePtrI32 = addOffsetToBaseAddr(rewriter, loc, baseAddr32,
linearOffset, elemByteSize);
// convert base pointer (i32) to LLVM pointer type
Value basePtrLLVM =
LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI32);
if (op.getSubgroupBlockIoAttr()) {
// if the attribute 'subgroup_block_io' is set to true, it lowers to
// xevm.blockload
Type intElemTy = rewriter.getIntegerType(elemBitWidth);
VectorType intVecTy =
VectorType::get(valOrResVecTy.getShape(), intElemTy);
if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
Value loadOp = xevm::BlockLoadOp::create(
rewriter, loc, intVecTy, basePtrLLVM, /*cache_control=*/nullptr);
if (intVecTy != valOrResVecTy) {
loadOp =
vector::BitCastOp::create(rewriter, loc, valOrResVecTy, loadOp);
}
rewriter.replaceOp(op, loadOp);
} else {
Value dataToStore = adaptor.getData();
if (valOrResVecTy != intVecTy) {
dataToStore =
vector::BitCastOp::create(rewriter, loc, intVecTy, dataToStore);
}
xevm::BlockStoreOp::create(rewriter, loc, basePtrLLVM, dataToStore,
nullptr);
rewriter.eraseOp(op);
}
return success();
}
if (valOrResVecTy.getNumElements() >= 1) {
auto chipOpt = xegpu::getChipStr(op);
if (!chipOpt ||
(*chipOpt != "pvc" && *chipOpt != "bmg" && *chipOpt != "cri")) {
// the lowering for chunk load only works for pvc, bmg or cri
return rewriter.notifyMatchFailure(
op, "The lowering is specific to pvc, bmg or cri.");
}
}
if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
// The load result type is taken from the type converter. This maps
// element types that are not directly representable in LLVM (e.g.
// f8E8M0FNU) to an integer storage type of the same bit width, and
// collapses single-element vectors to a scalar, since LLVM load/store
// does not support vectors of size 1.
Type loadTy =
this->getTypeConverter()->convertType(op.getResult().getType());
auto loadOp = LLVM::LoadOp::create(rewriter, loc, loadTy, basePtrLLVM);
rewriter.replaceOp(op, loadOp);
} else {
LLVM::StoreOp::create(rewriter, loc, adaptor.getData(), basePtrLLVM);
rewriter.eraseOp(op);
}
return success();
}
};
class PrefetchToXeVMPattern : public OpConversionPattern<xegpu::PrefetchOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(xegpu::PrefetchOp op, xegpu::PrefetchOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto ctxt = rewriter.getContext();
Value basePtrI64 = adaptor.getSource();
// Base pointer is passed as i32 or i64 by adaptor, cast to i64 if needed.
if (basePtrI64.getType() != rewriter.getI64Type())
basePtrI64 = arith::ExtUIOp::create(rewriter, loc, rewriter.getI64Type(),
basePtrI64);
Value offsets = adaptor.getOffsets();
if (offsets) {
VectorType offsetsVecTy = dyn_cast<VectorType>(offsets.getType());
if (offsetsVecTy) {
// Offset needs be scalar.
return rewriter.notifyMatchFailure(op,
"Expected offsets to be a scalar.");
} else {
int64_t elemBitWidth{0};
int64_t elemByteSize;
// Element byte size can come from two sources:
if (auto memRefTy = dyn_cast<MemRefType>(op.getSourceType())) {
// If memref is available, we use its element type to
// determine element byte size.
elemBitWidth = memRefTy.getElementType().getIntOrFloatBitWidth();
} else {
// Otherwise, we use the provided offset byte alignment.
elemByteSize = *op.getOffsetAlignByte();
}
if (elemBitWidth != 0) {
if (elemBitWidth % 8 != 0)
return rewriter.notifyMatchFailure(
op, "Expected element type bit width to be multiple of 8.");
elemByteSize = elemBitWidth / 8;
}
basePtrI64 = addOffsetToBaseAddr(rewriter, loc, basePtrI64, offsets,
elemByteSize);
}
}
// Default memory space is global.
LLVM::LLVMPointerType ptrTypeLLVM = LLVM::LLVMPointerType::get(
ctxt, getNumericXeVMAddrSpace(xegpu::MemorySpace::Global));
// If source is a memref, we use its memory space.
if (auto memRefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
FailureOr<unsigned> addrSpace =
getNumericMemorySpace(memRefTy.getMemorySpace());
if (failed(addrSpace))
return rewriter.notifyMatchFailure(
op, "Unsupported memref memory space attribute.");
if (*addrSpace != 0)
ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, *addrSpace);
}
// Convert base pointer (i64) to LLVM pointer type.
Value ptrLLVM =
LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);
// Create the prefetch op with cache control attribute.
xevm::PrefetchOp::create(
rewriter, loc, ptrLLVM,
xevm::LoadCacheControlAttr::get(
ctxt, translateLoadXeGPUCacheHint(op.getL1Hint(), op.getL3Hint())));
rewriter.eraseOp(op);
return success();
}
};
class FenceToXeVMPattern : public OpConversionPattern<xegpu::FenceOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(xegpu::FenceOp op, xegpu::FenceOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
xevm::MemScope memScope{xevm::MemScope::WORKGROUP};
switch (op.getFenceScope()) {
case xegpu::FenceScope::Workgroup:
memScope = xevm::MemScope::WORKGROUP;
break;
case xegpu::FenceScope::GPU:
memScope = xevm::MemScope::DEVICE;
break;
}
xevm::AddrSpace addrSpace{xevm::AddrSpace::GLOBAL};
switch (op.getMemoryKind()) {
case xegpu::MemorySpace::Global:
addrSpace = xevm::AddrSpace::GLOBAL;
break;
case xegpu::MemorySpace::SLM:
addrSpace = xevm::AddrSpace::SHARED;
break;
}
xevm::MemfenceOp::create(rewriter, loc, memScope, addrSpace);
rewriter.eraseOp(op);
return success();
}
};
static auto encodePrecision = [](Type type) -> xevm::ElemType {
if (type.isBF16())
return xevm::ElemType::BF16;
else if (type.isF16())
return xevm::ElemType::F16;
else if (type.isTF32())
return xevm::ElemType::TF32;
else if (type.isInteger(8)) {
if (type.isUnsignedInteger())
return xevm::ElemType::U8;
return xevm::ElemType::S8;
} else if (type.isF32())
return xevm::ElemType::F32;
else if (type.isInteger(32))
return xevm::ElemType::S32;
else if (type.isF8E5M2())
return xevm::ElemType::BF8;
else if (type.isF8E4M3FN())
return xevm::ElemType::F8;
else if (mlir::isa<Float4E2M1FNType>(type))
return xevm::ElemType::E2M1;
llvm_unreachable("add more support for ElemType");
};
static unsigned getNumOperandsPerDword(xevm::ElemType pTy) {
switch (pTy) {
case xevm::ElemType::TF32:
return 1;
case xevm::ElemType::BF16:
case xevm::ElemType::F16:
return 2;
case xevm::ElemType::U8:
case xevm::ElemType::S8:
case xevm::ElemType::F8:
case xevm::ElemType::BF8:
return 4;
case xevm::ElemType::E2M1:
return 8;
default:
llvm_unreachable("unsupported xevm::ElemType");
}
}
class DpasToXeVMPattern : public OpConversionPattern<xegpu::DpasOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(xegpu::DpasOp op, xegpu::DpasOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto ctxt = rewriter.getContext();
auto aTy = cast<VectorType>(op.getLhs().getType());
auto bTy = cast<VectorType>(op.getRhs().getType());
auto resultType = cast<VectorType>(op.getResultType());
// get the correct dpasInst by getting info from chip
auto chipStr = xegpu::getChipStr(op);
if (!chipStr)
return rewriter.notifyMatchFailure(op, "cannot determine target chip");
const auto *uArch = mlir::xegpu::uArch::getUArch(*chipStr);
if (!uArch)
return rewriter.notifyMatchFailure(op, "unsupported target uArch");
auto *dpasInst = const_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc *>(
llvm::dyn_cast_or_null<xegpu::uArch::SubgroupMatrixMultiplyAcc>(
uArch->getInstruction(
xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc)));
if (!dpasInst)
return rewriter.notifyMatchFailure(op,
"DPAS not supported by target uArch");
auto checkSupportedTypes = [&](VectorType vecTy,
xegpu::uArch::MMAOpndKind kind) -> bool {
auto supported = dpasInst->getSupportedTypes(*ctxt, kind);
return llvm::find(supported, vecTy.getElementType()) != supported.end();
};
if (!checkSupportedTypes(aTy, xegpu::uArch::MMAOpndKind::MatrixA))
return rewriter.notifyMatchFailure(
op, "A-matrix element type not supported by target uArch");
if (!checkSupportedTypes(bTy, xegpu::uArch::MMAOpndKind::MatrixB))
return rewriter.notifyMatchFailure(
op, "B-matrix element type not supported by target uArch");
// NOTE: Supported types for MatrixC and MatrixD are identical
if (!checkSupportedTypes(resultType, xegpu::uArch::MMAOpndKind::MatrixD))
return rewriter.notifyMatchFailure(
op, "result/accumulator element type not supported by target uArch");
xevm::ElemType precATy = encodePrecision(aTy.getElementType());
xevm::ElemType precBTy = encodePrecision(bTy.getElementType());
Value c = op.getAcc();
if (!c) {
auto elementTy = resultType.getElementType();
Attribute initValueAttr;
if (isa<FloatType>(elementTy))
initValueAttr = FloatAttr::get(elementTy, 0.0);
else
initValueAttr = IntegerAttr::get(elementTy, 0);
c = arith::ConstantOp::create(
rewriter, loc, DenseElementsAttr::get(resultType, initValueAttr));
}
Value aVec = op.getLhs();
Value bVec = op.getRhs();
auto cvecty = cast<VectorType>(c.getType());
xevm::ElemType precCTy = encodePrecision(cvecty.getElementType());
xevm::ElemType precDTy = encodePrecision(resultType.getElementType());
VectorType cNty =
VectorType::get(cvecty.getNumElements(), cvecty.getElementType());
if (cvecty != cNty)
c = vector::ShapeCastOp::create(rewriter, loc, cNty, c);
Value dpasRes = xevm::MMAOp::create(
rewriter, loc, cNty, aVec, bVec, c,
xevm::MMAShapeAttr::get(ctxt, cvecty.getNumElements(), executionSize,
systolicDepth *
getNumOperandsPerDword(precATy)),
xevm::MMATypesAttr::get(ctxt, precDTy, precATy, precBTy, precCTy));
if (cvecty != cNty)
dpasRes = vector::ShapeCastOp::create(rewriter, loc, resultType, dpasRes);
rewriter.replaceOp(op, dpasRes);
return success();
}
};
static std::optional<LLVM::AtomicBinOp>
matchSimpleAtomicOp(arith::AtomicRMWKind arithKind) {
switch (arithKind) {
case arith::AtomicRMWKind::addf:
return LLVM::AtomicBinOp::fadd;
case arith::AtomicRMWKind::addi:
return LLVM::AtomicBinOp::add;
case arith::AtomicRMWKind::assign:
return LLVM::AtomicBinOp::xchg;
case arith::AtomicRMWKind::maximumf:
return LLVM::AtomicBinOp::fmax;
case arith::AtomicRMWKind::maxs:
return LLVM::AtomicBinOp::max;
case arith::AtomicRMWKind::maxu:
return LLVM::AtomicBinOp::umax;
case arith::AtomicRMWKind::minimumf:
return LLVM::AtomicBinOp::fmin;
case arith::AtomicRMWKind::mins:
return LLVM::AtomicBinOp::min;
case arith::AtomicRMWKind::minu:
return LLVM::AtomicBinOp::umin;
case arith::AtomicRMWKind::ori:
return LLVM::AtomicBinOp::_or;
case arith::AtomicRMWKind::andi:
return LLVM::AtomicBinOp::_and;
default:
return std::nullopt;
}
}
class AtomicRMWToXeVMPattern : public OpConversionPattern<xegpu::AtomicRMWOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(xegpu::AtomicRMWOp op, xegpu::AtomicRMWOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto ctxt = rewriter.getContext();
auto tdesc = op.getTensorDesc().getType();
auto ptrTypeLLVM = LLVM::LLVMPointerType::get(
ctxt, getNumericXeVMAddrSpace(tdesc.getMemorySpace()));
Value basePtrI64 = arith::IndexCastOp::create(
rewriter, loc, rewriter.getI64Type(), adaptor.getTensorDesc());
Value basePtrLLVM =
LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);
VectorType srcOrDstVecTy = cast<VectorType>(op.getValue().getType());
VectorType srcOrDstFlatVecTy = VectorType::get(
srcOrDstVecTy.getNumElements(), srcOrDstVecTy.getElementType());
Value srcFlatVec = vector::ShapeCastOp::create(
rewriter, loc, srcOrDstFlatVecTy, op.getValue());
auto atomicKind = matchSimpleAtomicOp(op.getKind());
assert(atomicKind.has_value());
Value resVec = srcFlatVec;
for (int i = 0; i < srcOrDstVecTy.getNumElements(); i++) {
auto val = vector::ExtractOp::create(rewriter, loc, resVec, i);
Value idx = LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64Type(),
rewriter.getI64IntegerAttr(i));
Value currPtr =
LLVM::GEPOp::create(rewriter, loc, ptrTypeLLVM,
srcOrDstVecTy.getElementType(), basePtrLLVM, idx);
Value newVal =
LLVM::AtomicRMWOp::create(rewriter, loc, atomicKind.value(), currPtr,
val, LLVM::AtomicOrdering::seq_cst);
resVec = vector::InsertOp::create(rewriter, loc, newVal, resVec, i);
}
rewriter.replaceOp(op, resVec);
return success();
}
};
class DpasMxToXeVMPattern : public OpConversionPattern<xegpu::DpasMxOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(xegpu::DpasMxOp op, xegpu::DpasMxOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto ctxt = rewriter.getContext();
auto aTy = op.getA().getType();
auto bTy = op.getB().getType();
auto resVecTy =
cast<VectorType>(getTypeConverter()->convertType(op.getType()));
auto chipStr = xegpu::getChipStr(op);
if (!chipStr)
return rewriter.notifyMatchFailure(op, "cannot determine target chip");
const auto *uArch = xegpu::uArch::getUArch(*chipStr);
if (!uArch)
return rewriter.notifyMatchFailure(op, "unsupported target uArch");
// TODO: Add supported shape check
xevm::ElemType precATy = encodePrecision(aTy.getElementType());
xevm::ElemType precBTy = encodePrecision(bTy.getElementType());
Value c = adaptor.getAcc();
if (!c) {
auto elementTy = resVecTy.getElementType();
Attribute initValueAttr;
if (isa<FloatType>(elementTy))
initValueAttr = FloatAttr::get(elementTy, 0.0);
else
initValueAttr = IntegerAttr::get(elementTy, 0);
c = arith::ConstantOp::create(
rewriter, loc, DenseElementsAttr::get(resVecTy, initValueAttr));
}
Value aVec = adaptor.getA();
Value bVec = adaptor.getB();
auto aVecTy = cast<VectorType>(aVec.getType());
auto bVecTy = cast<VectorType>(bVec.getType());
if (aVecTy.getElementTypeBitWidth() == 4)
aVec = vector::BitCastOp::create(
rewriter, loc,
VectorType::get(aVecTy.getNumElements() / 2, rewriter.getI8Type()),
aVec);
if (bVecTy.getElementTypeBitWidth() == 4)
bVec = vector::BitCastOp::create(
rewriter, loc,
VectorType::get(bVecTy.getNumElements() / 2, rewriter.getI8Type()),
bVec);
auto cVecTy = cast<VectorType>(c.getType());
xevm::ElemType precCTy = encodePrecision(cVecTy.getElementType());
xevm::ElemType precDTy = encodePrecision(resVecTy.getElementType());
Value scaleA = adaptor.getScaleA();
Value scaleB = adaptor.getScaleB();
Value dpasMxRes = xevm::MMAMxOp::create(
rewriter, loc, resVecTy, aVec, bVec, scaleA, scaleB, c,
xevm::MMAShapeAttr::get(ctxt, cVecTy.getNumElements(), executionSize,
systolicDepth *
getNumOperandsPerDword(precATy)),
xevm::MMATypesAttr::get(ctxt, precDTy, precATy, precBTy, precCTy));
rewriter.replaceOp(op, dpasMxRes);
return success();
}
};
//===----------------------------------------------------------------------===//
// arith.extf / arith.truncf to xevm.extf / xevm.truncf
//===----------------------------------------------------------------------===//
//
// Micro-scaling (MX) GEMM lowering breaks arith.scaling_extf/scaling_truncf
// into plain arith.extf/arith.truncf whose narrow side uses one of the MX float
// formats (f8E5M2, f8E4M3FN or f4E2M1FN). These narrow floats have no native
// LLVM support, so the conversions are mapped onto the dedicated xevm.extf /
// xevm.truncf ops which lower to hardware builtins. The f8E8M0FNU scale type is
// intentionally not handled here: it is expanded into integer arithmetic by
// arith-expand before this pass runs.
// xevm.extf / xevm.truncf only convert between the MX narrow floats and
// f16/bf16, and the underlying builtins operate on exactly 16 f16/bf16 values.
static constexpr int64_t kXeVMExtfTruncfNumElems = 16;
// Maps a narrow MX float element type to the matching xevm.extf source enum.
static std::optional<xevm::ExtfSrcElemTypes> getExtfNarrowType(Type etype) {
if (isa<Float8E5M2Type>(etype))
return xevm::ExtfSrcElemTypes::BF8;
if (isa<Float8E4M3FNType>(etype))
return xevm::ExtfSrcElemTypes::F8;
if (isa<Float4E2M1FNType>(etype))
return xevm::ExtfSrcElemTypes::E2M1;
return std::nullopt;
}
// Maps a narrow MX float element type to the matching xevm.truncf dest enum.
static std::optional<xevm::TruncfDstElemTypes> getTruncfNarrowType(Type etype) {
if (isa<Float8E5M2Type>(etype))
return xevm::TruncfDstElemTypes::BF8;
if (isa<Float8E4M3FNType>(etype))
return xevm::TruncfDstElemTypes::F8;
if (isa<Float4E2M1FNType>(etype))
return xevm::TruncfDstElemTypes::E2M1;
return std::nullopt;
}
// Returns true if `op` is an arith.extf that can be lowered to xevm.extf, i.e.
// a rank-1 widening from an MX narrow float to a 16-element f16/bf16 vector.
static bool isXeVMExtf(arith::ExtFOp op) {
auto srcTy = dyn_cast<VectorType>(op.getIn().getType());
auto dstTy = dyn_cast<VectorType>(op.getType());
if (!srcTy || !dstTy || srcTy.getRank() != 1 || dstTy.getRank() != 1)
return false;
if (dstTy.getNumElements() != kXeVMExtfTruncfNumElems)
return false;
Type dstETy = dstTy.getElementType();
if (!dstETy.isF16() && !dstETy.isBF16())
return false;
return getExtfNarrowType(srcTy.getElementType()).has_value();
}
// Returns true if `op` is an arith.truncf that can be lowered to xevm.truncf,
// i.e. a rank-1 truncation from an f16/bf16 vector to an MX narrow float. The
// source has to hold a whole number of the fixed-size groups xevm.truncf
// converts at a time; wider vectors are converted in several steps.
static bool isXeVMTruncf(arith::TruncFOp op) {
auto srcTy = dyn_cast<VectorType>(op.getIn().getType());
auto dstTy = dyn_cast<VectorType>(op.getType());
if (!srcTy || !dstTy || srcTy.getRank() != 1 || dstTy.getRank() != 1)
return false;
int64_t numElems = srcTy.getNumElements();
if (numElems == 0 || numElems % kXeVMExtfTruncfNumElems != 0)
return false;
Type srcETy = srcTy.getElementType();
if (!srcETy.isF16() && !srcETy.isBF16())
return false;
return getTruncfNarrowType(dstTy.getElementType()).has_value();
}
class ExtfToXeVMPattern : public OpConversionPattern<arith::ExtFOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(arith::ExtFOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
if (!isXeVMExtf(op))
return rewriter.notifyMatchFailure(op, "not a xevm.extf compatible extf");
Location loc = op.getLoc();
MLIRContext *ctx = op.getContext();
auto srcVecTy = cast<VectorType>(op.getIn().getType());
auto dstVecTy = cast<VectorType>(op.getType());
xevm::ExtfSrcElemTypes srcEnum =
*getExtfNarrowType(srcVecTy.getElementType());
xevm::ExtfDstElemTypes dstEnum = dstVecTy.getElementType().isF16()
? xevm::ExtfDstElemTypes::F16
: xevm::ExtfDstElemTypes::BF16;
// The narrow float operand has already been type-converted to an integer
// vector of the same bit width (i4 for fp4, i8 for fp8). xevm.extf takes
// the values packed into an i8 vector, so re-pack fp4 (i4) operands.
Value src = adaptor.getIn();
auto convSrcTy = cast<VectorType>(src.getType());
if (convSrcTy.getElementTypeBitWidth() == 4)
src = vector::BitCastOp::create(
rewriter, loc,
VectorType::get(convSrcTy.getNumElements() / 2, rewriter.getI8Type()),
src);
Type resTy = getTypeConverter()->convertType(dstVecTy);
Value res = xevm::ExtfOp::create(
rewriter, loc, resTy, src, xevm::ExtfSrcElemTypeAttr::get(ctx, srcEnum),
xevm::ExtfDstElemTypeAttr::get(ctx, dstEnum));
rewriter.replaceOp(op, res);
return success();
}
};
class TruncfToXeVMPattern : public OpConversionPattern<arith::TruncFOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(arith::TruncFOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
if (!isXeVMTruncf(op))
return rewriter.notifyMatchFailure(op,
"not a xevm.truncf compatible truncf");
Location loc = op.getLoc();
MLIRContext *ctx = op.getContext();
auto srcVecTy = cast<VectorType>(op.getIn().getType());
auto dstVecTy = cast<VectorType>(op.getType());
xevm::TruncfSrcElemTypes srcEnum = srcVecTy.getElementType().isF16()
? xevm::TruncfSrcElemTypes::F16
: xevm::TruncfSrcElemTypes::BF16;
xevm::TruncfDstElemTypes dstEnum =
*getTruncfNarrowType(dstVecTy.getElementType());
auto srcEnumAttr = xevm::TruncfSrcElemTypeAttr::get(ctx, srcEnum);
auto dstEnumAttr = xevm::TruncfDstElemTypeAttr::get(ctx, dstEnum);
// xevm.truncf lowers to instructions that convert a fixed number of
// elements at a time, so a wider source is converted one group at a time
// and the packed results are concatenated. Each group produces the narrow
// floats packed into an i8 vector.
int64_t numGroups = srcVecTy.getNumElements() / kXeVMExtfTruncfNumElems;
int64_t groupBytes =
kXeVMExtfTruncfNumElems * dstVecTy.getElementTypeBitWidth() / 8;
Type groupTy = VectorType::get(groupBytes, rewriter.getI8Type());
Value src = adaptor.getIn();
Value packed;
if (numGroups == 1) {
packed = xevm::TruncfOp::create(rewriter, loc, groupTy, src, srcEnumAttr,
dstEnumAttr);
} else {
auto packedTy =
VectorType::get(groupBytes * numGroups, rewriter.getI8Type());
packed = arith::ConstantOp::create(rewriter, loc, packedTy,
rewriter.getZeroAttr(packedTy));
for (int64_t group = 0; group < numGroups; group++) {
Value slice = vector::ExtractStridedSliceOp::create(
rewriter, loc, src, group * kXeVMExtfTruncfNumElems,
kXeVMExtfTruncfNumElems, /*strides=*/1);
Value converted = xevm::TruncfOp::create(rewriter, loc, groupTy, slice,
srcEnumAttr, dstEnumAttr);
packed = vector::InsertStridedSliceOp::create(
rewriter, loc, converted, packed, group * groupBytes,
/*strides=*/1);
}
}
// Re-shape to the type-converted result type (i4 vector for fp4).
Type resTy = getTypeConverter()->convertType(dstVecTy);
if (packed.getType() != resTy)
packed = vector::BitCastOp::create(rewriter, loc, resTy, packed);
rewriter.replaceOp(op, packed);
return success();
}
};
// Lowers `xegpu.lane_shuffle` to `xevm.bitcast_shuffle`.
//
// `xevm.bitcast_shuffle` concatenates the components of its operand across the
// subgroup, the first component of every lane first, and then hands chunks the
// size of a result component back out to the lanes in order. Numbering the
// elements of a `vector<NxT>` fragment held by lane `i` of a subgroup of size
// `S` by their logical position, that concatenation is exactly the `pack` mode
// input numbering `j * S + i`. Taking the result as a single `N * width(T)` bit
// scalar then hands lane `i` the logical positions `i * N .. i * N + N - 1`,
// which is the `pack` mode output numbering.
//
// So `pack` is a vector-to-scalar `xevm.bitcast_shuffle` followed by a bitcast
// back to the fragment type, and `unpack`, being its inverse, is a bitcast to
// the scalar followed by a scalar-to-vector `xevm.bitcast_shuffle`.
//
// `xevm.bitcast_shuffle` only accepts the integer types `i8`, `i16`, `i32` and
// `i64`, since it is bit-preserving and so does not depend on how the bits are
// interpreted. A fragment of a floating point type is therefore bitcast to a
// same-width integer vector on the way in and back on the way out.
class LaneShuffleToXeVMPattern
: public OpConversionPattern<xegpu::LaneShuffleOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(xegpu::LaneShuffleOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto vecTy = dyn_cast<VectorType>(adaptor.getSource().getType());
if (!vecTy)
return rewriter.notifyMatchFailure(op, "Expected a vector fragment.");
// The shuffle redistributes whole bytes between the lanes, so sub-byte
// element types, fp4 in particular, cannot be shuffled. Widths without a
// matching integer type the op accepts are rejected for the same reason.
unsigned elemBits = vecTy.getElementTypeBitWidth();
if (elemBits != 8 && elemBits != 16 && elemBits != 32 && elemBits != 64)
return rewriter.notifyMatchFailure(
op, "Expected an element type of 8, 16, 32 or 64 bits.");
int64_t fragmentBits = vecTy.getNumElements() * elemBits;
if (fragmentBits > 64 || !llvm::isPowerOf2_64(fragmentBits))
return rewriter.notifyMatchFailure(
op, "Expected a fragment of 8, 16, 32 or 64 bits.");
Location loc = op.getLoc();
Type packedTy = rewriter.getIntegerType(fragmentBits);
// The integer vector type the shuffle actually operates on. Equal to the
// fragment type when that is already an integer vector.
VectorType shuffleTy =
VectorType::get(vecTy.getShape(), rewriter.getIntegerType(elemBits));
Value res;
if (op.getMode() == xegpu::LaneShuffleMode::Pack) {
Value src = adaptor.getSource();
if (shuffleTy != vecTy)
src = LLVM::BitcastOp::create(rewriter, loc, shuffleTy, src);
res = xevm::BitcastShuffleOp::create(rewriter, loc, packedTy, src);
res = LLVM::BitcastOp::create(rewriter, loc, vecTy, res);
} else {
Value packed =
LLVM::BitcastOp::create(rewriter, loc, packedTy, adaptor.getSource());
res = xevm::BitcastShuffleOp::create(rewriter, loc, shuffleTy, packed);
if (shuffleTy != vecTy)
res = LLVM::BitcastOp::create(rewriter, loc, vecTy, res);
}
rewriter.replaceOp(op, res);
return success();
}
};
//===----------------------------------------------------------------------===//
// Pass Definition
//===----------------------------------------------------------------------===//
struct ConvertXeGPUToXeVMPass
: public impl::ConvertXeGPUToXeVMPassBase<ConvertXeGPUToXeVMPass> {
using Base::Base;
void runOnOperation() override {
MLIRContext *context = &getContext();
// XeVM type converter is based on LLVM type converter with the
// following customizations.
// First, type conversion rules are added for xegpu custom types,
// TensorDescType and MemDescType.
// Second, MemRefType is lowered to single integer type
// Third, VectorType of single element or 0D is converted to vector
// element type. Otherwise, vector type is flatten to 1D.
LowerToLLVMOptions options(context);
options.overrideIndexBitwidth(this->use64bitIndex ? 64 : 32);
LLVMTypeConverter typeConverter(context, options);
Type xevmIndexType = typeConverter.convertType(IndexType::get(context));
Type i32Type = IntegerType::get(context, 32);
typeConverter.addConversion([&](VectorType type) -> Type {
auto elemType = typeConverter.convertType(type.getElementType());
// If the vector rank is 0 or has a single element, return the element
unsigned rank = type.getRank();
if (rank == 0 || type.getNumElements() == 1)
return elemType;
// Otherwise, convert the vector to a flat vector type.
int64_t sum = llvm::product_of(type.getShape());
return VectorType::get(sum, elemType);
});
typeConverter.addConversion([&](xegpu::TensorDescType type) -> Type {
if (type.getRank() == 1)
return xevmIndexType;
return VectorType::get(8, i32Type);
});
// SLM access related type conversions.
// TODO: LLVM DLTI provides clean way of representing different pointer size
// based on address space. Currently pointer size of SLM access is hard
// coded to 32bit. Update to use DLTI when switching overall XeGPU lowering
// to use DLTI instead of use64bitIndex option used above.
// Convert MemDescType into i32 for SLM
typeConverter.addConversion(
[&](xegpu::MemDescType type) -> Type { return i32Type; });
typeConverter.addConversion([&](MemRefType type) -> Type {
return isSharedMemRef(type) ? i32Type : xevmIndexType;
});
// LLVM type converter puts unrealized casts for the following cases:
// add materialization casts to handle them.
// Materialization to convert memref to i64 or i32 depending on global/SLM
// Applies only to target materialization.
// Note: int type to memref materialization is not required as xegpu ops
// currently do not produce memrefs as result.
auto memrefToIntMaterializationCast = [](OpBuilder &builder, Type type,
ValueRange inputs,
Location loc) -> Value {
if (inputs.size() != 1)
return {};
auto input = inputs.front();
if (auto memrefTy = dyn_cast<MemRefType>(input.getType())) {
unsigned rank = memrefTy.getRank();
Type indexType = builder.getIndexType();
int64_t intOffsets;
SmallVector<int64_t> intStrides;
Value addr;
Value offset;
if (succeeded(memrefTy.getStridesAndOffset(intStrides, intOffsets)) &&
ShapedType::isStatic(intOffsets)) {
addr = memref::ExtractAlignedPointerAsIndexOp::create(builder, loc,
input);
offset = arith::ConstantOp::create(builder, loc,
builder.getIndexAttr(intOffsets));
} else {
// Result types: [base_memref, offset, stride0, stride1, ...,
// strideN-1, size0, size1, ..., sizeN-1]
SmallVector<Type> resultTypes{
MemRefType::get({}, memrefTy.getElementType(),
MemRefLayoutAttrInterface(),
memrefTy.getMemorySpace()),
indexType};
// strides + sizes
resultTypes.append(2 * rank, indexType);
auto meta = memref::ExtractStridedMetadataOp::create(
builder, loc, resultTypes, input);
addr = memref::ExtractAlignedPointerAsIndexOp::create(
builder, loc, meta.getBaseBuffer());
offset = meta.getOffset();
}
auto addrCasted =
arith::IndexCastUIOp::create(builder, loc, type, addr);
auto offsetCasted =
arith::IndexCastUIOp::create(builder, loc, type, offset);
// Compute the final address: base address + byte offset
auto byteSize = arith::ConstantOp::create(
builder, loc, type,
builder.getIntegerAttr(type,
memrefTy.getElementTypeBitWidth() / 8));
auto byteOffset =
arith::MulIOp::create(builder, loc, offsetCasted, byteSize);
auto addrWithOffset =
arith::AddIOp::create(builder, loc, addrCasted, byteOffset);
return addrWithOffset.getResult();
}
return {};
};
// Materialization to convert ui64 to i64
// Applies only to target materialization.
// Note: i64 to ui64 materialization is not required as xegpu ops
// currently do not produce ui64 as result.
auto ui64ToI64MaterializationCast = [](OpBuilder &builder, Type type,
ValueRange inputs,
Location loc) -> Value {
if (inputs.size() != 1)
return {};
auto input = inputs.front();
if (input.getType() == builder.getIntegerType(64, false)) {
Value cast =
index::CastUOp::create(builder, loc, builder.getIndexType(), input)
.getResult();
return arith::IndexCastUIOp::create(builder, loc, type, cast)
.getResult();
}
return {};
};
// Materialization to convert ui32 to i32
// Applies only to target materialization.
// Note: i32 to ui32 materialization is not required as xegpu ops
// currently do not produce ui32 as result.
auto ui32ToI32MaterializationCast = [](OpBuilder &builder, Type type,
ValueRange inputs,
Location loc) -> Value {
if (inputs.size() != 1)
return {};
auto input = inputs.front();
if (input.getType() == builder.getIntegerType(32, false)) {
Value cast =
index::CastUOp::create(builder, loc, builder.getIndexType(), input)
.getResult();
return arith::IndexCastUIOp::create(builder, loc, type, cast)
.getResult();
}
return {};
};
// Materialization to convert between vector types
// - Add shape cast for different shapes
// - Add bitcast for different element types
// Applies to both source and target materialization.
auto vectorToVectorMaterializationCast = [](OpBuilder &builder, Type type,
ValueRange inputs,
Location loc) -> Value {
if (inputs.size() != 1)
return {};
auto input = inputs.front();
if (auto vecTy = dyn_cast<VectorType>(input.getType())) {
if (auto targetVecTy = dyn_cast<VectorType>(type)) {
Value cast = input;
// If the target type has a different shape, add a shape cast
// If the target type has a different element type, add a bitcast
if (targetVecTy.getShape() != vecTy.getShape()) {
cast = vector::ShapeCastOp::create(
builder, loc,
VectorType::get(targetVecTy.getShape(),
vecTy.getElementType()),
cast)
.getResult();
}
if (targetVecTy.getElementType() != vecTy.getElementType()) {
cast = vector::BitCastOp::create(builder, loc, targetVecTy, cast)
.getResult();
}
return cast;
}
}
return {};
};
// Materialization to convert
// - single element vector to single element of vector element type
// Applies only to target materialization.
auto vectorToSingleElementMaterializationCast =
[](OpBuilder &builder, Type type, ValueRange inputs,
Location loc) -> Value {
if (inputs.size() != 1)
return {};
auto input = inputs.front();
if (auto vecTy = dyn_cast<VectorType>(input.getType())) {
// Source needs to be single element vector
auto rank = vecTy.getRank();
if (rank != 0 && vecTy.getNumElements() != 1)
return {};
auto inElemTy = vecTy.getElementType();
// extract scalar
Value cast = input;
if (rank == 0) {
cast = vector::ExtractOp::create(builder, loc, cast, {}).getResult();
} else {
cast = vector::ExtractOp::create(builder, loc, cast,
SmallVector<int64_t>(rank, 0))
.getResult();
}
// Extracted element type may need conversion
// Two cases
// 1. Index type to integer type
// 2. Other element type mismatch
if (inElemTy.isIndex()) {
cast = arith::IndexCastUIOp::create(builder, loc, type, cast)
.getResult();
} else if (inElemTy != type) {
cast = arith::BitcastOp::create(builder, loc, type, cast).getResult();
}
return cast;
}
return {};
};
// Materialization to convert
// - single element of vector element type to single element vector
// If result type of original op is single element vector and lowered type
// is scalar. This materialization cast creates a single element vector by
// First convert element type if needed and then broadcast to single
// element vector.
// Applies only to source materialization.
auto singleElementToVectorMaterializationCast =
[](OpBuilder &builder, Type type, ValueRange inputs,
Location loc) -> Value {
if (inputs.size() != 1)
return {};
auto input = inputs.front();
auto inTy = input.getType();
if (!inTy.isIntOrFloat())
return {};
// If the target type is a vector of rank 0 or single element vector
// of element type matching input type, broadcast input to target type.
if (auto vecTy = dyn_cast<VectorType>(type)) {
if (vecTy.getRank() != 0 && vecTy.getNumElements() != 1)
return {};
auto outElemTy = vecTy.getElementType();
Value cast = input;
if (outElemTy.isIndex()) {
cast = arith::IndexCastUIOp::create(builder, loc,
builder.getIndexType(), cast)
.getResult();
} else if (inTy != outElemTy) {
cast = arith::BitcastOp::create(builder, loc, outElemTy, cast)
.getResult();
}
return vector::BroadcastOp::create(builder, loc, vecTy, cast)
.getResult();
}
return {};
};
typeConverter.addSourceMaterialization(
singleElementToVectorMaterializationCast);
typeConverter.addSourceMaterialization(vectorToVectorMaterializationCast);
typeConverter.addTargetMaterialization(memrefToIntMaterializationCast);
typeConverter.addTargetMaterialization(ui32ToI32MaterializationCast);
typeConverter.addTargetMaterialization(ui64ToI64MaterializationCast);
typeConverter.addTargetMaterialization(
vectorToSingleElementMaterializationCast);
typeConverter.addTargetMaterialization(vectorToVectorMaterializationCast);
ConversionTarget target(*context);
target.addLegalDialect<xevm::XeVMDialect, LLVM::LLVMDialect,
vector::VectorDialect, arith::ArithDialect,
memref::MemRefDialect, gpu::GPUDialect,
index::IndexDialect>();
target.addIllegalDialect<xegpu::XeGPUDialect>();
// arith.extf/arith.truncf between MX narrow floats and f16/bf16 are routed
// to xevm.extf/xevm.truncf; all other arith float casts stay legal.
target.addDynamicallyLegalOp<arith::ExtFOp>(
[](arith::ExtFOp op) { return !isXeVMExtf(op); });
target.addDynamicallyLegalOp<arith::TruncFOp>(
[](arith::TruncFOp op) { return !isXeVMTruncf(op); });
RewritePatternSet patterns(context);
populateXeGPUToXeVMConversionPatterns(typeConverter, patterns);
scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter,
patterns, target);
if (failed(applyPartialConversion(getOperation(), target,
std::move(patterns))))
signalPassFailure();
}
};
} // namespace
//===----------------------------------------------------------------------===//
// Pattern Population
//===----------------------------------------------------------------------===//
void mlir::populateXeGPUToXeVMConversionPatterns(
const LLVMTypeConverter &typeConverter, RewritePatternSet &patterns) {
patterns.add<CreateNdDescToXeVMPattern,
LoadStorePrefetchNdToXeVMPattern<xegpu::LoadNdOp>,
LoadStorePrefetchNdToXeVMPattern<xegpu::StoreNdOp>,
LoadStorePrefetchNdToXeVMPattern<xegpu::PrefetchNdOp>>(
typeConverter, patterns.getContext());
patterns.add<AtomicRMWToXeVMPattern, PrefetchToXeVMPattern,
LoadStoreToXeVMPattern<xegpu::LoadGatherOp>,
LoadStoreToXeVMPattern<xegpu::StoreScatterOp>>(
typeConverter, patterns.getContext());
patterns.add<LoadStoreMatrixToXeVMPattern<xegpu::LoadMatrixOp>,
LoadStoreMatrixToXeVMPattern<xegpu::StoreMatrixOp>,
CreateMemDescOpPattern>(typeConverter, patterns.getContext());
patterns.add<FenceToXeVMPattern, DpasToXeVMPattern>(typeConverter,
patterns.getContext());
patterns.add<DpasMxToXeVMPattern>(typeConverter, patterns.getContext());
patterns.add<ExtfToXeVMPattern, TruncfToXeVMPattern>(typeConverter,
patterns.getContext());
patterns.add<LaneShuffleToXeVMPattern>(typeConverter, patterns.getContext());
}