Skip to content

[MLIR][XeGPU][VectorToXeGPU] Add lowering from vector.gather/scatter to xegpu.load/store #158024

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Sep 19, 2025
Merged

[MLIR][XeGPU][VectorToXeGPU] Add lowering from vector.gather/scatter to xegpu.load/store #158024

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d97a3c2
[MLIR][XeGPU] Add lowering from vector.gather/scatter to xegpu.load/s…
dchigarev Sep 10, 2025
4d2c284
Add alignment handling
dchigarev Sep 11, 2025
62c5c38
Align lowering with new utils behavior
dchigarev Sep 12, 2025
78b057d
Handle non-unit inner stride
dchigarev Sep 16, 2025
4ce2415
Correct memref.subview test
dchigarev Sep 16, 2025
eb5fbc7
Remove non-necessary comment
dchigarev Sep 16, 2025
54a8299
Remove 'gatherScatterPreconditions' function
dchigarev Sep 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 146 additions & 11 deletions mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,17 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
return success();
}

// Common preconditions for the lowering of vector.gather and vector.scatter:
// 1. Source is a memref.
static LogicalResult gatherScatterPreconditions(PatternRewriter &rewriter,
Operation *op, Type baseType) {
auto srcTy = dyn_cast<MemRefType>(baseType);
if (!srcTy)
return rewriter.notifyMatchFailure(op, "Expects memref source");

return success();
}

static xegpu::CreateNdDescOp
createNdDescriptor(PatternRewriter &rewriter, Location loc,
xegpu::TensorDescType descType, TypedValue<MemRefType> src,
Expand Down Expand Up @@ -183,11 +194,15 @@ static void adjustStridesForPermutation(AffineMap permMap,
// Computes memory strides and a memref offset for vector transfer operations,
// handling both static and dynamic memrefs while applying permutation
// transformations for XeGPU lowering.
template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
std::decay_t<OpType>, vector::TransferReadOp, vector::TransferWriteOp,
vector::GatherOp, vector::ScatterOp>::value>>
static std::pair<SmallVector<Value>, Value>
computeMemrefMeta(VectorTransferOpInterface xferOp, PatternRewriter &rewriter) {
computeMemrefMeta(OpType xferOp, PatternRewriter &rewriter) {
SmallVector<Value> strides;
Value baseMemref = xferOp.getBase();
AffineMap permMap = xferOp.getPermutationMap();
MemRefType memrefType = dyn_cast<MemRefType>(baseMemref.getType());

Location loc = xferOp.getLoc();
Expand All @@ -197,9 +212,14 @@ computeMemrefMeta(VectorTransferOpInterface xferOp, PatternRewriter &rewriter) {
SmallVector<int64_t> intStrides;
if (failed(memrefType.getStridesAndOffset(intStrides, offset)))
return {{}, offsetVal};
// Wrap static strides as MLIR values
for (int64_t s : intStrides)
strides.push_back(arith::ConstantIndexOp::create(rewriter, loc, s));
bool hasDynamicStrides = llvm::any_of(intStrides, [](int64_t strideVal) {
return ShapedType::isDynamic(strideVal);
});

if (!hasDynamicStrides)
for (int64_t s : intStrides)
strides.push_back(arith::ConstantIndexOp::create(rewriter, loc, s));

if (!ShapedType::isDynamic(offset))
offsetVal = arith::ConstantIndexOp::create(rewriter, loc, offset);
}
Expand Down Expand Up @@ -232,8 +252,14 @@ computeMemrefMeta(VectorTransferOpInterface xferOp, PatternRewriter &rewriter) {
if (!offsetVal)
offsetVal = meta.getOffset();
}
// Adjust strides according to the permutation map (e.g., for transpose)
adjustStridesForPermutation(permMap, strides);

if constexpr (llvm::is_one_of<std::decay_t<OpType>, vector::TransferReadOp,
vector::TransferWriteOp>::value) {
AffineMap permMap = xferOp.getPermutationMap();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The permutation map could exist for the gather/store's memref, and I don't understand why we should treat them differently here.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

here we're accessing a permutation map of the operation (transfer_read/write). vector.gather/scatter don't have a permutation map, that's why we're skipping it in that case

If a memref has its own permutation this should be handled automatically by memref.extract_strided_metadata

// Adjust strides according to the permutation map (e.g., for transpose)
adjustStridesForPermutation(permMap, strides);
}

return {strides, offsetVal};
}

Expand Down Expand Up @@ -339,9 +365,51 @@ static Value computeOffsets(VectorTransferOpInterface xferOp,
return localOffsets;
}

// Compute the element-wise offsets for vector.gather or vector.scatter ops.
//
// This function linearizes the base offsets of the gather/scatter operation
// and combines them with the per-element indices to produce a final vector of
// memory offsets.
template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
std::decay_t<OpType>, vector::GatherOp, vector::ScatterOp>::value>>
static Value computeOffsets(PatternRewriter &rewriter, OpType gatScatOp,
ArrayRef<Value> strides, Value baseOffset) {
Location loc = gatScatOp.getLoc();
SmallVector<Value> offsets = gatScatOp.getOffsets();
for (size_t i = 0; i < offsets.size(); ++i) {
Value offsetContrib =
arith::MulIOp::create(rewriter, loc, offsets[i], strides[i]);
baseOffset =
arith::AddIOp::create(rewriter, loc, baseOffset, offsetContrib);
}
Value indices = gatScatOp.getIndices();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think these indices need to multiple with stride of innermost dim, if we allow the non-unit innermost dim stride. I believe this is the only change we need to support it.

VectorType vecType = cast<VectorType>(indices.getType());

Value strideVector =
vector::BroadcastOp::create(rewriter, loc, vecType, strides.back())
.getResult();
Value stridedIndices =
arith::MulIOp::create(rewriter, loc, strideVector, indices).getResult();

Value baseVector =
vector::BroadcastOp::create(
rewriter, loc,
VectorType::get(vecType.getShape(), rewriter.getIndexType()),
baseOffset)
.getResult();
return arith::AddIOp::create(rewriter, loc, baseVector, stridedIndices)
.getResult();
}

template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
std::decay_t<OpType>, vector::TransferReadOp, vector::TransferWriteOp,
vector::GatherOp, vector::ScatterOp>::value>>
// Convert memref to i64 base pointer
static Value memrefToIndexPtr(VectorTransferOpInterface xferOp,
PatternRewriter &rewriter) {
static Value memrefToIndexPtr(OpType xferOp, PatternRewriter &rewriter) {
Location loc = xferOp.getLoc();
auto indexPtr = memref::ExtractAlignedPointerAsIndexOp::create(
rewriter, loc, xferOp.getBase())
Expand Down Expand Up @@ -539,6 +607,71 @@ struct TransferWriteLowering
}
};

struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
using OpRewritePattern<vector::GatherOp>::OpRewritePattern;

LogicalResult matchAndRewrite(vector::GatherOp gatherOp,
PatternRewriter &rewriter) const override {
if (failed(gatherScatterPreconditions(rewriter, gatherOp,
gatherOp.getBase().getType())))
return failure();

Location loc = gatherOp.getLoc();
VectorType vectorType = gatherOp.getVectorType();

auto meta = computeMemrefMeta(gatherOp, rewriter);
if (meta.first.empty())
return rewriter.notifyMatchFailure(gatherOp, "Failed to compute strides");

Value localOffsets =
computeOffsets(rewriter, gatherOp, meta.first, meta.second);
Value flatMemref = memrefToIndexPtr(gatherOp, rewriter);

auto xeGatherOp = xegpu::LoadGatherOp::create(
rewriter, loc, vectorType, flatMemref, localOffsets, gatherOp.getMask(),
/*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{});

auto selectOp =
arith::SelectOp::create(rewriter, loc, gatherOp.getMask(),
xeGatherOp.getResult(), gatherOp.getPassThru());
rewriter.replaceOp(gatherOp, selectOp.getResult());
return success();
}
};

struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
using OpRewritePattern<vector::ScatterOp>::OpRewritePattern;

LogicalResult matchAndRewrite(vector::ScatterOp scatterOp,
PatternRewriter &rewriter) const override {
if (failed(gatherScatterPreconditions(rewriter, scatterOp,
scatterOp.getBase().getType())))
return failure();

Location loc = scatterOp.getLoc();
auto meta = computeMemrefMeta(scatterOp, rewriter);
if (meta.first.empty())
return rewriter.notifyMatchFailure(scatterOp,
"Failed to compute strides");

Value localOffsets =
computeOffsets(rewriter, scatterOp, meta.first, meta.second);
Value flatMemref = memrefToIndexPtr(scatterOp, rewriter);

xegpu::StoreScatterOp::create(rewriter, loc, scatterOp.getValueToStore(),
flatMemref, localOffsets, scatterOp.getMask(),
/*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{});
rewriter.eraseOp(scatterOp);
return success();
}
};

struct LoadLowering : public OpRewritePattern<vector::LoadOp> {
using OpRewritePattern<vector::LoadOp>::OpRewritePattern;

Expand Down Expand Up @@ -654,6 +787,8 @@ struct ConvertVectorToXeGPUPass

void mlir::populateVectorToXeGPUConversionPatterns(
RewritePatternSet &patterns) {
patterns.add<TransferReadLowering, TransferWriteLowering, LoadLowering,
StoreLowering, ContractionLowering>(patterns.getContext());
patterns
.add<TransferReadLowering, TransferWriteLowering, LoadLowering,
ScatterLowering, GatherLowering, StoreLowering, ContractionLowering>(
patterns.getContext());
}
Loading