blob: 0457f8128b90880eff1b7ce5d6a4fd299d3717d7 [file] [log] [blame]
//===- XeGPUUnroll.cpp - patterns to do unrolling ---------------*- 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 contains patterns for unrolling XeGPU operations. It follows a
// similar concept and design as vector unroll patterns, serving as a complement
// to them.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Debug.h"
#include <numeric>
namespace mlir {
namespace xegpu {
#define GEN_PASS_DEF_XEGPUUNROLL
#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
} // namespace xegpu
} // namespace mlir
#define DEBUG_TYPE "xegpu-unroll"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
using namespace mlir;
namespace {
template <typename SourceOp>
struct UnrollPattern : public OpRewritePattern<SourceOp> {
UnrollPattern(MLIRContext *context, const xegpu::UnrollOptions &options,
PatternBenefit benefit = 1)
: OpRewritePattern<SourceOp>(context, benefit), options(options) {}
protected:
/// Return the target shape for the given `op`. Return std::nullopt if the
/// op shouldn't be or cannot be unrolled.
std::optional<SmallVector<int64_t>> getTargetShape(Operation *op) const {
LDBG("");
LDBG("Get unroll shape for: " << *op);
if (options.filterConstraint && failed(options.filterConstraint(op))) {
LDBG("--no filter constraint -> BAIL");
return std::nullopt;
}
assert(options.nativeShape &&
"expects the native shape for native shape call back function.");
auto nativeShape = options.nativeShape(op);
return nativeShape;
}
SmallVector<Type> getUnrolledTypes(ShapedType type,
ArrayRef<int64_t> tileShape) const {
return options.getUnrolledTypes(type, tileShape);
}
/// Emulate the the unpack behavior using insert_strided_slice for VectorType
/// values and unrealized_conversion_cast for TensorDescType values.
Value unpack(ValueRange srcs, Type destTy, ArrayRef<int64_t> blockSize,
Location loc, PatternRewriter &rewriter) const {
if (auto vecTy = dyn_cast<VectorType>(destTy)) {
assert(vecTy.getRank() == static_cast<int64_t>(blockSize.size()) &&
"Expecting blockSize size to match the rank of destTy.");
auto shape = vecTy.getShape();
return xegpu::createVectorWithShapeFromValues(rewriter, loc, srcs, shape);
}
if (isa<xegpu::TensorDescType>(destTy)) {
auto attr = NamedAttribute(rewriter.getStringAttr(unpackAttrName),
rewriter.getUnitAttr());
auto blkAttr = NamedAttribute(rewriter.getStringAttr(blockAttrName),
rewriter.getDenseI64ArrayAttr(blockSize));
auto castOp = rewriter.create<UnrealizedConversionCastOp>(
loc, destTy, srcs, ArrayRef<NamedAttribute>({attr, blkAttr}));
return castOp.getResult(0);
}
llvm_unreachable("Unexpected destTy.");
return Value();
}
/// Emulate the the pack behavior using extract_strided_slice for VectorType
/// values and unrealized_conversion_cast for TensorDescType values.
SmallVector<Value> pack(Value src, TypeRange destTypes,
ArrayRef<int64_t> blockSize, Location loc,
PatternRewriter &rewriter) const {
if (auto vecTy = dyn_cast<VectorType>(src.getType())) {
assert(vecTy.getRank() == static_cast<int64_t>(blockSize.size()) &&
"Expecting blockSize size to match the rank of src.");
return xegpu::extractVectorsWithShapeFromValue(rewriter, loc, src,
blockSize);
}
if (isa<xegpu::TensorDescType>(src.getType())) {
auto attr = NamedAttribute(rewriter.getStringAttr(packAttrName),
rewriter.getUnitAttr());
auto blkAttr = NamedAttribute(rewriter.getStringAttr(blockAttrName),
rewriter.getDenseI64ArrayAttr(blockSize));
auto castOp = rewriter.create<UnrealizedConversionCastOp>(
loc, destTypes, src, ArrayRef<NamedAttribute>({attr, blkAttr}));
return castOp.getResults();
}
llvm_unreachable("Unexpected src type.");
return SmallVector<Value>();
}
private:
const char *const packAttrName = "__xegpu_blocking_pack__";
const char *const unpackAttrName = "__xegpu_blocking_unpack__";
const char *const blockAttrName = "__xegpu_blocking_tile_shape__";
xegpu::UnrollOptions options;
};
struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
using UnrollPattern<xegpu::CreateNdDescOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::CreateNdDescOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
xegpu::TensorDescType tdescTy = op.getType();
int64_t rank = tdescTy.getRank();
ArrayRef<int64_t> shape = tdescTy.getShape();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
auto newTdescTy = getUnrolledTypes(tdescTy, *targetShape)[0];
auto addi = [&](OpFoldResult a, int64_t b) -> Value {
std::optional<int64_t> maybeInt = getConstantIntValue(a);
if (maybeInt) {
return rewriter.create<arith::ConstantIndexOp>(loc, *maybeInt + b);
} else {
auto aV = llvm::cast<Value>(a);
auto bV = rewriter.create<arith::ConstantIndexOp>(loc, b);
return rewriter.createOrFold<arith::AddIOp>(loc, aV, bV);
}
};
SmallVector<OpFoldResult> mixedOffsets = op.getMixedOffsets();
// For n-D memrefs where n > rank, we need to handle the last `rank`
// dimensions only, and keep the first `n-rank` dimensions as is.
SmallVector<OpFoldResult> oldOffsets = llvm::to_vector(
llvm::drop_begin(mixedOffsets, mixedOffsets.size() - rank));
auto validIdxes =
llvm::seq<int64_t>(mixedOffsets.size() - rank, mixedOffsets.size());
SmallVector<Value> newOps;
for (SmallVector<int64_t> offsets :
StaticTileOffsetRange(shape, *targetShape)) {
for (auto [idx, oldOff, offset] :
llvm::zip(validIdxes, oldOffsets, offsets))
mixedOffsets[idx] = addi(oldOff, offset);
auto newOp = rewriter.create<xegpu::CreateNdDescOp>(
loc, newTdescTy, op.getSource(), mixedOffsets, op.getMixedSizes(),
op.getMixedStrides());
newOps.push_back(newOp);
}
Value castOp = unpack(newOps, tdescTy, *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
}
};
struct UnrollUpdateNdOffsetOp : public UnrollPattern<xegpu::UpdateNdOffsetOp> {
using UnrollPattern<xegpu::UpdateNdOffsetOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::UpdateNdOffsetOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
xegpu::TensorDescType tdescTy = op.getTensorDescType();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdesc = pack(
op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
SmallVector<Value> newOps;
for (auto t : convertedTdesc) {
auto newOp = rewriter.create<xegpu::UpdateNdOffsetOp>(
loc, t.getType(), t, op.getOffsets(), op.getConstOffsets());
newOps.push_back(newOp);
}
Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
}
};
struct UnrollPrefetchNdOp : public UnrollPattern<xegpu::PrefetchNdOp> {
using UnrollPattern<xegpu::PrefetchNdOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::PrefetchNdOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
xegpu::TensorDescType tdescTy = op.getTensorDescType();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdesc = pack(
op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
for (auto t : convertedTdesc)
rewriter.create<xegpu::PrefetchNdOp>(loc, TypeRange(), t, op->getAttrs());
rewriter.eraseOp(op);
return success();
}
};
struct UnrollLoadNdOp : public UnrollPattern<xegpu::LoadNdOp> {
using UnrollPattern<xegpu::LoadNdOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::LoadNdOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
VectorType valueTy = op.getType();
xegpu::TensorDescType tdescTy = op.getTensorDescType();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
Type elemTy = tdescTy.getElementType();
VectorType newValueTy = valueTy.cloneWith(*targetShape, elemTy);
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdescs = pack(
op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
SmallVector<Value> newOps;
for (auto t : convertedTdescs) {
auto newOp =
rewriter.create<xegpu::LoadNdOp>(loc, newValueTy, t, op->getAttrs());
newOps.push_back(newOp);
}
Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
}
};
struct UnrollStoreNdOp : public UnrollPattern<xegpu::StoreNdOp> {
using UnrollPattern<xegpu::StoreNdOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::StoreNdOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
VectorType valueTy = op.getValueType();
xegpu::TensorDescType tdescTy = op.getTensorDescType();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
SmallVector<Type> convertedValTypes =
getUnrolledTypes(valueTy, *targetShape);
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedValues =
pack(op.getValue(), convertedValTypes, *targetShape, loc, rewriter);
SmallVector<Value> convertedTdescs = pack(
op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
for (auto [v, t] : llvm::zip(convertedValues, convertedTdescs))
rewriter.create<xegpu::StoreNdOp>(loc, v, t, op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr());
rewriter.eraseOp(op);
return success();
}
};
struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
using UnrollPattern<xegpu::DpasOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::DpasOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
// expecting every operands is a 2D Vector
if (llvm::any_of(op->getOperandTypes(), [&](Type type) {
auto vecTy = dyn_cast<VectorType>(type);
return !vecTy || vecTy.getRank() != 2;
}))
return failure();
// A vector of 3 elements should be returned, representing M, K, N
// respectively.
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape || targetShape->size() != 3)
return failure();
auto M = (*targetShape)[0];
auto K = (*targetShape)[1];
auto N = (*targetShape)[2];
int64_t aBlockSize[2] = {M, K};
int64_t bBlockSize[2] = {K, N};
int64_t cBlockSize[2] = {M, N};
auto packWrapper = [&](TypedValue<VectorType> val,
ArrayRef<int64_t> blockSize) {
VectorType type = val.getType();
std::optional<SmallVector<int64_t>> grids =
computeShapeRatio(type.getShape(), blockSize);
assert(grids && "Expecting grids to be computed.");
auto numNewOps = computeProduct(*grids);
if (numNewOps == 1)
return SmallVector<Value>({val});
VectorType newVecTy = type.cloneWith(blockSize, type.getElementType());
SmallVector<Type> convertedTypes(numNewOps, newVecTy);
SmallVector<Value> values =
pack(val, convertedTypes, blockSize, loc, rewriter);
return values;
};
auto a = op.getLhs();
auto b = op.getRhs();
auto c = op.getAcc();
auto aShape = a.getType().getShape();
auto bShape = b.getType().getShape();
SmallVector<Value> aVals, bVals, cVals;
aVals = packWrapper(a, aBlockSize);
bVals = packWrapper(b, bBlockSize);
if (c)
cVals = packWrapper(c, cBlockSize);
// Skip the operation if every operand has an invalid blocking size (empty)
// or if the original shape matches the blocking size (size == 1).
auto ranges = c ? SmallVector<ValueRange>({aVals, bVals, cVals})
: SmallVector<ValueRange>({aVals, bVals});
if (llvm::any_of(ranges, [](auto &v) { return v.size() == 0; }) ||
llvm::all_of(ranges, [](auto &v) { return v.size() == 1; }))
return failure();
VectorType resultTy = op.getResult().getType();
auto vecTy = VectorType::get(cBlockSize, resultTy.getElementType());
int64_t mIters = aShape[0] / M;
int64_t kIters = aShape[1] / K;
int64_t nIters = bShape[1] / N;
SmallVector<Value> newOps;
for (int64_t i = 0; i < mIters; ++i) {
for (int64_t j = 0; j < nIters; ++j) {
Value tmpC;
if (c)
tmpC = cVals[i * nIters + j]; // init with acc
for (int64_t k = 0; k < kIters; ++k) {
Value aVec = aVals[i * kIters + k];
Value bVec = bVals[k * nIters + j];
SmallVector<Value> operands({aVec, bVec});
if (tmpC)
operands.push_back(tmpC);
tmpC = rewriter.create<xegpu::DpasOp>(loc, vecTy, operands,
op->getAttrs());
}
newOps.push_back(tmpC);
}
}
Value castOp = unpack(newOps, resultTy, cBlockSize, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
}
};
struct UnrollCreateDescOp : public UnrollPattern<xegpu::CreateDescOp> {
using UnrollPattern<xegpu::CreateDescOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::CreateDescOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
xegpu::TensorDescType tdescTy = op.getType();
TypedValue<::mlir::VectorType> indiceVec = op.getOffsets();
VectorType indiceVecTy = indiceVec.getType();
if (!tdescTy.isScattered())
return failure();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
SmallVector<int64_t> targetIndiceShape(*targetShape);
int64_t originalChunkSize = tdescTy.getChunkSize();
// IndiceVec is 1 dim lower than tdescTy when chunkSize is larger than 1.
if (originalChunkSize > 1)
targetIndiceShape.pop_back();
auto newTdescTy = getUnrolledTypes(tdescTy, *targetShape)[0];
SmallVector<Type> convertedIndiceTypes =
getUnrolledTypes(indiceVecTy, targetIndiceShape);
SmallVector<Value> convertedIndiceVec =
pack(indiceVec, convertedIndiceTypes, targetIndiceShape, loc, rewriter);
SmallVector<Value> newOps;
// More indices is need when chunkSize > 1. Since a big load from one
// address could be break into multiple small loads.
if (originalChunkSize > 1) {
int64_t blockedChunkSize = targetShape->back();
int64_t numNewChunks = originalChunkSize / blockedChunkSize;
for (auto [indice, indiceType] :
llvm::zip(convertedIndiceVec, convertedIndiceTypes)) {
for (int64_t i = 0; i < numNewChunks; ++i) {
// Compute the offset
Value inc = rewriter.create<arith::ConstantIndexOp>(
loc, i * blockedChunkSize);
Value incVec = rewriter.create<vector::SplatOp>(loc, indiceType, inc);
Value offsetIndice =
rewriter.create<arith::AddIOp>(loc, indice, incVec);
auto newOp = rewriter.create<xegpu::CreateDescOp>(
loc, newTdescTy, op.getSource(), offsetIndice);
newOps.push_back(newOp);
}
}
} else {
for (auto indice : convertedIndiceVec) {
auto newOp = rewriter.create<xegpu::CreateDescOp>(
loc, newTdescTy, op.getSource(), indice);
newOps.push_back(newOp);
}
}
Value castOp = unpack(newOps, tdescTy, *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
}
};
struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
using UnrollPattern<xegpu::LoadGatherOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::LoadGatherOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
VectorType valueTy = llvm::dyn_cast<VectorType>(op.getValue().getType());
xegpu::TensorDescType tdescTy = op.getTensorDescType();
if (!tdescTy.isScattered())
return failure();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
SmallVector<int64_t> targetMaskShape(*targetShape);
int64_t originalChunkSize = tdescTy.getChunkSize();
VectorType maskTy = llvm::dyn_cast<VectorType>(op.getMask().getType());
Type elemTy = tdescTy.getElementType();
VectorType newValueTy = valueTy.cloneWith(*targetShape, elemTy);
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdescs = pack(
op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
SmallVector<Type> convertedMaskTypes;
SmallVector<Value> convertedMasks;
if (originalChunkSize > 1) {
targetMaskShape.pop_back();
convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
SmallVector<Value> convertedMasks1D = pack(
op.getMask(), convertedMaskTypes, targetMaskShape, loc, rewriter);
int64_t blockedChunkSize = targetShape->back();
int64_t numNewChunks = originalChunkSize / blockedChunkSize;
for (auto mask : convertedMasks1D) {
for (int64_t i = 0; i < numNewChunks; ++i)
convertedMasks.push_back(mask);
}
// This is to handle the transpose effect when chunkSize > 1.
std::swap((*targetShape)[0], (*targetShape)[1]);
newValueTy = valueTy.cloneWith(*targetShape, elemTy);
} else {
convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
convertedMasks = pack(op.getMask(), convertedMaskTypes, targetMaskShape,
loc, rewriter);
}
SmallVector<Value> newOps;
for (auto [t, m] : llvm::zip(convertedTdescs, convertedMasks)) {
auto newOp = rewriter.create<xegpu::LoadGatherOp>(
loc, newValueTy, t, m, op.getTransposeAttr(), op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr());
newOps.push_back(newOp);
}
Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
}
};
struct UnrollPrefetchOp : public UnrollPattern<xegpu::PrefetchOp> {
using UnrollPattern<xegpu::PrefetchOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::PrefetchOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
xegpu::TensorDescType tdescTy = op.getTensorDescType();
if (!tdescTy.isScattered())
return failure();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdesc = pack(
op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
for (auto t : convertedTdesc)
rewriter.create<xegpu::PrefetchOp>(loc, TypeRange(), t, op->getAttrs());
rewriter.eraseOp(op);
return success();
}
};
struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
using UnrollPattern<xegpu::StoreScatterOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::StoreScatterOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
VectorType valueTy = llvm::dyn_cast<VectorType>(op.getValue().getType());
xegpu::TensorDescType tdescTy = op.getTensorDescType();
if (!tdescTy.isScattered())
return failure();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
SmallVector<int64_t> targetIndiceShape(*targetShape);
int64_t originalChunkSize = tdescTy.getChunkSize();
VectorType maskTy = llvm::dyn_cast<VectorType>(op.getMask().getType());
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdescs = pack(
op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
SmallVector<Type> convertedMaskTypes;
SmallVector<Value> convertedMasks;
if (originalChunkSize > 1) {
int64_t blockedChunkSize = targetShape->back();
int64_t numNewChunks = originalChunkSize / blockedChunkSize;
convertedMaskTypes = getUnrolledTypes(maskTy, (*targetShape)[0]);
SmallVector<Value> convertedMasks1D = pack(
op.getMask(), convertedMaskTypes, (*targetShape)[0], loc, rewriter);
for (auto mask : convertedMasks1D) {
for (int64_t i = 0; i < numNewChunks; ++i) {
convertedMasks.push_back(mask);
}
}
// This is to handle the transpose effect when chunkSize > 1.
std::swap((*targetShape)[0], (*targetShape)[1]);
} else {
convertedMaskTypes = getUnrolledTypes(maskTy, *targetShape);
convertedMasks =
pack(op.getMask(), convertedMaskTypes, *targetShape, loc, rewriter);
}
SmallVector<Type> convertedValTypes =
getUnrolledTypes(valueTy, *targetShape);
SmallVector<Value> convertedValues =
pack(op.getValue(), convertedValTypes, *targetShape, loc, rewriter);
for (size_t i = 0; i < convertedValues.size(); ++i) {
Value v = convertedValues[i];
Value t = convertedTdescs[i];
Value m = op.getMask() ? convertedMasks[i] : nullptr;
rewriter.create<xegpu::StoreScatterOp>(
loc, v, t, m, op.getTransposeAttr(), op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr());
}
rewriter.eraseOp(op);
return success();
}
};
struct UnrollUpdateOffsetOp : public UnrollPattern<xegpu::UpdateOffsetOp> {
using UnrollPattern<xegpu::UpdateOffsetOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::UpdateOffsetOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
xegpu::TensorDescType tdescTy = op.getTensorDescType();
if (tdescTy.getRank() > 2)
return failure();
if (!tdescTy.isScattered())
return failure();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdesc = pack(
op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
TypedValue<::mlir::VectorType> offsetVec = op.getOffsets();
VectorType offsetVecTy = offsetVec.getType();
SmallVector<Type> convertedOffsetTypes;
SmallVector<Value> convertedOffsetVec;
SmallVector<Value> newOps;
int64_t originalChunkSize = tdescTy.getChunkSize();
if (originalChunkSize > 1) {
SmallVector<int64_t> shape1D(targetShape->begin(),
targetShape->end() - 1);
convertedOffsetTypes = getUnrolledTypes(offsetVecTy, shape1D);
SmallVector<Value> convertedOffsetVec1D =
pack(offsetVec, convertedOffsetTypes, shape1D, loc, rewriter);
int64_t blockedChunkSize = targetShape->back();
int64_t numNewChunks = originalChunkSize / blockedChunkSize;
for (auto offset : convertedOffsetVec1D) {
for (int64_t i = 0; i < numNewChunks; ++i) {
convertedOffsetVec.push_back(offset);
}
}
} else {
convertedOffsetTypes = getUnrolledTypes(offsetVecTy, *targetShape);
convertedOffsetVec =
pack(offsetVec, convertedOffsetTypes, *targetShape, loc, rewriter);
}
for (auto [t, o] : llvm::zip(convertedTdesc, convertedOffsetVec)) {
auto newOp =
rewriter.create<xegpu::UpdateOffsetOp>(loc, t.getType(), t, o);
newOps.push_back(newOp);
}
Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
}
};
} // namespace
void mlir::xegpu::populateXeGPUUnrollPatterns(
RewritePatternSet &patterns, const xegpu::UnrollOptions &options) {
patterns.add<UnrollCreateNdOp, UnrollUpdateNdOffsetOp, UnrollPrefetchNdOp,
UnrollLoadNdOp, UnrollStoreNdOp, UnrollDpasOp,
UnrollCreateDescOp, UnrollLoadGatherOp, UnrollStoreScatterOp,
UnrollPrefetchOp, UnrollUpdateOffsetOp>(patterns.getContext(),
options);
}