[Mlir-commits] [mlir] [mlir][XeGPU] Distribute coalesced gather/scatter via lane_data (PR #218834)
Md Abdullah Shahneous Bari
llvmlistbot at llvm.org
Tue Aug 25 21:30:42 PDT 2026
https://github.com/mshahneo created https://github.com/llvm/llvm-project/pull/218834
Reworks `SgToLaneLoadGather` / `SgToLaneStoreScatter` so a **coalesced**
gather/scatter lowers to a chunked (contiguous) memory access driven **purely by
the layout's `lane_data`** — with no `chunk_size` attribute anywhere.
Stacked on top of #205122 (`chunk_size` removal); it should be reviewed/merged
after that PR.
### Motivation
Once `chunk_size` is removed (#205122), the SIMT distribution has no attribute
telling it that a lane owns a contiguous run. That information already lives in
the layout: `lane_data[FCD]`. This PR makes distribution read it directly, so
the coalesced form is expressed entirely through the layout and the value type.
### What it does
When `lane_data[FCD] = D > 1` covers the lane's **entire** per-lane fragment —
one round, i.e. `lane_layout[FCD] * D == FCD extent`, so the distributed
per-lane vector has exactly `D` elements — the lane owns a single contiguous run
`{base, base+1, …, base+D-1}`. In that case distribution:
- extracts the **first** offset and mask of the group (scalar `base` + scalar
mask), and
- emits a `vector<D>` value.
The chunk size is implied by the value type, which the XeVM lowering already
consumes; the now-unused per-lane offset computations are DCE'd during lowering.
The **round-robin** case (`lane_data[FCD] == 1`, multiple strided rounds — e.g.
a reduction source where lane `l` owns `{l, l+SG, l+2·SG, …}`) is strided, not
contiguous, and must not be coalesced. The `innerLaneData > 1` guard excludes it.
>From 266f39cdc814f4dc37d852ccb11299fb290ae567 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Tue, 28 Jul 2026 17:18:09 +0000
Subject: [PATCH 1/2] [mlir][xegpu] Remove the chunk_size attribute from
load/store gather/scatter
The chunk_size attribute on xegpu.load (LoadGatherOp) and xegpu.store
(StoreScatterOp) is removed entirely, along with every code path that
special-cased chunk_size > 1. Without chunk_size, the mask and offsets
always describe the same lane grid as the value, so:
- The verifier (isValidGatherScatterBufferParams) validates purely off
types: offsets and mask must both be scalar or both be vectors of the
same shape; a scalar/size-1 mask means one lane loads/stores the whole
value as a contiguous block (any value shape); otherwise the value
shape must equal the mask shape.
- XeGPUUnroll drops the chunked-unroll branch (the whole reason for the
removal) and keeps the plain congruent path.
- XeGPUUtils no longer drops the innermost dim from the mask/offset
layout; they share the value's layout.
- PropagateLayout / XeGPULayoutImpl drop the now-dead contigChunkSize
parameter (it was always fed by getChunkSize().value_or(1)) and the
inferMaskOffsetLayoutForScatterIO helper; gather/scatter is one element
per lane (maxChunkSize == 1).
- SgToLaneDistribute always distributes the 1-D gather/scatter form.
The separate `contiguity` attribute and its analysis are untouched.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 42 +++----
.../XeGPU/Transforms/XeGPULayoutImpl.h | 14 +--
.../VectorToXeGPU/VectorToXeGPU.cpp | 4 -
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 93 ++++++--------
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 54 ++++----
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 18 +--
.../Transforms/XeGPUSgToLaneDistribute.cpp | 40 ++----
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 116 +++---------------
.../Transforms/XeGPUWgToSgDistribute.cpp | 16 +--
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 17 +--
mlir/test/Dialect/XeGPU/invalid.mlir | 59 +++------
mlir/test/Dialect/XeGPU/ops.mlir | 37 ++----
.../XeGPU/propagate-layout-subgroup.mlir | 4 +-
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 4 +-
.../XeGPU/sg-to-lane-distribute-unit.mlir | 36 ------
.../Dialect/XeGPU/sg-to-lane-distribute.mlir | 61 ---------
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 82 +++----------
.../Dialect/XeGPU/xegpu-unroll-patterns.mlir | 56 +--------
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 42 +++----
19 files changed, 184 insertions(+), 611 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index cb39e25142c46..ef79132925d26 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -708,15 +708,12 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
let description = [{ It (aka. load) load data per each lane. The output
describes the data being loaded at the subgroup level, so its size is
- consistent with the number of lanes in a subgroup. When the chunk size
- is larger than 2, the output vector is a 2D vector, with dim-0 correspoding
- to lanes, and dim-1 corresponding to the chunk size loaded by each lane.
+ consistent with the number of lanes in a subgroup. At lane level each lane
+ may load a 1D vector of contiguous elements; that per-lane width is implied
+ by the result type and does not need a separate attribute.
The mask operand masks out memory access so that it is safe to pass out-of-boundary
addresses/offsets as long as they are masked. Each mask element applies to one lane.
- In lane level, the result is a 1D vector that represents the data to be loaded by
- each lane. If size is not 1, size should be equal to the chunk size.
-
This operation serves as an anchor through which users assign a layout attribute
to govern computation distribution.
@@ -733,8 +730,6 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
mask is a vector of size equal to the subgroup size, or 1 at lane level.
scalar mask is also valid for lane level.
- - `chunk_size`: [optional] represents contiguous number of elements to load from per work item.
-
- `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache.
- `layout`: [optional] Describes the expected layout of the `tensor_desc` operand or the result
@@ -758,8 +753,9 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
```
Example 2 (lane level):
- lane level only accepts the offsets variant. chunk_size can be inferred from result
- type. In this example, chunk_size is 8.
+ lane level only accepts the offsets variant. Each lane loads a contiguous 1D
+ vector whose length is given by the result type; here each lane loads 8
+ elements.
```mlir
%2 = xegpu.load %1[%2], %0 <{l1_hint = #xegpu.cache_hint<cached>,
l2_hint = #xegpu.cache_hint<uncached>,
@@ -771,7 +767,7 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
let arguments = (ins XeGPU_GatherScatterBaseAddrType:$source,
AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
- AnyTypeOf<[XeGPU_MaskType, I1]>:$mask, OptionalAttr<I64Attr>:$chunk_size,
+ AnyTypeOf<[XeGPU_MaskType, I1]>:$mask,
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
@@ -817,13 +813,11 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
let builders = [
OpBuilder<(ins "Type": $value, "Value": $source,
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
- "IntegerAttr": $chunk_size,
"xegpu::CachePolicyAttr": $l1_hint,
"xegpu::CachePolicyAttr": $l2_hint,
"xegpu::CachePolicyAttr": $l3_hint)>,
OpBuilder<(ins "Type": $value, "Value": $source,
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
- "IntegerAttr": $chunk_size,
"xegpu::CachePolicyAttr": $l1_hint,
"xegpu::CachePolicyAttr": $l2_hint,
"xegpu::CachePolicyAttr": $l3_hint,
@@ -837,14 +831,9 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
let summary = "store data to scattered memory locations.";
let description =
[{ It (aka. store) stores data to scattered memory locations. The value is
- typically a 1D vector. But when the chunk size of the TensorDesc is larger than 1, it will be
- a 2D vector instead. For the later case, dim-1 of the value correspods to the simd lanes
- and the dim-0 of the value corresponds to the chunk size stored per lane. So `store_scatter`
- has transpose effect, which is similar to `load_gather`. Therefore, a transpose attribute is
- introduced on purpose, making sure users are aware of this implicit transformation.
-
- In lane level, the result is a 1D vector that represents the data to be stored by
- each lane. If size is not 1, size should be equal to the chunk size.
+ typically a 1D vector. At lane level each lane may store a 1D vector of
+ contiguous elements; that per-lane width is implied by the value type and does
+ not need a separate attribute.
This operation serves as an anchor through which users assign a layout attribute
to govern computation distribution.
@@ -864,8 +853,6 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
mask is a vector of size equal to the subgroup size, or 1 at lane level.
scalar mask is also valid for lane level.
- - `chunk_size`: [optional] represents contiguous number of elements to store to per work item.
-
- `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache.
- `layout`: [optional] Describes the expected layout of the `tensor_desc` operand or the value
@@ -888,8 +875,9 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
```
Example 2 (Lane level):
- Lane level IR only accepts the offsets variant. chunk_size can be inferred from value
- type. In this example, chunk_size is 8.
+ Lane level IR only accepts the offsets variant. Each lane stores a contiguous
+ 1D vector whose length is given by the value type; here each lane stores 8
+ elements.
```mlir
xegpu.store %0, %1[%2], %3 <{l1_hint = #xegpu.cache_hint<uncached>,
l2_hint = #xegpu.cache_hint<write_back>,
@@ -902,7 +890,7 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
let arguments = (ins XeGPU_ValueOrScalarType:$value,
XeGPU_GatherScatterBaseAddrType:$dest,
AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
- AnyTypeOf<[XeGPU_MaskType, I1]>:$mask, OptionalAttr<I64Attr>:$chunk_size,
+ AnyTypeOf<[XeGPU_MaskType, I1]>:$mask,
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
@@ -947,13 +935,11 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
let builders = [
OpBuilder<(ins "Value": $value, "Value": $dest,
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
- "IntegerAttr": $chunk_size,
"xegpu::CachePolicyAttr": $l1_hint,
"xegpu::CachePolicyAttr": $l2_hint,
"xegpu::CachePolicyAttr": $l3_hint)>,
OpBuilder<(ins "Value": $value, "Value": $dest,
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
- "IntegerAttr": $chunk_size,
"xegpu::CachePolicyAttr": $l1_hint,
"xegpu::CachePolicyAttr": $l2_hint,
"xegpu::CachePolicyAttr": $l3_hint,
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index defc389006af5..4fb8cf65d8857 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -166,12 +166,6 @@ DistributeLayoutAttr inferExtractSourceLayout(DistributeLayoutAttr resLayout,
ArrayRef<int64_t> resShape,
ArrayRef<int64_t> srcShape);
-/// Infers the layout attribute for mask and offset operand for Chunked load
-/// and store, given the anchor layout attribute for the value being load/store.
-DistributeLayoutAttr
-inferMaskOffsetLayoutForScatterIO(DistributeLayoutAttr payloadLayout,
- int chunkSize);
-
/// Infers the source layout attribute for an operand using result layout
/// attribute
DistributeLayoutAttr
@@ -277,9 +271,10 @@ DistributeLayoutAttr setupInsertStridedSliceResultLayout(
DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for a load gather operation.
-DistributeLayoutAttr setupLoadGatherAnchorLayout(
- LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize,
- DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
+DistributeLayoutAttr
+setupLoadGatherAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
+ DistributeLayoutAttr consumerLayout,
+ const uArch::uArch *uArch);
/// Sets up the anchor layout for load matrix operation.
DistributeLayoutAttr setupLoadMatrixAnchorLayout(
@@ -290,7 +285,6 @@ DistributeLayoutAttr setupLoadMatrixAnchorLayout(
/// `numSg` is only used for Subgroup-kind layouts.
DistributeLayoutAttr setupStoreScatterAnchorLayout(LayoutKind layoutKind,
VectorType vectorTy,
- int contigChunkSize,
int numSg,
const uArch::uArch *uArch);
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 9863206f14fe1..a17f955def4ec 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -461,7 +461,6 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
vectorShape);
auto gatherOp = xegpu::LoadGatherOp::create(
rewriter, loc, vectorType, flatMemref, localOffsets, mask,
- /*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
@@ -496,7 +495,6 @@ static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
vectorShape);
xegpu::StoreScatterOp::create(rewriter, loc, writeOp.getVector(), flatMemref,
localOffsets, mask,
- /*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
@@ -741,7 +739,6 @@ struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
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{},
@@ -776,7 +773,6 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
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{},
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 686e4342215b0..4129dee368a93 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -64,53 +64,39 @@ static bool isWriteHintOrNone(const CachePolicyAttr &attr) {
static LogicalResult
isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
- VectorType valueTy, int64_t chunkSize,
+ VectorType valueTy,
function_ref<InFlightDiagnostic()> emitError) {
auto maskVecTy = dyn_cast<VectorType>(maskTy);
auto offsetsVecTy = dyn_cast<VectorType>(offsetsTy);
+
+ // The offsets and mask describe the same lane grid: either both are scalar
+ // (single lane) or both are vectors carrying one offset and one mask bit per
+ // lane, so they must have the same shape.
+ if (static_cast<bool>(maskVecTy) != static_cast<bool>(offsetsVecTy))
+ return emitError() << "Expecting offsets and mask to both be scalar or "
+ "both be vectors.";
+ if (maskVecTy && getShapeOf(maskTy) != getShapeOf(offsetsTy))
+ return emitError() << "Expecting offsets and mask to have the same shape.";
+
+ // Scalar payload (e.g. `index, i1 -> f16`): the offsets and mask must be
+ // scalar as well.
if (!valueTy) {
- if (chunkSize > 1)
- return emitError() << "Expecting chunk size == 1 for scalar result";
if (maskVecTy || offsetsVecTy)
return emitError() << "Expecting scalar mask and offsets.";
- else if (maskVecTy && offsetsVecTy)
- return emitError() << "Expecting a vector type result.";
- return success();
- }
-
- auto valueSize = valueTy.getNumElements();
- // SIMT mode with scalar mask and offsets.
- if (!maskVecTy && !offsetsVecTy) {
- if (valueSize != chunkSize)
- return emitError() << "value elements must match chunk size "
- << chunkSize;
return success();
}
- auto maskShape = getShapeOf(maskTy);
- auto valueShape = getShapeOf(valueTy);
- if (!maskVecTy)
- return emitError() << "Expecting a vector type mask.";
- int64_t maskSize = maskVecTy.getNumElements();
-
- if (chunkSize > 1) {
- if ((valueTy.getRank() == 1) && (valueSize != chunkSize))
- return emitError() << "value elements must match chunk size "
- << chunkSize;
- } else {
- if (valueSize != maskSize)
- return emitError()
- << "Mask should match value except the chunk size dim.";
- }
- llvm::SmallVector<int64_t> expectedMaskShape(valueShape);
+ // A scalar or size-1 mask means a single lane loads/stores the whole value as
+ // a contiguous 1D block, so any value shape is valid.
+ int64_t maskSize = maskVecTy ? maskVecTy.getNumElements() : 1;
if (maskSize == 1)
return success();
- if (chunkSize > 1)
- expectedMaskShape.pop_back();
- if (expectedMaskShape != maskShape)
- return emitError() << "Mask should match value except the chunk size dim.";
+ // Otherwise this is a congruent gather/scatter: one element per lane, so the
+ // value must have exactly the same shape as the mask.
+ if (getShapeOf(valueTy) != getShapeOf(maskTy))
+ return emitError() << "Value shape must match mask shape.";
return success();
}
@@ -605,7 +591,6 @@ LogicalResult LoadGatherOp::verify() {
return emitOpError("invalid l3_hint: ") << getL3HintAttr();
auto srcTy = getSourceType();
- uint64_t chunkSize = static_cast<int64_t>(getChunkSize().value_or(1));
auto memTy = dyn_cast<MemRefType>(srcTy);
if (memTy && (getElementType() != memTy.getElementType()))
@@ -620,14 +605,14 @@ LogicalResult LoadGatherOp::verify() {
if (failed(isValidContiguity(getContiguity(), offsetsTy,
[&]() { return emitOpError(); })))
return failure();
- return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
+ return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy,
[&]() { return emitOpError(); });
}
void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
Type valueType, Value source,
ArrayRef<OpFoldResult> offsets, Value mask,
- IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint,
+ xegpu::CachePolicyAttr l1_hint,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint) {
auto loc = source.getLoc();
@@ -636,15 +621,15 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto values = getValueOrCreateConstantIndexOp(builder, loc, offsets);
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
- build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
- l2_hint, l3_hint, /*anchor_layout=*/nullptr,
+ build(builder, state, valueType, source, offset, mask, l1_hint, l2_hint,
+ l3_hint, /*anchor_layout=*/nullptr,
/*contiguity=*/nullptr);
}
void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
Type valueType, Value source,
ArrayRef<OpFoldResult> offsets, Value mask,
- IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint,
+ xegpu::CachePolicyAttr l1_hint,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint,
DistributeLayoutAttr layout) {
@@ -654,8 +639,8 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto values = getValueOrCreateConstantIndexOp(builder, loc, offsets);
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
- build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
- l2_hint, l3_hint, layout, /*contiguity=*/nullptr);
+ build(builder, state, valueType, source, offset, mask, l1_hint, l2_hint,
+ l3_hint, layout, /*contiguity=*/nullptr);
}
//===----------------------------------------------------------------------===//
@@ -675,7 +660,6 @@ LogicalResult StoreScatterOp::verify() {
return emitOpError("invalid l3_hint: ") << getL3HintAttr();
auto destTy = getDestType();
- uint64_t chunkSize = static_cast<int64_t>(getChunkSize().value_or(1));
auto memTy = dyn_cast<MemRefType>(destTy);
if (memTy && (getElementType() != memTy.getElementType()))
@@ -690,14 +674,13 @@ LogicalResult StoreScatterOp::verify() {
if (failed(isValidContiguity(getContiguity(), offsetsTy,
[&]() { return emitOpError(); })))
return failure();
- return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
+ return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy,
[&]() { return emitOpError(); });
}
void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
Value value, Value dest,
ArrayRef<OpFoldResult> offsets, Value mask,
- IntegerAttr chunk_size,
xegpu::CachePolicyAttr l1_hint,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint) {
@@ -708,15 +691,17 @@ void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
// Call the correct builder overload that does not expect result types.
- build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
- l3_hint, /*anchor_layout=*/nullptr, /*contiguity=*/nullptr);
+ build(builder, state, value, dest, offset, mask, l1_hint, l2_hint, l3_hint,
+ /*anchor_layout=*/nullptr, /*contiguity=*/nullptr);
}
-void StoreScatterOp::build(
- OpBuilder &builder, OperationState &state, Value value, Value dest,
- ArrayRef<OpFoldResult> offsets, Value mask, IntegerAttr chunk_size,
- xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint,
- xegpu::CachePolicyAttr l3_hint, DistributeLayoutAttr layout) {
+void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
+ Value value, Value dest,
+ ArrayRef<OpFoldResult> offsets, Value mask,
+ xegpu::CachePolicyAttr l1_hint,
+ xegpu::CachePolicyAttr l2_hint,
+ xegpu::CachePolicyAttr l3_hint,
+ DistributeLayoutAttr layout) {
auto loc = dest.getLoc();
int64_t size = static_cast<int64_t>(offsets.size());
auto type = VectorType::get(size, builder.getIndexType());
@@ -724,8 +709,8 @@ void StoreScatterOp::build(
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
// Call the correct builder overload that does not expect result types.
- build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
- l3_hint, layout, /*contiguity=*/nullptr);
+ build(builder, state, value, dest, offset, mask, l1_hint, l2_hint, l3_hint,
+ layout, /*contiguity=*/nullptr);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index b11f7ecd3df06..d7c098db42a21 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1002,17 +1002,6 @@ xegpu::DistributeLayoutAttr xegpu::inferResultLayoutFromSourceForNonAnchorOp(
return nullptr;
}
-/// Infers the layout attribute for mask and offset operand for Chunked load
-/// and store, given the anchor layout attribute for the value being load/store.
-xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
- xegpu::DistributeLayoutAttr payloadLayout, int chunkSize) {
- auto rank = payloadLayout.getRank();
- if (chunkSize > 1)
- return payloadLayout.dropDims(
- llvm::to_vector(llvm::seq<int64_t>(rank - 1, rank)));
- return payloadLayout;
-}
-
//===----------------------------------------------------------------------===//
// Layout derivation helpers: factorize sgCount into
// sg_layout candidates, then
@@ -1849,13 +1838,20 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
///
/// For Subgroup layout, uses the consumer layout directly.
///
-/// For InstData layout, takes consumer's inst_data as-is. lane_layout and
-/// lane_data are taken from the consumer when present; otherwise the helper
-/// derives the standard scatter-style default (subgroupSize lanes on the
-/// innermost dim, per-lane vector capped by maxChunkSize).
+/// For InstData layout, takes consumer's inst_data as-is; lane_layout and
+/// lane_data are taken from the consumer.
+///
+/// For Lane layout, lane_layout/lane_data are taken from the consumer.
///
-/// For Lane layout, lane_layout/lane_data are taken from the consumer when
-/// present; otherwise derived from the same default.
+/// TODO: `maxChunkSize` (the per-lane contiguous-chunk cap from the uArch's max
+/// lane access size, or 1 for a plain gather) is currently NOT consumed on the
+/// load path: a load always has a consumer whose layout dictates lane_layout /
+/// lane_data, so the cap is inherited implicitly. It is retained here for the
+/// planned LoadMatrix/StoreMatrix coalescing, which will need to *derive* a
+/// chunked default when no consumer layout is available -- mirroring the store
+/// path, where `setupGenericStoreAnchorLayout` already applies it via
+/// `computeScatterIOLaneLayoutAndData(..., maxChunkSize)`. Do not remove it as
+/// dead code before that lands.
static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize,
@@ -1871,6 +1867,9 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
SmallVector<int64_t> consumerLaneData =
consumerLayout.getEffectiveLaneDataAsInt();
+ // Take lane_layout / lane_data from the consumer. (See the TODO above: the
+ // `maxChunkSize`-capped derive path used by the store side is not yet needed
+ // here, but will be for matrix coalescing.)
SmallVector<int64_t> laneLayout;
SmallVector<int64_t> laneData;
assert(!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
@@ -1892,17 +1891,16 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
/// Sets up the anchor layout for a load gather operation.
xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
- xegpu::LayoutKind layoutKind, VectorType resVecTy, int contigChunkSize,
+ xegpu::LayoutKind layoutKind, VectorType resVecTy,
xegpu::DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
ArrayRef<int64_t> resShape = resVecTy.getShape();
auto context = resVecTy.getContext();
- const auto *uArchInstruction = dyn_cast<xegpu::uArch::LoadGatherInstruction>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize =
- std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
+ // Gather loads one element per lane, so the innermost per-lane vector is
+ // capped at a single element.
+ int maxChunkSize = 1;
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
maxChunkSize, resShape, subgroupSize);
@@ -1972,18 +1970,16 @@ static xegpu::DistributeLayoutAttr setupGenericStoreAnchorLayout(
/// Sets up the anchor layout for a store scatter operation.
xegpu::DistributeLayoutAttr
xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy, int contigChunkSize,
- int numSg, const uArch::uArch *uArch) {
+ VectorType srcVecTy, int numSg,
+ const uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
ArrayRef<int64_t> srcShape = srcVecTy.getShape();
auto context = srcVecTy.getContext();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::StoreScatterInstruction>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize =
- std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
+ // Scatter stores one element per lane, so the innermost per-lane vector is
+ // capped at a single element.
+ int maxChunkSize = 1;
return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
srcShape, subgroupSize, numSg);
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 19d02e02b02ba..92a523c6bceb2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1234,7 +1234,6 @@ void LayoutInfoPropagation::visitLoadGatherOp(
if (!uArch)
return;
VectorType resVecTy = load.getValueType();
- int chunkSize = load.getChunkSize().value_or(1);
LayoutInfo resLayoutInfo = results[0]->getValue();
if (!resLayoutInfo.isAssigned())
@@ -1268,14 +1267,12 @@ void LayoutInfoPropagation::visitLoadGatherOp(
return;
}
requiredAnchorLayoutAttr = xegpu::setupLoadGatherAnchorLayout(
- layoutKind, resVecTy, chunkSize, consumerLayoutAttr, uArch);
+ layoutKind, resVecTy, consumerLayoutAttr, uArch);
load.setLayoutAttr(requiredAnchorLayoutAttr);
}
- assert((chunkSize <= 1) || (layoutKind != xegpu::LayoutKind::Subgroup));
- auto maskLayoutAttr = xegpu::inferMaskOffsetLayoutForScatterIO(
- requiredAnchorLayoutAttr, chunkSize);
- LayoutInfo maskLayoutInfo = makeLayoutInfo(maskLayoutAttr);
+ // The mask and offset operands share the value's anchor layout.
+ LayoutInfo maskLayoutInfo = makeLayoutInfo(requiredAnchorLayoutAttr);
auto loadLayoutInfo = makeLayoutInfo(requiredAnchorLayoutAttr);
// Propagate the new layout to the tensor descriptor operand.
@@ -1299,7 +1296,6 @@ void LayoutInfoPropagation::visitStoreScatterOp(
if (!uArch)
return;
VectorType srcVecTy = storeScatter.getValueType();
- int chunkSize = storeScatter.getChunkSize().value_or(1);
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
@@ -1331,7 +1327,7 @@ void LayoutInfoPropagation::visitStoreScatterOp(
if (failed(numSgOrErr))
return;
requiredAnchorLayoutAttr = xegpu::setupStoreScatterAnchorLayout(
- layoutKind, srcVecTy, chunkSize, numSgOrErr.value_or(0), uArch);
+ layoutKind, srcVecTy, numSgOrErr.value_or(0), uArch);
if (!requiredAnchorLayoutAttr) {
markFailure(storeScatter,
"Failed to determine required layout for store scatter.");
@@ -1341,10 +1337,8 @@ void LayoutInfoPropagation::visitStoreScatterOp(
}
LayoutInfo srcLayoutInfo = makeLayoutInfo(requiredAnchorLayoutAttr);
- assert((chunkSize <= 1) || (layoutKind != xegpu::LayoutKind::Subgroup));
- auto maskLayoutAttr = xegpu::inferMaskOffsetLayoutForScatterIO(
- requiredAnchorLayoutAttr, chunkSize);
- LayoutInfo maskLayoutInfo = makeLayoutInfo(maskLayoutAttr);
+ // The mask and offset operands share the value's anchor layout.
+ LayoutInfo maskLayoutInfo = makeLayoutInfo(requiredAnchorLayoutAttr);
// Propagate the payload operand layout
propagateIfChanged(operands[0], operands[0]->meet(srcLayoutInfo));
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index bc85140f9f121..45e0bac55c275 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -497,7 +497,7 @@ struct SgToLanePrefetchNd : public OpConversionPattern<xegpu::PrefetchNdOp> {
/// Distributes a subgroup-level LoadGather (xegpu.load) op to lane-level.
///
-/// Example 1 (1D, no chunk size):
+/// Example 1 (1D):
/// layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
/// %mask = producer_op : vector<16xi1>
/// %offset = producer_op : vector<16xindex>
@@ -509,15 +509,7 @@ struct SgToLanePrefetchNd : public OpConversionPattern<xegpu::PrefetchNdOp> {
/// %0 = xegpu.load %src[%offset], %mask : memref<256xf16>,
/// vector<1xindex>, vector<1xi1> -> vector<1xf16>
///
-/// Example 2 (2D with chunk size, same mask & offset):
-/// layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
-/// %0 = xegpu.load %src[%offset], %mask <{chunk_size=8}> :
-/// memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
-/// Distributed to:
-/// %0 = xegpu.load %src[%offset], %mask <{chunk_size=8}> :
-/// memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-///
-/// Example 3 (3D with leading unit dims):
+/// Example 2 (3D with leading unit dims):
/// layout = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>
/// %mask = producer_op : vector<1x1x16xi1>
/// %offset = producer_op : vector<1x1x16xindex>
@@ -543,8 +535,7 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
return failure();
// Check that leading dimensions are unit.
- int chunkSize = op.getChunkSize().value_or(1);
- int effectiveVecRank = (chunkSize == 1) ? 1 : 2;
+ int effectiveVecRank = 1;
ArrayRef<int64_t> shape = origResultTy.getShape();
if (llvm::any_of(
shape.take_front(origResultTy.getRank() - effectiveVecRank),
@@ -581,8 +572,8 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
Value distSource = adaptor.getSource();
auto newOp = xegpu::LoadGatherOp::create(
rewriter, op.getLoc(), distResultTy1D, distSource, distOffsets,
- distMask, op.getChunkSizeAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/nullptr, /*contiguity=*/nullptr);
+ distMask, op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(),
+ /*layout=*/nullptr, /*contiguity=*/nullptr);
Value result = newOp->getResult(0);
if (distResultTy1D != distResultTy)
@@ -1024,7 +1015,7 @@ struct SgToLaneStoreMatrix : public OpConversionPattern<xegpu::StoreMatrixOp> {
/// Distributes a subgroup-level StoreScatter (xegpu.store) op to
/// lane-level.
///
-/// Example 1 (1D, no chunk size):
+/// Example 1 (1D):
/// layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
/// %mask = producer_op : vector<16xi1>
/// %offset = producer_op : vector<16xindex>
@@ -1036,15 +1027,7 @@ struct SgToLaneStoreMatrix : public OpConversionPattern<xegpu::StoreMatrixOp> {
/// xegpu.store %payload, %src[%offset], %mask : vector<1xf16>,
/// memref<256xf16>, vector<1xindex>, vector<1xi1>
///
-/// Example 2 (2D with chunk size, same mask & offset):
-/// layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
-/// xegpu.store %payload, %src[%offset], %mask <{chunk_size=8}> :
-/// vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
-/// Distributed to:
-/// xegpu.store %payload, %src[%offset], %mask <{chunk_size=8}> :
-/// vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
-///
-/// Example 3 (3D with leading unit dims):
+/// Example 2 (3D with leading unit dims):
/// layout = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>
/// %mask = producer_op : vector<1x1x16xi1>
/// %offset = producer_op : vector<1x1x16xindex>
@@ -1071,8 +1054,7 @@ struct SgToLaneStoreScatter
return failure();
// Check that all leading dimensions are unit dimensions.
- int chunkSize = op.getChunkSize().value_or(1);
- int effectiveVecRank = (chunkSize == 1) ? 1 : 2;
+ int effectiveVecRank = 1;
ArrayRef<int64_t> shape = origValueTy.getShape();
if (llvm::any_of(shape.take_front(origValueTy.getRank() - effectiveVecRank),
[](int64_t d) { return d != 1; }))
@@ -1112,9 +1094,9 @@ struct SgToLaneStoreScatter
Value distDest = adaptor.getDest();
xegpu::StoreScatterOp::create(rewriter, op.getLoc(), distValue, distDest,
- distOffsets, distMask, op.getChunkSizeAttr(),
- op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/nullptr,
+ distOffsets, distMask, op.getL1HintAttr(),
+ op.getL2HintAttr(), op.getL3HintAttr(),
+ /*layout=*/nullptr,
/*contiguity=*/nullptr);
rewriter.eraseOp(op);
return success();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 5914e23f1f11a..be3a1b66c8b7d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -563,11 +563,6 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
return failure();
SmallVector<int64_t> targetMaskShape(*targetShape);
- int64_t chunkSize = 1;
- if (auto chunkSizeAttr = op->getAttr("chunk_size")) {
- if (auto intAttr = llvm::dyn_cast<IntegerAttr>(chunkSizeAttr))
- chunkSize = intAttr.getInt();
- }
// Unroll mask and offsets with correct shape
VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
@@ -575,50 +570,15 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
Type elemTy = valueTy.getElementType();
VectorType newValueTy = VectorType::get(*targetShape, elemTy);
- SmallVector<Type> convertedMaskTypes;
- SmallVector<Value> convertedMasks;
- SmallVector<Type> convertedOffsetTypes;
- SmallVector<Value> convertedOffsets;
-
- if (chunkSize > 1) {
- // For chunked loads, mask and offsets have one less dimension
- targetMaskShape.pop_back();
- int64_t blockedChunkSize = targetShape->back();
- int64_t numNewChunks = chunkSize / blockedChunkSize;
- chunkSize = blockedChunkSize;
-
- convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
- convertedOffsetTypes = getUnrolledTypes(offsetsTy, targetMaskShape);
-
- SmallVector<Value> convertedMasksBase =
- pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
- SmallVector<Value> convertedOffsetsBase =
- pack(offsets, convertedOffsetTypes, targetMaskShape, loc, rewriter);
-
- for (auto maskVal : convertedMasksBase)
- convertedMasks.append(numNewChunks, maskVal);
-
- for (auto [baseOffset, offsetType] :
- llvm::zip(convertedOffsetsBase, convertedOffsetTypes)) {
- for (int64_t i = 0; i < numNewChunks; ++i) {
- Value inc = arith::ConstantIndexOp::create(rewriter, loc,
- i * blockedChunkSize);
- Value incVec =
- vector::BroadcastOp::create(rewriter, loc, offsetType, inc);
- Value offsetVal =
- arith::AddIOp::create(rewriter, loc, baseOffset, incVec);
- convertedOffsets.push_back(offsetVal);
- }
- }
- } else {
- convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
- convertedMasks =
- pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
+ SmallVector<Type> convertedMaskTypes =
+ getUnrolledTypes(maskTy, targetMaskShape);
+ SmallVector<Value> convertedMasks =
+ pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
- convertedOffsetTypes = getUnrolledTypes(offsetsTy, *targetShape);
- convertedOffsets =
- pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
- }
+ SmallVector<Type> convertedOffsetTypes =
+ getUnrolledTypes(offsetsTy, *targetShape);
+ SmallVector<Value> convertedOffsets =
+ pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
auto layout = op.getLayoutAttr();
if (layout)
@@ -627,8 +587,7 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
SmallVector<Value> newOps;
for (auto [o, m] : llvm::zip(convertedOffsets, convertedMasks)) {
auto newOp = xegpu::LoadGatherOp::create(
- rewriter, loc, newValueTy, op.getSource(), o, m,
- rewriter.getI64IntegerAttr(chunkSize), op.getL1HintAttr(),
+ rewriter, loc, newValueTy, op.getSource(), o, m, op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr(), layout,
/*contiguity=*/nullptr);
newOps.push_back(newOp);
@@ -657,59 +616,19 @@ struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
if (!targetShape)
return failure();
- int64_t chunkSize = 1;
- if (auto chunkSizeAttr = op->getAttr("chunk_size")) {
- if (auto intAttr = llvm::dyn_cast<IntegerAttr>(chunkSizeAttr))
- chunkSize = intAttr.getInt();
- }
-
SmallVector<int64_t> targetMaskShape(*targetShape);
VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
VectorType offsetsTy = llvm::dyn_cast<VectorType>(offsets.getType());
- SmallVector<Type> convertedMaskTypes;
- SmallVector<Value> convertedMasks;
- SmallVector<Type> convertedOffsetTypes;
- SmallVector<Value> convertedOffsets;
-
- if (chunkSize > 1) {
- targetMaskShape.pop_back();
- int64_t blockedChunkSize = targetShape->back();
- int64_t numNewChunks = chunkSize / blockedChunkSize;
- chunkSize = blockedChunkSize;
-
- convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
- convertedOffsetTypes = getUnrolledTypes(offsetsTy, targetMaskShape);
-
- SmallVector<Value> convertedMasksBase =
- pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
- SmallVector<Value> convertedOffsetsBase =
- pack(offsets, convertedOffsetTypes, targetMaskShape, loc, rewriter);
-
- for (auto maskVal : convertedMasksBase)
- convertedMasks.append(numNewChunks, maskVal);
-
- for (auto [baseOffset, offsetType] :
- llvm::zip(convertedOffsetsBase, convertedOffsetTypes)) {
- for (int64_t i = 0; i < numNewChunks; ++i) {
- Value inc = arith::ConstantIndexOp::create(rewriter, loc,
- i * blockedChunkSize);
- Value incVec =
- vector::BroadcastOp::create(rewriter, loc, offsetType, inc);
- Value offsetVal =
- arith::AddIOp::create(rewriter, loc, baseOffset, incVec);
- convertedOffsets.push_back(offsetVal);
- }
- }
- } else {
- convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
- convertedMasks =
- pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
+ SmallVector<Type> convertedMaskTypes =
+ getUnrolledTypes(maskTy, targetMaskShape);
+ SmallVector<Value> convertedMasks =
+ pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
- convertedOffsetTypes = getUnrolledTypes(offsetsTy, *targetShape);
- convertedOffsets =
- pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
- }
+ SmallVector<Type> convertedOffsetTypes =
+ getUnrolledTypes(offsetsTy, *targetShape);
+ SmallVector<Value> convertedOffsets =
+ pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
SmallVector<Type> convertedValTypes =
getUnrolledTypes(valueTy, *targetShape);
@@ -723,7 +642,6 @@ struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
for (auto [v, o, m] :
llvm::zip(convertedValues, convertedOffsets, convertedMasks)) {
xegpu::StoreScatterOp::create(rewriter, loc, v, op.getDest(), o, m,
- rewriter.getI64IntegerAttr(chunkSize),
op.getL1HintAttr(), op.getL2HintAttr(),
op.getL3HintAttr(), layout,
/*contiguity=*/nullptr);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 0e8a386fb08b6..07c067ce076f2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -828,14 +828,12 @@ struct WgToSgLoadGatherOp : public OpConversionPattern<xegpu::LoadGatherOp> {
}
SmallVector<Value> newLoadOps;
- auto chunkSizeAttr =
- rewriter.getI64IntegerAttr(op.getChunkSize().value_or(1));
VectorType newTy = VectorType::get(sgShape, resultType.getElementType());
for (auto [offsets, mask] :
llvm::zip(adaptor.getOffsets(), adaptor.getMask())) {
auto newLayout = layout.dropSgLayoutAndData();
auto newLoadOp = xegpu::LoadGatherOp::create(
- rewriter, loc, newTy, op.getSource(), offsets, mask, chunkSizeAttr,
+ rewriter, loc, newTy, op.getSource(), offsets, mask,
op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(), newLayout,
/*contiguity=*/nullptr);
newLoadOps.push_back(newLoadOp);
@@ -875,16 +873,12 @@ struct WgToSgStoreScatterOp
"offsets have not been distributed");
}
- auto chunkSizeOpt = op.getChunkSize();
- int64_t chunkSize = chunkSizeOpt ? static_cast<int64_t>(*chunkSizeOpt) : 1;
- auto chunkSizeAttr = rewriter.getI64IntegerAttr(chunkSize);
for (auto [val, offs, mask] : llvm::zip(
adaptor.getValue(), adaptor.getOffsets(), adaptor.getMask())) {
- xegpu::StoreScatterOp::create(rewriter, loc, val, op.getDest(), offs,
- mask, chunkSizeAttr, op.getL1HintAttr(),
- op.getL2HintAttr(), op.getL3HintAttr(),
- layout.dropSgLayoutAndData(),
- /*contiguity=*/nullptr);
+ xegpu::StoreScatterOp::create(
+ rewriter, loc, val, op.getDest(), offs, mask, op.getL1HintAttr(),
+ op.getL2HintAttr(), op.getL3HintAttr(), layout.dropSgLayoutAndData(),
+ /*contiguity=*/nullptr);
}
rewriter.eraseOp(op);
return success();
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 76269cf193d13..67097cfd3b6aa 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -235,22 +235,9 @@ xegpu::getDistributeLayoutAttr(const OpOperand &opr) {
if (isa<xegpu::StoreNdOp, xegpu::StoreMatrixOp>(op) && (idx < 2))
return layout;
- if (isa<xegpu::StoreScatterOp>(op)) {
- xegpu::StoreScatterOp store(op);
- int chunkSize = store.getChunkSize().value_or(1);
- if (layout && idx >= 2 && chunkSize > 1)
- return layout.dropDims(llvm::to_vector(
- llvm::seq<int64_t>(layout.getRank() - 1, layout.getRank())));
+ // For gather/scatter ops the mask and offsets share the value's layout.
+ if (isa<xegpu::StoreScatterOp, xegpu::LoadGatherOp>(op))
return layout;
- }
- if (isa<xegpu::LoadGatherOp>(op)) {
- xegpu::LoadGatherOp load(op);
- int chunkSize = load.getChunkSize().value_or(1);
- if (layout && idx >= 1 && chunkSize > 1)
- return layout.dropDims(llvm::to_vector(
- llvm::seq<int64_t>(layout.getRank() - 1, layout.getRank())));
- return layout;
- }
}
std::string layoutName = xegpu::getTemporaryLayoutName(opr);
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index 932cf6ff9b201..33aa414c69e4f 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -244,9 +244,9 @@ func.func @load_gather_vc_2(%src: memref<?xf32>) {
func.func @load_gather_vc_3(%src: memref<?xf32>) {
%offsets = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
%mask = arith.constant dense<1>: vector<8xi1>
- // expected-error at +1 {{Mask should match value except the chunk size dim}}
- %2 = xegpu.load %src[%offsets], %mask <{chunk_size = 2}>
- : memref<?xf32>, vector<4xindex>, vector<8xi1> -> vector<4x2xf32>
+ // expected-error at +1 {{Expecting offsets and mask to have the same shape}}
+ %2 = xegpu.load %src[%offsets], %mask
+ : memref<?xf32>, vector<4xindex>, vector<8xi1> -> vector<4xf32>
return
}
@@ -254,8 +254,8 @@ func.func @load_gather_vc_3(%src: memref<?xf32>) {
func.func @load_gather_simt_1(%src: memref<?xf32>) {
%0 = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
%1 = arith.constant dense<1>: vector<4xi1>
- // expected-error at +1 {{value elements must match chunk size}}
- %2 = xegpu.load %src[%0], %1 <{chunk_size = 2}>
+ // expected-error at +1 {{Value shape must match mask shape}}
+ %2 = xegpu.load %src[%0], %1
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<6xf32>
return
}
@@ -275,10 +275,10 @@ func.func @store_scatter_vc_2(%dst: memref<?xf32>) {
func.func @store_scatter_vc_3(%dst: memref<?xf32>) {
%0 = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
%1 = arith.constant dense<1>: vector<8xi1>
- %2 = arith.constant dense<2.9>: vector<4x2xf32>
- // expected-error at +1 {{Mask should match value except the chunk size dim}}
- xegpu.store %2, %dst[%0], %1 <{chunk_size = 2}>
- : vector<4x2xf32>, memref<?xf32>, vector<4xindex>, vector<8xi1>
+ %2 = arith.constant dense<2.9>: vector<4xf32>
+ // expected-error at +1 {{Expecting offsets and mask to have the same shape}}
+ xegpu.store %2, %dst[%0], %1
+ : vector<4xf32>, memref<?xf32>, vector<4xindex>, vector<8xi1>
return
}
@@ -287,8 +287,8 @@ func.func @store_scatter_simt_1(%dst: memref<?xf32>) {
%0 = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
%1 = arith.constant dense<1>: vector<4xi1>
%2 = arith.constant dense<2.9>: vector<6xf32>
- // expected-error at +1 {{value elements must match chunk size}}
- xegpu.store %2, %dst[%0], %1 <{chunk_size = 2}>
+ // expected-error at +1 {{Value shape must match mask shape}}
+ xegpu.store %2, %dst[%0], %1
: vector<6xf32>, memref<?xf32>, vector<4xindex>, vector<4xi1>
return
}
@@ -321,30 +321,10 @@ func.func @prefetch_offset_wi_5(%src: i64) {
func.func @load_gather_offset_sg(%src: memref<?xf16>) {
%offsets = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
%mask = arith.constant dense<1>: vector<8xi1>
- // expected-error at +1 {{Mask should match value except the chunk size dim}}
+ // expected-error at +1 {{Expecting offsets and mask to have the same shape}}
%2 = xegpu.load %src[%offsets], %mask
: memref<?xf16>, vector<4xindex>, vector<8xi1>
- -> vector<4x2xf16>
- return
-}
-
-// -----
-func.func @load_gather_offset_wi(%src: ui64) {
- %mask = arith.constant dense<1>: vector<1xi1>
- %offsets = arith.constant dense<[0]> : vector<1xindex>
- // expected-error at +1 {{value elements must match chunk size}}
- %2 = xegpu.load %src[%offsets], %mask <{chunk_size = 2}> : ui64, vector<1xindex>, vector<1xi1> -> vector<3xf32>
- return
-}
-
-// -----
-func.func @store_scatter_offset_wi_1(%src: memref<?xf16>) {
- %val = arith.constant dense<2.9>: vector<4xf16>
- %offsets = arith.constant dense<[0]> : vector<1xindex>
- %mask = arith.constant dense<1>: vector<1xi1>
- // expected-error at +1 {{Mask should match value except the chunk size dim}}
- xegpu.store %val, %src[%offsets], %mask
- : vector<4xf16>, memref<?xf16>, vector<1xindex>, vector<1xi1>
+ -> vector<4xf16>
return
}
@@ -375,16 +355,7 @@ func.func @load_gather_offset_wi_4(%src: !xegpu.tensor_desc<1x2xf16>) {
%mask = arith.constant dense<1>: vector<1xi1>
%offsets = arith.constant dense<[0]> : vector<1xindex>
// expected-error at +1 {{op operand #0 must be 1D memref}}
- %2 = xegpu.load %src[%offsets], %mask <{chunk_size = 2}> : !xegpu.tensor_desc<1x2xf16>, vector<1xindex>, vector<1xi1> -> vector<2xf16>
- return
-}
-
-// -----
-func.func @load_gather_offset_wi_2(%src: ui64) {
- %mask = arith.constant dense<1>: vector<1xi1>
- %offsets = arith.constant dense<[0]> : vector<1xindex>
- // expected-error at +1 {{value elements must match chunk size}}
- %2 = xegpu.load %src[%offsets], %mask <{chunk_size = 2}> : ui64, vector<1xindex>, vector<1xi1> -> vector<3xf16>
+ %2 = xegpu.load %src[%offsets], %mask : !xegpu.tensor_desc<1x2xf16>, vector<1xindex>, vector<1xi1> -> vector<2xf16>
return
}
@@ -393,7 +364,7 @@ func.func @load_gather_offset_wi_1(%src: memref<4x4xf32>) {
%mask = arith.constant dense<1>: vector<1xi1>
%offsets = arith.constant dense<[0]> : vector<1xindex>
// expected-error at +1 {{op operand #0 must be 1D memref}}
- %2 = xegpu.load %src[%offsets], %mask <{chunk_size = 2}> : memref<4x4xf32>, vector<1xindex>, vector<1xi1> -> vector<2xf32>
+ %2 = xegpu.load %src[%offsets], %mask : memref<4x4xf32>, vector<1xindex>, vector<1xi1> -> vector<2xf32>
return
}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index 5e795017e2122..a6283c8afd208 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -375,8 +375,8 @@ gpu.func @prefetch_nd_3d(%src: memref<4x8x16xf16>) {
// CHECK: gpu.func @simt_load_4(%[[arg0:.*]]: memref<256xf16>, %[[arg1:.*]]: vector<1xindex>, %[[arg2:.*]]: vector<1xi1>) {
gpu.func @simt_load_4(%arg0: memref<256xf16>, %arg1: vector<1xindex>, %arg2: vector<1xi1>) {
- // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{chunk_size = 8 : i64}> : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
- %0 = xegpu.load %arg0[%arg1], %arg2 <{chunk_size = 8 : i64}> : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
+ // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
+ %0 = xegpu.load %arg0[%arg1], %arg2 : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
gpu.return
}
@@ -389,8 +389,8 @@ gpu.func @simt_load_5(%arg0: memref<256xf16>, %arg1: vector<1xindex>, %arg2: vec
// CHECK: gpu.func @simt_load_6(%[[arg0:.*]]: memref<256xf16>, %[[arg1:.*]]: index, %[[arg2:.*]]: i1) {
gpu.func @simt_load_6(%arg0: memref<256xf16>, %arg1: index, %arg2: i1) {
- // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{chunk_size = 8 : i64}> : memref<256xf16>, index, i1 -> vector<8xf16>
- %0 = xegpu.load %arg0[%arg1], %arg2 <{chunk_size = 8 : i64}> : memref<256xf16>, index, i1 -> vector<8xf16>
+ // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] : memref<256xf16>, index, i1 -> vector<8xf16>
+ %0 = xegpu.load %arg0[%arg1], %arg2 : memref<256xf16>, index, i1 -> vector<8xf16>
gpu.return
}
@@ -401,27 +401,17 @@ gpu.func @simt_load_7(%arg0: memref<256xf16>, %arg1: index, %arg2: i1) {
gpu.return
}
-// CHECK: gpu.func @subgroup_load_offset_1(%arg0: memref<?xf16>) {
-gpu.func @subgroup_load_offset_1(%src: memref<?xf16>) {
- %offset = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
- %mask = arith.constant dense<1>: vector<4xi1>
- //CHECK: %[[R1:.*]] = xegpu.load %arg0[%cst], %cst_0 <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : memref<?xf16>, vector<4xindex>, vector<4xi1> -> vector<4x2xf16>
- %val = xegpu.load %src[%offset], %mask <{chunk_size=2, l1_hint = #xegpu.cache_hint<cached>}>
- : memref<?xf16>, vector<4xindex>, vector<4xi1> -> vector<4x2xf16>
- gpu.return
-}
-
// CHECK: gpu.func @simt_store_4(%[[arg0:.*]]: vector<8xf16>, %[[arg1:.*]]: memref<256xf16>, %[[arg2:.*]]: vector<1xindex>, %[[arg3:.*]]: vector<1xi1>) {
gpu.func @simt_store_4(%arg0: vector<8xf16>, %arg1: memref<256xf16>, %arg2: vector<1xindex>, %arg3: vector<1xi1>) {
- // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
- xegpu.store %arg0, %arg1[%arg2], %arg3 <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
+ xegpu.store %arg0, %arg1[%arg2], %arg3 : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
gpu.return
}
// CHECK: gpu.func @simt_store_5(%[[arg0:.*]]: vector<8xf16>, %[[arg1:.*]]: memref<256xf16>, %[[arg2:.*]]: index, %[[arg3:.*]]: i1) {
gpu.func @simt_store_5(%arg0: vector<8xf16>, %arg1: memref<256xf16>, %arg2: index, %arg3: i1) {
- // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, index, i1
- xegpu.store %arg0, %arg1[%arg2], %arg3 <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, index, i1
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] : vector<8xf16>, memref<256xf16>, index, i1
+ xegpu.store %arg0, %arg1[%arg2], %arg3 : vector<8xf16>, memref<256xf16>, index, i1
gpu.return
}
@@ -439,17 +429,6 @@ gpu.func @simt_store_7(%arg0: f16, %arg1: memref<256xf16>, %arg2: index, %arg3:
gpu.return
}
-// CHECK: gpu.func @subgroup_store_offset_1(%arg0: memref<?xf16>) {
-gpu.func @subgroup_store_offset_1(%dest: memref<?xf16>) {
- %val = arith.constant dense<2.9>: vector<4x2xf16>
- %offset = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
- %mask = arith.constant dense<1>: vector<4xi1>
- //CHECK: xegpu.store %[[R0:.*]], %arg0[%cst_0], %cst_1 <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<4x2xf16>, memref<?xf16>, vector<4xindex>, vector<4xi1>
- xegpu.store %val, %dest[%offset], %mask <{chunk_size=2, l1_hint = #xegpu.cache_hint<cached>}>
- : vector<4x2xf16>, memref<?xf16>, vector<4xindex>, vector<4xi1>
- gpu.return
-}
-
// CHECK: gpu.func @load_contiguity(%[[arg0:.*]]: i64, %[[arg1:.*]]: vector<16xindex>, %[[arg2:.*]]: vector<16xi1>) {
gpu.func @load_contiguity(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
// A user-provided `contiguity` round-trips through the optional op attribute.
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index 02d59f80892e7..d7ecf59048b81 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -596,9 +596,9 @@ gpu.module @test {
%val = arith.constant dense<25.5> : vector<256xf16>
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
- // CHECK: xegpu.store %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
+ // CHECK: xegpu.store %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}} <{l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
// CHECK-SAME: : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
- xegpu.store %val, %dest[%offset], %mask {chunk_size = 1, l1_hint = #xegpu.cache_hint<cached>}
+ xegpu.store %val, %dest[%offset], %mask {l1_hint = #xegpu.cache_hint<cached>}
: vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index b4e69173c1144..5b48c8ce759c9 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -641,8 +641,8 @@ gpu.module @test{
%5 = vector.broadcast %3 : index to vector<1xindex>
%6 = arith.addi %4, %5 : vector<1xindex>
%7 = vector.broadcast %6 : vector<1xindex> to vector<1x1x1x16xindex>
- xegpu.store %cst, %0[%7], %cst_0 <{chunk_size = 1 : i64}> : vector<1x1x1x16xf32>, i64, vector<1x1x1x16xindex>, vector<1x1x1x16xi1>
- xegpu.store %cst, %0[%7], %cst_0 <{chunk_size = 1 : i64}> : vector<1x1x1x16xf32>, i64, vector<1x1x1x16xindex>, vector<1x1x1x16xi1>
+ xegpu.store %cst, %0[%7], %cst_0 : vector<1x1x1x16xf32>, i64, vector<1x1x1x16xindex>, vector<1x1x1x16xi1>
+ xegpu.store %cst, %0[%7], %cst_0 : vector<1x1x1x16xf32>, i64, vector<1x1x1x16xindex>, vector<1x1x1x16xi1>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index 0f3cea16797f1..d3ccfaa54433e 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -277,42 +277,6 @@ gpu.func @prefetch_nd() {
gpu.return
}
-// CHECK-LABEL: gpu.func @scatter_load_chunksize
-// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
-// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<8xf16> to vector<1x8xf16>
-gpu.func @scatter_load_chunksize(%src: memref<256xf16>) {
- %offset = arith.constant dense<12> : vector<16xindex>
- %mask = arith.constant dense<true> : vector<16xi1>
- %0 = xegpu.load %src[%offset], %mask
- <{chunk_size = 8, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- gpu.return
-}
-
-// CHECK-LABEL: gpu.func @scatter_store_chunksize
-// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
-// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK: %[[C1:.*]] = vector.shape_cast %[[LOAD]] : vector<8xf16> to vector<1x8xf16>
-// CHECK: %[[C2:.*]] = vector.shape_cast %[[C1]] : vector<1x8xf16> to vector<8xf16>
-// CHECK: xegpu.store %[[C2]], %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
-gpu.func @scatter_store_chunksize(%src: memref<256xf16>) {
- %offset = arith.constant dense<12> : vector<16xindex>
- %mask = arith.constant dense<true> : vector<16xi1>
- %0 = xegpu.load %src[%offset], %mask
- <{chunk_size = 8, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %0, %src[%offset], %mask
- <{chunk_size = 8, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
- : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
- gpu.return
-}
-
// CHECK-LABEL: gpu.func @scatter_load
// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
index 85b939b6d054a..ad046a7e432ed 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
@@ -250,67 +250,6 @@ gpu.module @xevm_module{
}
}
-// -----
-// CHECK-LABEL: gpu.func @scatter_ops_scf_yield
-// CHECK: (%{{.*}}: memref<256xf16>, %[[PREDICATE:[a-zA-Z0-9]+]]: i1) {
-// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.200000e+01> : vector<1x8xf16>
-// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK-DAG: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
-// CHECK: %[[IF:.*]] = scf.if %[[PREDICATE]] -> (vector<1x8xf16>) {
-// CHECK-NEXT: %[[LD:.*]] = xegpu.load %{{.*}}[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK-NEXT: %[[LD_CAST:.*]] = vector.shape_cast %[[LD]] : vector<8xf16> to vector<1x8xf16>
-// CHECK-NEXT: scf.yield %[[LD_CAST]] : vector<1x8xf16>
-// CHECK-NEXT: } else {
-// CHECK-NEXT: scf.yield %[[CST]] : vector<1x8xf16>
-// CHECK-NEXT: }
-// CHECK-NEXT: %[[IF_CAST:.*]] = vector.shape_cast %[[IF]] : vector<1x8xf16> to vector<8xf16>
-// CHECK-NEXT: xegpu.store %[[IF_CAST]], %{{.*}}[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
-gpu.module @xevm_module{
- gpu.func @scatter_ops_scf_yield(%src: memref<256xf16>, %pred : i1) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %loaded = scf.if %pred -> (vector<16x8xf16>) {
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {
- layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
- } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- scf.yield %3 : vector<16x8xf16>
- } else {
- %3 = arith.constant dense<12.> : vector<16x8xf16>
- scf.yield %3 : vector<16x8xf16>
- }
- xegpu.store %loaded, %src[%offset], %1 <{chunk_size=8}> {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
- gpu.return
- }
-}
-
-// -----
-// CHECK-LABEL: gpu.func @scatter_ops_scf_non_yield({{.*}}) {
-// CHECK: %[[PREDICATE:.*]] = llvm.mlir.poison : i1
-// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
-// CHECK: scf.if %[[PREDICATE]] {
-// CHECK-NEXT: %[[LOADED:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK-NEXT: xegpu.store %[[LOADED]], %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
-// CHECK-NEXT: }
-gpu.module @xevm_module{
- gpu.func @scatter_ops_scf_non_yield(%src: memref<256xf16>) {
- %pred = llvm.mlir.poison : i1
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- scf.if %pred {
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {
- layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
- } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=8}> {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
- }
- gpu.return
- }
-}
-
// -----
// CHECK-LABEL: gpu.func @mma_transpose_b(
// CHECK: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<16x8xi32>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index 490f2a622624e..3cc17e9b643e2 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -334,8 +334,8 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184,
192, 200, 208, 216, 224, 232, 240, 248
]> : vector<32xindex>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%cst], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%cst], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -353,8 +353,8 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184,
192, 200, 208, 216, 224, 232, 240, 248
]> : vector<32xindex>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%cst], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%cst], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -369,8 +369,8 @@ gpu.module @test_kernel {
//CHECK: arith.addi [[step]], [[cst]] : vector<16xindex>
%step = vector.step : vector<32xindex>
%mask = vector.create_mask %c16 : vector<32xi1>
- %ld = xegpu.load %src[%step], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%step], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%step], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%step], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -557,7 +557,7 @@ gpu.module @test_kernel {
// -----
gpu.module @test_kernel {
// CHECK-LABEL: load_with_offsets
- // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.func @load_with_offsets(%src: ui64) -> vector<32xf32> {
%cst = arith.constant dense<[
0, 8, 16, 24, 32, 40, 48, 56,
@@ -568,7 +568,7 @@ gpu.module @test_kernel {
%c17 = arith.constant 17: index
%mask = vector.create_mask %c17 : vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %ld = xegpu.load %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %ld : vector<32xf32>
}
@@ -577,7 +577,7 @@ gpu.module @test_kernel {
// -----
gpu.module @test_kernel {
// CHECK-LABEL: store_with_offsets
- // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
+ // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
gpu.func @store_with_offsets(%src: ui64) {
%cst = arith.constant dense<[
0, 8, 16, 24, 32, 40, 48, 56,
@@ -590,62 +590,12 @@ gpu.module @test_kernel {
%mask = vector.create_mask %c17 : vector<32xi1>
%st_vec = arith.constant dense<1023.0>: vector<32xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ xegpu.store %st_vec, %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
-// -----
-gpu.module @test_kernel {
- // CHECK-LABEL: load_with_offsets_chunk
- // CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<32x4xf32>
- // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
- // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
- // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
- // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
- // CHECK-COUNT-4: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16x2xf32>
- gpu.func @load_with_offsets_chunk(%src: ui64) -> vector<32x4xf32> {
- %cst = arith.constant dense<[
- 0, 8, 16, 24, 32, 40, 48, 56,
- 64, 72, 80, 88, 96, 104, 112, 120,
- 128, 136, 144, 152, 160, 168, 176, 184,
- 192, 200, 208, 216, 224, 232, 240, 248
- ]> : vector<32xindex>
-
- %c17 = arith.constant 17: index
- %mask = vector.create_mask %c17 : vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
- gpu.return %ld : vector<32x4xf32>
- }
-}
-
-// -----
-gpu.module @test_kernel {
- // CHECK-LABEL: store_with_offsets_chunk
- // CHECK: [[cst:%.+]] = arith.constant dense<1.023000e+03> : vector<16x2xf32
- // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
- // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
- // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
- // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
- // CHECK-COUNT-4: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16x2xf32>, ui64, vector<16xindex>, vector<16xi1>
- gpu.func @store_with_offsets_chunk(%src: ui64) {
- %cst = arith.constant dense<[
- 0, 8, 16, 24, 32, 40, 48, 56,
- 64, 72, 80, 88, 96, 104, 112, 120,
- 128, 136, 144, 152, 160, 168, 176, 184,
- 192, 200, 208, 216, 224, 232, 240, 248
- ]> : vector<32xindex>
-
- %c17 = arith.constant 17: index
- %mask = vector.create_mask %c17 : vector<32xi1>
-
- %st_vec = arith.constant dense<1023.>: vector<32x4xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
- gpu.return
- }
-}
-
// -----
gpu.module @test_kernel {
// CHECK-LABEL: preserve_unit_dim_of_load_inst_data
@@ -654,8 +604,8 @@ gpu.module @test_kernel {
// CHECK: [[cst_0:%.+]] = arith.constant dense<true> : vector<1x1x16xi1>
// CHECK: [[cst_1:%.+]] = arith.constant dense<{{.*}}> : vector<1x1x16xindex>
// CHECK: [[cst_2:%.+]] = arith.constant dense<{{.*}}> : vector<1x1x16xindex>
- // CHECK: [[ld_0:%.+]] = xegpu.load [[arg0]][[[cst_1]]], [[cst_0]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
- // CHECK: [[ld_1:%.+]] = xegpu.load [[arg0]][[[cst_2]]], [[cst_0]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
+ // CHECK: [[ld_0:%.+]] = xegpu.load [[arg0]][[[cst_1]]], [[cst_0]] <{l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
+ // CHECK: [[ld_1:%.+]] = xegpu.load [[arg0]][[[cst_2]]], [[cst_0]] <{l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
// CHECK: [[ins_0:%.+]] = vector.insert_strided_slice [[ld_0]], [[cst]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x1x16xf32> into vector<1x1x32xf32>
// CHECK: [[ins_1:%.+]] = vector.insert_strided_slice [[ld_1]], [[ins_0]] offsets = [0, 0, 16], strides = [1, 1, 1] : vector<1x1x16xf32> into vector<1x1x32xf32>
gpu.func @preserve_unit_dim_of_load_inst_data(%src: ui64) -> vector<1x1x32xf32> {
@@ -667,7 +617,7 @@ gpu.module @test_kernel {
]]> : vector<1x1x32xindex>
%mask = arith.constant dense<true> : vector<1x1x32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [1, 1, 16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %ld = xegpu.load %src[%cst], %mask {layout = #xegpu.layout<inst_data = [1, 1, 16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
gpu.return %ld : vector<1x1x32xf32>
}
@@ -735,10 +685,10 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]]
]> : vector<1x1x32xindex>
%mask = arith.constant dense<true> : vector<1x1x32xi1>
- %a = xegpu.load %A[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
- %b = xegpu.load %B[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %a = xegpu.load %A[%cst], %mask {layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %b = xegpu.load %B[%cst], %mask {layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
%addf = arith.addf %a, %b : vector<1x1x32xf32>
- xegpu.store %addf, %C[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : vector<1x1x32xf32>, ui64, vector<1x1x32xindex>, vector<1x1x32xi1>
+ xegpu.store %addf, %C[%cst], %mask {layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : vector<1x1x32xf32>, ui64, vector<1x1x32xindex>, vector<1x1x32xi1>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
index f23ae46905652..9811a730e20f8 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
@@ -128,7 +128,7 @@ gpu.module @test {
//-----
// CHECK-LABEL: load_with_offsets
// CHECK-SAME: [[arg0:%.+]]: ui64
- // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.func @load_with_offsets(%src: ui64) -> vector<32xf32> {
%cst = arith.constant dense<[
0, 8, 16, 24, 32, 40, 48, 56,
@@ -139,7 +139,7 @@ gpu.module @test {
%c17 = arith.constant 17: index
%mask = vector.create_mask %c17: vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %ld = xegpu.load %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %ld : vector<32xf32>
}
@@ -147,7 +147,7 @@ gpu.module @test {
//-----
// CHECK-LABEL: store_with_offsets
// CHECK-SAME: [[arg0:%.+]]: ui64
- // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
+ // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
gpu.func @store_with_offsets(%src: ui64) {
%cst = arith.constant dense<[
0, 8, 16, 24, 32, 40, 48, 56,
@@ -160,59 +160,11 @@ gpu.module @test {
%mask = vector.create_mask %c17: vector<32xi1>
%st_vec = arith.constant dense<1023.0>: vector<32xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ xegpu.store %st_vec, %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
-//-----
- // CHECK-LABEL: load_with_offsets_chunk
- // CHECK-SAME: [[arg0:%.+]]: ui64
- // CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<32x4xf32>
- // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
- // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
- // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
- // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
- // CHECK-COUNT-4: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16x2xf32>
- gpu.func @load_with_offsets_chunk(%src: ui64) -> vector<32x4xf32> {
- %cst = arith.constant dense<[
- 0, 8, 16, 24, 32, 40, 48, 56,
- 64, 72, 80, 88, 96, 104, 112, 120,
- 128, 136, 144, 152, 160, 168, 176, 184,
- 192, 200, 208, 216, 224, 232, 240, 248
- ]> : vector<32xindex>
-
- %c17 = arith.constant 17: index
- %mask = vector.create_mask %c17: vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
- gpu.return %ld : vector<32x4xf32>
- }
-
-//-----
- // CHECK-LABEL: store_with_offsets_chunk
- // CHECK-SAME: [[arg0:%.+]]: ui64
- // CHECK: [[cst:%.+]] = arith.constant dense<1.023000e+03> : vector<16x2xf32
- // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
- // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
- // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
- // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
- // CHECK-COUNT-4: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16x2xf32>, ui64, vector<16xindex>, vector<16xi1>
- gpu.func @store_with_offsets_chunk(%src: ui64) {
- %cst = arith.constant dense<[
- 0, 8, 16, 24, 32, 40, 48, 56,
- 64, 72, 80, 88, 96, 104, 112, 120,
- 128, 136, 144, 152, 160, 168, 176, 184,
- 192, 200, 208, 216, 224, 232, 240, 248
- ]> : vector<32xindex>
-
- %c17 = arith.constant 17: index
- %mask = vector.create_mask %c17: vector<32xi1>
-
- %st_vec = arith.constant dense<1023.>: vector<32x4xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
- gpu.return
- }
-
//-----
// CHECK-LABEL: load_nd_store_nd
// CHECK-SAME: [[arg0:%.+]]: memref<256x318xf32>
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 9267f3bbbfbda..3b8626693cfb9 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -313,11 +313,11 @@ gpu.module @test_distribution {
gpu.func @load_gather(%src : memref<?xf16>) {
// CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<32x4xindex>
// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<32x4xi1>
- // CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}>
+ // CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{l1_hint = #xegpu.cache_hint<cached>}>
// CHECK-SAME: : memref<?xf16>, vector<32x4xindex>, vector<32x4xi1> -> vector<32x4xf16>
%offset = arith.constant dense<0> : vector<256x16xindex>
%mask = arith.constant dense<1> : vector<256x16xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>, l1_hint = #xegpu.cache_hint<cached>}
+ %load = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>, l1_hint = #xegpu.cache_hint<cached>}
: memref<?xf16>, vector<256x16xindex>, vector<256x16xi1> -> vector<256x16xf16>
gpu.return
}
@@ -328,31 +328,17 @@ gpu.module @test_distribution {
// CHECK: %[[VAL:.*]] = arith.constant dense<2.550000e+01> : vector<8xf16>
// CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<8xindex>
// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8xi1>
- // CHECK: xegpu.store %[[VAL]], %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<inst_data = [8]>}>
+ // CHECK: xegpu.store %[[VAL]], %[[ARG0]][%[[CST]]], %[[MASK]] <{l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<inst_data = [8]>}>
// CHECK-SAME: : vector<8xf16>, memref<256xf16>, vector<8xindex>, vector<8xi1>
%val = arith.constant dense<25.5> : vector<256xf16>
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
- xegpu.store %val, %dest[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>,
+ xegpu.store %val, %dest[%offset], %mask {layout = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>,
l1_hint = #xegpu.cache_hint<cached>}
: vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
gpu.return
}
- // CHECK-LABEL: @load_with_non_unit_chunk_size
- // CHECK-SAME: %[[ARG0:.*]]: memref<?xf16>
- gpu.func @load_with_non_unit_chunk_size(%src : memref<?xf16>) {
- // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<8xindex>
- // CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8xi1>
- // CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 4 : i64, l1_hint = #xegpu.cache_hint<cached>}>
- // CHECK-SAME: : memref<?xf16>, vector<8xindex>, vector<8xi1> -> vector<8x4xf16>
- %offset = arith.constant dense<0> : vector<256xindex>
- %mask = arith.constant dense<1> : vector<256xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 4, layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [8, 4]>, l1_hint = #xegpu.cache_hint<cached>}
- : memref<?xf16>, vector<256xindex>, vector<256xi1> -> vector<256x4xf16>
- gpu.return
- }
-
// CHECK-LABEL: distribute_load_matrix
// CHECK-SAME: [[arg0:%.+]]: memref<32768xi8, 3>
gpu.func @distribute_load_matrix(%arg0: memref<32768xi8, 3>) {
@@ -760,9 +746,9 @@ gpu.module @test_distribution {
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
- // CHECK: %[[LOAD:.*]] = xegpu.load {{.*}} <{chunk_size = 1 : i64, layout = #xegpu.slice<#xegpu.layout<inst_data = [8, 16]>, dims = [0]>}>
+ // CHECK: %[[LOAD:.*]] = xegpu.load {{.*}} <{layout = #xegpu.slice<#xegpu.layout<inst_data = [8, 16]>, dims = [0]>}>
// CHECK-SAME: memref<4096xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- %3 = xegpu.load %2[%offset], %mask {chunk_size = 1, layout = #xegpu.slice<#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>, dims = [0]> } : memref<4096xf32>, vector<256xindex>, vector<256xi1> -> vector<256xf32>
+ %3 = xegpu.load %2[%offset], %mask {layout = #xegpu.slice<#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>, dims = [0]> } : memref<4096xf32>, vector<256xindex>, vector<256xi1> -> vector<256xf32>
// CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[LOAD]] : vector<32xf32> to vector<32x32xf32>
%4 = vector.broadcast %3 : vector<256xf32> to vector<256x256xf32>
@@ -779,7 +765,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.000000e+00> : vector<1x32xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : vector<1x1x32xindex>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<1x1x32xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0:.*]][%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0:.*]][%[[CST_0]]], %[[CST_1]] : memref<?xf32>, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<0.000000e+00> : vector<1x32xf32>
// CHECK-DAG: %[[LOCAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_2]] [1] : vector<1x1x32xf32> to vector<1x32xf32>
// CHECK-DAG: %[[CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<1x32xf32> to vector<1x1x32xf32>
@@ -796,7 +782,7 @@ gpu.module @test_distribution {
%cst_3 = arith.constant dense<1.0> : vector<1x32xf32>
%offset = arith.constant dense<0> : vector<1x32x32xindex>
%mask = arith.constant dense<true> : vector<1x32x32xi1>
- %14 = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [1, 32, 1], sg_data = [1, 1, 32]>} : memref<?xf32>, vector<1x32x32xindex>, vector<1x32x32xi1> -> vector<1x32x32xf32>
+ %14 = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [1, 32, 1], sg_data = [1, 1, 32]>} : memref<?xf32>, vector<1x32x32xindex>, vector<1x32x32xi1> -> vector<1x32x32xf32>
%15 = vector.multi_reduction <add>, %14, %cst_3 [1] : vector<1x32x32xf32> to vector<1x32xf32>
%anchor = xegpu.convert_layout %15
<{
@@ -852,7 +838,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<1x1xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : vector<1x1x32x32xindex>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<1x1x32x32xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<1x1x32x32xindex>, vector<1x1x32x32xi1> -> vector<1x1x32x32xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST_0]]], %[[CST_1]] : memref<?xf32>, vector<1x1x32x32xindex>, vector<1x1x32x32xi1> -> vector<1x1x32x32xf32>
// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<0.000000e+00> : vector<1x1xf32>
// CHECK-DAG: %[[LOCAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_2]] [2, 3] : vector<1x1x32x32xf32> to vector<1x1xf32>
// CHECK-DAG: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<1x1xf32> to vector<1x1x1x1xf32>
@@ -869,7 +855,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<0.0> : vector<2x2xf32>
%offset = arith.constant dense<0> : vector<2x2x128x128xindex>
%mask = arith.constant dense<true> : vector<2x2x128x128xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [1, 1, 32, 32]>} : memref<?xf32>, vector<2x2x128x128xindex>, vector<2x2x128x128xi1> -> vector<2x2x128x128xf32>
+ %load = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [1, 1, 32, 32]>} : memref<?xf32>, vector<2x2x128x128xindex>, vector<2x2x128x128xi1> -> vector<2x2x128x128xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [2, 3] : vector<2x2x128x128xf32> to vector<2x2xf32>
%anchor = xegpu.convert_layout %reduce
<{
@@ -884,7 +870,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<16x16xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : vector<16x16x32x32xindex>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<16x16x32x32xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<16x16x32x32xindex>, vector<16x16x32x32xi1> -> vector<16x16x32x32xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST_0]]], %[[CST_1]] : memref<?xf32>, vector<16x16x32x32xindex>, vector<16x16x32x32xi1> -> vector<16x16x32x32xf32>
// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<0.000000e+00> : vector<16x16xf32>
// CHECK-DAG: %[[LOCAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_2]] [2, 3] : vector<16x16x32x32xf32> to vector<16x16xf32>
// CHECK-DAG: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<16x16xf32> to vector<16x16x1x1xf32>
@@ -901,7 +887,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<0.0> : vector<32x32xf32>
%offset = arith.constant dense<0> : vector<32x32x128x128xindex>
%mask = arith.constant dense<true> : vector<32x32x128x128xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [16, 16, 32, 32]>} : memref<?xf32>, vector<32x32x128x128xindex>, vector<32x32x128x128xi1> -> vector<32x32x128x128xf32>
+ %load = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [16, 16, 32, 32]>} : memref<?xf32>, vector<32x32x128x128xindex>, vector<32x32x128x128xi1> -> vector<32x32x128x128xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [2, 3] : vector<32x32x128x128xf32> to vector<32x32xf32>
%anchor = xegpu.convert_layout %reduce
<{
@@ -997,7 +983,7 @@ gpu.module @test_distribution {
gpu.func @convert_layout_3D(%arg0: memref<?xf32>) {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : vector<1x32x16xindex>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<true> : vector<1x32x16xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST]]], %[[CST_0]] <{chunk_size = 1 : i64, layout = #xegpu.layout<inst_data = [1, 16, 16]>}> : memref<?xf32>, vector<1x32x16xindex>, vector<1x32x16xi1> -> vector<1x32x16xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST]]], %[[CST_0]] <{layout = #xegpu.layout<inst_data = [1, 16, 16]>}> : memref<?xf32>, vector<1x32x16xindex>, vector<1x32x16xi1> -> vector<1x32x16xf32>
// CHECK-DAG: %[[ALLOCA:.*]] = memref.alloca() : memref<1048576xi8, 3>
// CHECK-DAG: %[[MDESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<1048576xi8, 3> -> !xegpu.mem_desc<8x128x256xf32>
// CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index
@@ -1026,7 +1012,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MDESC]][%[[LOAD_OFF_Z]], %[[LOAD_OFF_Y]], %[[LOAD_OFF_X]]] <{layout = #xegpu.layout<inst_data = [1, 16, 16]>}> : !xegpu.mem_desc<8x128x256xf32>, index, index, index -> vector<1x16x32xf32>
%offset = arith.constant dense<0> : vector<8x128x256xindex>
%mask = arith.constant dense<true> : vector<8x128x256xi1>
- %1 = xegpu.load %arg0[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>} : memref<?xf32>, vector<8x128x256xindex>, vector<8x128x256xi1> -> vector<8x128x256xf32>
+ %1 = xegpu.load %arg0[%offset], %mask {layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>} : memref<?xf32>, vector<8x128x256xindex>, vector<8x128x256xi1> -> vector<8x128x256xf32>
%2 = xegpu.convert_layout %1 <{input_layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>,
target_layout = #xegpu.layout<sg_layout = [8, 8, 8], sg_data = [1, 16, 32], inst_data = [1, 16, 16]>}> : vector<8x128x256xf32>
%anchor = xegpu.convert_layout %2
>From 30e73ea6d24c64f8c712fef0889380a058ffe7c2 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Tue, 25 Aug 2026 22:55:33 +0000
Subject: [PATCH 2/2] [mlir][XeGPU] Distribute coalesced gather/scatter via
lane_data
Rework SgToLaneLoadGather / SgToLaneStoreScatter to lower a coalesced
gather/scatter to the chunked memory access purely from the layout's
lane_data, without relying on a chunk_size attribute.
When lane_data[FCD] = D > 1 covers the lane's entire per-lane fragment (one
round: lane_layout[FCD] * D == FCD extent, so the distributed per-lane vector
has exactly D elements), the lane owns a single contiguous run
{base, base+1, ..., base+D-1}. Distribution takes the first offset/mask of the
group (scalar base + scalar mask) and emits a value vector<D>; the chunk size
is implied by the value type, which the XeVM lowering already consumes. The
dropped per-lane offset computations become dead code removed during lowering.
The round-robin case (lane_data[FCD] == 1, multiple strided rounds) is excluded
by the lane_data > 1 guard.
This builds on the chunk_size removal (PR #205122); the distribution is now
driven entirely by lane_data with no chunk_size attribute anywhere.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 48 +++++++++++++++++
.../Dialect/XeGPU/sg-to-lane-distribute.mlir | 54 +++++++++++++++++++
2 files changed, 102 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 45e0bac55c275..0d437e0a29993 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -570,6 +570,35 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
castValueTo(rewriter, cast<TypedValue<VectorType>>(distMask), maskTy1D);
Value distSource = adaptor.getSource();
+
+ // Coalesced case: the layout assigns a genuine contiguous chunk on the FCD
+ // via `lane_data[FCD] = D > 1`, and that chunk is the lane's *entire*
+ // per-lane fragment (one round: `lane_layout[FCD] * D == FCD extent`, so
+ // the distributed per-lane vector has exactly D elements). Only then does
+ // the lane own a single run of D contiguous elements
+ // `{base, base+1, ..., base+D-1}`, which the XeVM lowering expects as a
+ // chunked access: a scalar base offset + scalar mask + a `vector<D>` value
+ // (the chunk size is implied by the value type). Take the first offset /
+ // mask of the contiguous group; the remaining offset computations become
+ // dead code and are removed during lowering.
+ //
+ // This must NOT fire for the round-robin case (`lane_data[FCD] = 1`,
+ // multiple rounds, e.g. a reduction source where lane l owns
+ // `{l, l+SG, l+2*SG, ...}`): there the per-lane elements are strided, not
+ // contiguous, so a chunked access would read the wrong elements. That case
+ // has `lane_data[FCD] == 1`, so the guard below excludes it.
+ int64_t laneElems = distResultTy1D.getNumElements();
+ int64_t innerLaneData = 1;
+ if (auto laneDataArr = layout.getEffectiveLaneDataAsInt();
+ !laneDataArr.empty())
+ innerLaneData = laneDataArr.back();
+ if (innerLaneData > 1 && laneElems == innerLaneData) {
+ distOffsets = vector::ExtractOp::create(
+ rewriter, op.getLoc(), distOffsets, ArrayRef<int64_t>{0});
+ distMask = vector::ExtractOp::create(rewriter, op.getLoc(), distMask,
+ ArrayRef<int64_t>{0});
+ }
+
auto newOp = xegpu::LoadGatherOp::create(
rewriter, op.getLoc(), distResultTy1D, distSource, distOffsets,
distMask, op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(),
@@ -1093,6 +1122,25 @@ struct SgToLaneStoreScatter
castValueTo(rewriter, cast<TypedValue<VectorType>>(distMask), maskTy1D);
Value distDest = adaptor.getDest();
+
+ // Coalesced case (mirrors SgToLaneLoadGather): a genuine contiguous chunk
+ // on the FCD, `lane_data[FCD] = D > 1` covering the lane's entire per-lane
+ // fragment (one round), stores as a chunked access: scalar base offset +
+ // scalar mask + a `vector<D>` value. Take the first offset / mask of the
+ // contiguous group; the dropped offset computations are DCE'd in lowering.
+ // Excludes the round-robin case (`lane_data[FCD] == 1`).
+ int64_t laneElems = distValueTy1D.getNumElements();
+ int64_t innerLaneData = 1;
+ if (auto laneDataArr = layout.getEffectiveLaneDataAsInt();
+ !laneDataArr.empty())
+ innerLaneData = laneDataArr.back();
+ if (innerLaneData > 1 && laneElems == innerLaneData) {
+ distOffsets = vector::ExtractOp::create(
+ rewriter, op.getLoc(), distOffsets, ArrayRef<int64_t>{0});
+ distMask = vector::ExtractOp::create(rewriter, op.getLoc(), distMask,
+ ArrayRef<int64_t>{0});
+ }
+
xegpu::StoreScatterOp::create(rewriter, op.getLoc(), distValue, distDest,
distOffsets, distMask, op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr(),
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
index ad046a7e432ed..25d89dd492306 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
@@ -512,3 +512,57 @@ gpu.module @xevm_module {
gpu.return
}
}
+
+// -----
+// Coalesced gather/scatter: the load/store layout has lane_data[FCD] = 2, i.e.
+// each lane owns 2 *contiguous* elements (one round: lane_layout[FCD] * 2 == 32
+// == FCD extent). Distribution must emit the chunked form the XeVM lowering
+// accepts: a scalar base offset + scalar mask + a value vector<2xf32> (the
+// chunk size is implied by the value type), taking element 0 of the per-lane
+// offsets/mask as the base -- NOT a 2-wide offsets/mask vector. This is driven
+// entirely by lane_data; the dropped offset lanes are DCE'd during lowering.
+gpu.module @xevm_module {
+ // CHECK-LABEL: gpu.func @coalesced_load_store
+ // CHECK: %[[LD:.*]] = xegpu.load %{{.*}}[%[[BASE:.*]]], %{{.*}} : i64, index, i1 -> vector<2xf32>
+ // CHECK: %[[MUL:.*]] = arith.mulf %[[LD]], %{{.*}} : vector<2xf32>
+ // CHECK: xegpu.store %[[MUL]], %{{.*}}[%[[BASE]]], %{{.*}} : vector<2xf32>, i64, index, i1
+ gpu.func @coalesced_load_store(%src: i64, %dst: i64) {
+ %step = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %src[%step], %mask <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %c = arith.constant dense<2.0> : vector<32xf32>
+ %p = arith.mulf %v, %c {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} : vector<32xf32>
+ xegpu.store %p, %dst[%step], %mask <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}>
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ gpu.return
+ }
+}
+
+// -----
+// Coalesced gather/scatter, 2-D: the FCD (innermost dim) carries the contiguous
+// chunk via lane_data = [1, 2] with a unit leading dim, so lane_layout[FCD] * 2
+// == 32 == FCD extent and each lane owns 2 contiguous elements. This must lower
+// to the SAME chunked form as the 1-D case -- scalar base offset + scalar mask
+// + a value vector<2xf32> (the leading unit dim is folded away) -- proving the
+// coalescing is driven by lane_data[FCD], independent of the value rank.
+gpu.module @xevm_module {
+ // CHECK-LABEL: gpu.func @coalesced_load_store_2d
+ // CHECK: %[[LD:.*]] = xegpu.load %{{.*}}[%[[BASE:.*]]], %{{.*}} : i64, index, i1 -> vector<2xf32>
+ // CHECK: %[[CAST:.*]] = vector.shape_cast %[[LD]] : vector<2xf32> to vector<1x2xf32>
+ // CHECK: %[[MUL:.*]] = arith.mulf %[[CAST]], %{{.*}} : vector<1x2xf32>
+ // CHECK: %[[CAST2:.*]] = vector.shape_cast %[[MUL]] : vector<1x2xf32> to vector<2xf32>
+ // CHECK: xegpu.store %[[CAST2]], %{{.*}}[%[[BASE]]], %{{.*}} : vector<2xf32>, i64, index, i1
+ gpu.func @coalesced_load_store_2d(%src: i64, %dst: i64) {
+ %step = vector.step : vector<32xindex>
+ %step2d = vector.shape_cast %step : vector<32xindex> to vector<1x32xindex>
+ %mask = arith.constant dense<true> : vector<1x32xi1>
+ %v = xegpu.load %src[%step2d], %mask <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}>
+ : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
+ %c = arith.constant dense<2.0> : vector<1x32xf32>
+ %p = arith.mulf %v, %c {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>} : vector<1x32xf32>
+ xegpu.store %p, %dst[%step2d], %mask <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}>
+ : vector<1x32xf32>, i64, vector<1x32xindex>, vector<1x32xi1>
+ gpu.return
+ }
+}
More information about the Mlir-commits
mailing list