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
+592 −11
Merged

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

Show file tree
Hide file tree
Changes from 3 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
146 changes: 138 additions & 8 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,24 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
return success();
}

// Common preconditions for the lowering of vector.gather and vector.scatter:
// 1. Source is a memref.
// 2. The innermost dimension of the memref is contiguous (stride == 1)
Copy link
Contributor

@Jianhui-Li Jianhui-Li Sep 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

what is the reason the memref must have stride==1? The HW should support non-unit-stride memref since the offset is per lane.

Copy link
Contributor

Choose a reason for hiding this comment

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

From vector.gather definition.
result[i,j] := if mask[i,j] then base[i0, i1, i2 + indices[i,j]]
else pass_thru[i,j]

My understading is that we first compute the base_offset of base[i0, i1, i2 ], and then compute the offset described by indices by computeing base[0, 0, indices[i, j]], and combine them to get the memory address. The computation uses the strides[] and the strides could be permutated or equal to 1.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

what is the reason the memref must have stride==1? The HW should support non-unit-stride memref since the offset is per lane.

my bad, I genuinely thought that vector.gather/scatter only support inner strides == 1. Added support for non-unit inner strides

static LogicalResult gatherScatterPreconditions(PatternRewriter &rewriter,
Operation *op, Type baseType) {
auto srcTy = dyn_cast<MemRefType>(baseType);
if (!srcTy)
return rewriter.notifyMatchFailure(op, "Expects memref source");

SmallVector<int64_t> strides;
int64_t offset;
if (failed(srcTy.getStridesAndOffset(strides, offset)) || strides.back() != 1)
return rewriter.notifyMatchFailure(
op, "Buffer must be contiguous in the innermost dimension");

return success();
}

static xegpu::CreateNdDescOp
createNdDescriptor(PatternRewriter &rewriter, Location loc,
xegpu::TensorDescType descType, TypedValue<MemRefType> src,
Expand Down Expand Up @@ -183,11 +201,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 Down Expand Up @@ -232,8 +254,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 +367,44 @@ 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 baseVector =
vector::BroadcastOp::create(
rewriter, loc,
VectorType::get(vecType.getShape(), rewriter.getIndexType()),
baseOffset)
.getResult();
return arith::AddIOp::create(rewriter, loc, baseVector, indices).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 +602,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 +782,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());
}
187 changes: 187 additions & 0 deletions mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,187 @@
// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | FileCheck %s

gpu.module @xevm_module {
gpu.func @load_1D_vector(%source: memref<8x16x32xf32>,
%off1: index, %off2: index, %off3: index,
%indices: vector<8xindex>, %mask: vector<8xi1>,
%pass_thru: vector<8xf32>) -> vector<8xf32> {
%0 = vector.gather %source[%off1, %off2, %off3][%indices], %mask,
%pass_thru : memref<8x16x32xf32>, vector<8xindex>, vector<8xi1>, vector<8xf32> into vector<8xf32>
gpu.return %0 : vector<8xf32>
}
// CHECK-LABEL: @load_1D_vector(
// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>
// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8xf32>) -> vector<8xf32> {
// CHECK-COUNT2: arith.muli {{.*}} : index
// CHECK-COUNT2: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8xi1>, vector<8xf32>
// CHECK: gpu.return %[[RES]] : vector<8xf32>
}

// -----
gpu.module @xevm_module {
gpu.func @load_2D_memref(%source: memref<8x32xf32>,
%off1: index, %off2: index,
%indices: vector<8xindex>, %mask: vector<8xi1>,
%pass_thru: vector<8xf32>) -> vector<8xf32> {
%0 = vector.gather %source[%off1, %off2][%indices], %mask,
%pass_thru : memref<8x32xf32>, vector<8xindex>, vector<8xi1>, vector<8xf32> into vector<8xf32>
gpu.return %0 : vector<8xf32>
}
// CHECK-LABEL: @load_2D_memref(
// CHECK-SAME: %[[SRC:.+]]: memref<8x32xf32>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>
// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8xf32>) -> vector<8xf32> {
// CHECK-COUNT1: arith.muli {{.*}} : index
// CHECK-COUNT1: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x32xf32> -> index
// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8xi1>, vector<8xf32>
// CHECK: gpu.return %[[RES]] : vector<8xf32>
}

// -----
gpu.module @xevm_module {
gpu.func @load_2D_vector(%source: memref<8x16x32xf32>,
%off0: index, %off1: index, %off2: index,
%indices: vector<8x16xindex>, %mask: vector<8x16xi1>,
%pass_thru: vector<8x16xf32>) -> vector<8x16xf32> {
%0 = vector.gather %source[%off0, %off1, %off2][%indices], %mask,
%pass_thru : memref<8x16x32xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32> into vector<8x16xf32>
gpu.return %0 : vector<8x16xf32>
}
// CHECK-LABEL: @load_2D_vector(
// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>
// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {
// CHECK-COUNT2: arith.muli {{.*}} : index
// CHECK-COUNT2: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8x16xi1>, vector<8x16xf32>
// CHECK: gpu.return %[[RES]] : vector<8x16xf32>
}

// -----
gpu.module @xevm_module {
gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,
%off0: index, %off1: index, %off2: index,
%indices: vector<8x16xindex>, %mask: vector<8x16xi1>,
%pass_thru: vector<8x16xf32>) -> vector<8x16xf32> {
%0 = vector.gather %source[%off0, %off1, %off2][%indices], %mask,
%pass_thru : memref<?x?x?xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32> into vector<8x16xf32>
gpu.return %0 : vector<8x16xf32>
}
// CHECK-LABEL: @load_dynamic_source(
// CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>
// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {
// CHECK: memref.extract_strided_metadata %[[SRC]]
// CHECK-COUNT2: arith.muli {{.*}} : index
// CHECK-COUNT2: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?xf32> -> index
// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8x16xi1>, vector<8x16xf32>
// CHECK: gpu.return %[[RES]] : vector<8x16xf32>
}

// -----
gpu.module @xevm_module {
gpu.func @load_dynamic_source2(%source: memref<?x8x16xf32>,
%off0: index, %off1: index, %off2: index,
%indices: vector<8x16xindex>, %mask: vector<8x16xi1>,
%pass_thru: vector<8x16xf32>) -> vector<8x16xf32> {
%0 = vector.gather %source[%off0, %off1, %off2][%indices], %mask,
%pass_thru : memref<?x8x16xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32> into vector<8x16xf32>
gpu.return %0 : vector<8x16xf32>
}
// CHECK-LABEL: @load_dynamic_source2(
// CHECK-SAME: %[[SRC:.+]]: memref<?x8x16xf32>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>
// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {
// CHECK-NOT: memref.extract_strided_metadata %[[SRC]]
// CHECK-COUNT2: arith.muli {{.*}} : index
// CHECK-COUNT2: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x8x16xf32> -> index
// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8x16xi1>, vector<8x16xf32>
// CHECK: gpu.return %[[RES]] : vector<8x16xf32>
}

// -----
gpu.module @xevm_module {
gpu.func @no_load_tensor(%source: tensor<32x64xf32>,
%off: index, %indices: vector<8x16xindex>,
%mask: vector<8x16xi1>, %pass_thru: vector<8x16xf32>) -> vector<8x16xf32> {
%0 = vector.gather %source[%off, %off][%indices], %mask,
%pass_thru : tensor<32x64xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32> into vector<8x16xf32>
gpu.return %0 : vector<8x16xf32>
}
// CHECK-LABEL: @no_load_tensor(
// CHECK: vector.gather
}

// -----
gpu.module @xevm_module {
gpu.func @gather_from_subview(%source: memref<4096x4096xf16>,
%off1: index, %off2: index,
%indices: vector<8xindex>,
%mask: vector<8xi1>,
%pass_thru: vector<8xf16>) -> vector<8xf16> {
%subview = memref.subview %source[%off1, %off2] [256, 256] [1, 1]
: memref<4096x4096xf16>
to memref<256x256xf16, strided<[4096, 1], offset: ?>>
%0 = vector.gather %subview[%off1, %off2][%indices], %mask, %pass_thru
Copy link
Contributor

Choose a reason for hiding this comment

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

Can we use a different value than off1 and off2? The off1 and off2 is suppose to be multiple of 256, so the %subview[%off1, %off2] would be out of boundary? It also makes hard to read the check code sequence, which I believe doesn't contain offset computation for vector.gather using off1/off2
.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Fixed. Now subview and the gather op use different values

: memref<256x256xf16, strided<[4096, 1], offset: ?>>,
vector<8xindex>, vector<8xi1>, vector<8xf16>
into vector<8xf16>
gpu.return %0 : vector<8xf16>
}
// CHECK-LABEL: @gather_from_subview(
// CHECK-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>,
// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>,
// CHECK-SAME: %[[PASS:.+]]: vector<8xf16>) -> vector<8xf16> {
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1]
// CHECK: %[[BB:.+]], %[[OFFSET:.+]],{{.*}},{{.*}} = memref.extract_strided_metadata %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index
// CHECK: arith.muli {{.*}} : index
// CHECK: arith.addi %[[OFFSET]]{{.*}} : index
// CHECK: %[[BASE_OFF:.+]] = arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast %[[BASE_OFF]] : index to vector<8xindex>
// CHECK: %[[LIN:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>
// CHECK: %[[BASE_IDX:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE_IDX]] : index to i64
// CHECK: %[[VEC:.+]] = xegpu.load %[[BASE_I64]]{{\[}}%[[LIN]]{{\]}}, %[[MASK]]
// CHECK-SAME: : i64, vector<8xindex>, vector<8xi1> -> vector<8xf16>
// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS]] : vector<8xi1>, vector<8xf16>
// CHECK: gpu.return %[[RES]] : vector<8xf16>
}
Loading