[Mlir-commits] [mlir] [mlir][xegpu] Remove chunk_size attribute; infer from types (PR #205122)
Md Abdullah Shahneous Bari
llvmlistbot at llvm.org
Fri Sep 11 18:31:59 PDT 2026
https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/205122
>From 401043341529a6c63f10c849adfe372235c51707 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Fri, 11 Sep 2026 21:00:34 +0000
Subject: [PATCH] [mlir][xegpu] Remove the chunk_size attribute from load/store
gather/scatter
Removes the `chunk_size` attribute from `xegpu.load` (LoadGatherOp) and
`xegpu.store` (StoreScatterOp), along with every code path that
special-cased `chunk_size > 1`. The per-lane access width is implied by
the value type, so the attribute carried no information the types did not
already have.
The `contiguity` attribute is the thing that says how much of an access is
contiguous. It stays, and it now feeds the per-lane chunk cap that
`chunk_size` used to feed.
One shape rule, everywhere: the mask, the offsets, and the value all have
the same shape, at every level. A scalar and a size-1 vector both stand
for a single element, so they stay interchangeable, but a scalar mask can
no longer carry a wider value. That form was the old `chunk_size` spelling
of a contiguous block access, and keeping it would have left two
contradictory ways to write the same lane-level access.
That leaves the congruent form as the only spelling, so XeGPUToXeVM has to
lower it. It previously required a scalar offset and a scalar mask. It now
coalesces: when the offsets are one ascending run and the mask is uniform,
the base offset and the mask bit come from element 0 and one block
llvm.load/llvm.store of the whole `vector<D>` is emitted. Both facts are
proven on the original operands, whose defining ops are still intact,
while the values come from the converted operands, so no extra casts are
needed. Two offset forms are matched, both of which the distribution
emits: `vector.from_elements %b, %b + 1, ...` and `vector.step`. A uniform
mask is a splat constant, a broadcast of a scalar, or a
`vector.from_elements` of one repeated value. When either proof fails the
access stays a real gather and the pattern fails to match, as before.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 101 ++++++++--------
.../XeGPU/Transforms/XeGPULayoutImpl.h | 10 +-
.../VectorToXeGPU/VectorToXeGPU.cpp | 4 -
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 44 +++++++
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 96 +++++++---------
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 33 +++---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 24 ++--
.../Transforms/XeGPUSgToLaneDistribute.cpp | 40 ++-----
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 108 +++---------------
.../Transforms/XeGPUWgToSgDistribute.cpp | 16 +--
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 17 +--
.../XeGPUToXeVM/failed_conversion.mlir | 16 +++
.../loadstore_gather_coalesce.mlir | 102 +++++++++++++++++
mlir/test/Dialect/XeGPU/invalid.mlir | 92 ++++++++-------
mlir/test/Dialect/XeGPU/ops.mlir | 49 +++-----
.../propagate-layout-subgroup-invalid.mlir | 2 +-
.../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 | 70 ++++++------
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 +++----
23 files changed, 459 insertions(+), 589 deletions(-)
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/loadstore_gather_coalesce.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 2821b01e880d5..d889fd9166ebe 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -706,16 +706,12 @@ def XeGPU_PrefetchOp : XeGPU_Op<"prefetch", [AnchorLayoutInterface]> {
def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayoutInterface]> {
let summary = "load a set of scattered data points from memory.";
- 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.
+ let description = [{ It (aka. load) load data per each lane.
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.
+ At every level the mask, offsets, and result may be nD vectors, and all three
+ must have the same shape.
This operation serves as an anchor through which users assign a layout attribute
to govern computation distribution.
@@ -726,17 +722,19 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
1D memref or pointer (ui64, ui32, i64 or i32).
- `offsets`: represents offsets from source.
- offsets is a vector of `index` type and vector length is either the subgroup size
- or 1 at lane level. scalar offset is also valid for lane level.
+ offsets is a vector of `index` type with the same shape as the result.
+ A scalar offset is also valid: one lane then reads only one element.
- `mask`: is a vector of `i1` type, which is used to mask out the memory access.
- 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.
+ mask has the same shape as the offsets and the result. Each element applies
+ to one lane. A scalar mask is also valid, and gates that one element.
- `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache.
+ - `contiguity`: [optional] the number of neighbouring elements each lane may
+ access in one instruction, as established by the contiguity analysis. It
+ must be at least 2 and must divide the innermost offsets dimension.
+
- `layout`: [optional] Describes the expected layout of the `tensor_desc` operand or the result
of load. Only valid at workgroup and subgroup levels.
@@ -750,28 +748,36 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
%a = memref.alloc() : memref<1024xf32>
%offsets = vector.step : vector<16xindex>
%mask = vector.constant_mask [16]: vector<16xi1>
- %val = xegpu.load %a[%offsets], %mask {l1_hint = #xegpu.cache_hint<cached>,
+ %val = xegpu.load %a[%offsets], %mask <{l1_hint = #xegpu.cache_hint<cached>,
l2_hint = #xegpu.cache_hint<cached>,
l3_hint = #xegpu.cache_hint<cached>,
- layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
- : memref<1024xf32>, vector<16xi1>, vector<16xindex> -> vector<16xf32>
+ layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}>
+ : memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
```
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. The mask, offsets, and result all
+ have the same shape; here each lane reads 8 elements.
```mlir
- %2 = xegpu.load %1[%2], %0 <{l1_hint = #xegpu.cache_hint<cached>,
+ %val = xegpu.load %src[%offsets], %mask <{l1_hint = #xegpu.cache_hint<cached>,
l2_hint = #xegpu.cache_hint<uncached>,
l3_hint = #xegpu.cache_hint<uncached>}>
- : memref<128xf32>, vector<1xindex>, vector<1xi1> -> vector<8xf32>
+ : memref<128xf32>, vector<8xindex>, vector<8xi1> -> vector<8xf32>
+ ```
+
+ Example 3 (subgroup level, contiguous runs):
+ `contiguity` says each lane may take 4 neighbouring elements in one access.
+ The shapes stay the same.
+ ```mlir
+ %val = xegpu.load %src[%offsets], %mask <{contiguity = 4 : i64}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
```
}];
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 +823,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,
@@ -836,15 +840,10 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorLayoutInterface]> {
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.
+ [{ It (aka. store) stores data to scattered memory locations.
- 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.
+ At every level the mask, offsets, and value may be nD vectors, and all three
+ must have the same shape.
This operation serves as an anchor through which users assign a layout attribute
to govern computation distribution.
@@ -857,17 +856,19 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
1D memref or pointer (ui64, ui32, i64 or i32).
- `offsets`: represents offsets from dest.
- offsets is a vector of `index` type and vector length is either the subgroup size
- or 1 at lane level. scalar offset is also valid for lane level.
+ offsets is a vector of `index` type with the same shape as the value.
+ A scalar offset is also valid: one lane then writes only one element.
- `mask`: is a vector of `i1` type, which is used to mask out the memory access.
- 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.
+ mask has the same shape as the offsets and the value. Each element applies
+ to one lane. A scalar mask is also valid, and gates that one element.
- `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache.
+ - `contiguity`: [optional] the number of neighbouring elements each lane may
+ access in one instruction, as established by the contiguity analysis. It
+ must be at least 2 and must divide the innermost offsets dimension.
+
- `layout`: [optional] Describes the expected layout of the `tensor_desc` operand or the value
to be stored. Only valid at workgroup and subgroup levels.
@@ -880,21 +881,29 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
%val = arith.constant dense<0.0> : vector<16xf32>
%offsets = vector.step : vector<16xindex>
%mask = vector.constant_mask [16]: vector<16xi1>
- xegpu.store %val, %a[%offsets], %mask {l1_hint = #xegpu.cache_hint<cached>,
+ xegpu.store %val, %a[%offsets], %mask <{l1_hint = #xegpu.cache_hint<cached>,
l2_hint = #xegpu.cache_hint<cached>,
l3_hint = #xegpu.cache_hint<cached>,
- layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
- : vector<16xf32>, memref<1024xf32>, vector<16xi1>, vector<16xindex>
+ layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}>
+ : vector<16xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
```
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. The mask, offsets, and value all
+ have the same shape; here each lane writes 8 elements.
```mlir
- xegpu.store %0, %1[%2], %3 <{l1_hint = #xegpu.cache_hint<uncached>,
+ xegpu.store %val, %dest[%offsets], %mask <{l1_hint = #xegpu.cache_hint<uncached>,
l2_hint = #xegpu.cache_hint<write_back>,
l3_hint = #xegpu.cache_hint<write_through>}>
- : vector<8xf32>, memref<256xf32>, vector<1xindex>, vector<1xi1>
+ : vector<8xf32>, memref<256xf32>, vector<8xindex>, vector<8xi1>
+ ```
+
+ Example 3 (Subgroup level, contiguous runs):
+ `contiguity` says each lane may take 4 neighbouring elements in one access.
+ The shapes stay the same.
+ ```mlir
+ xegpu.store %val, %dest[%offsets], %mask <{contiguity = 4 : i64}>
+ : vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
```
}];
@@ -902,7 +911,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 +956,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 9d0ee049ad936..e23bce46b16ee 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -170,12 +170,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
@@ -281,6 +275,8 @@ DistributeLayoutAttr setupInsertStridedSliceResultLayout(
DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for a load gather operation.
+/// `contigChunkSize` is the per-lane contiguous run the offsets allow, from
+/// the op's `contiguity` attribute (1 when it is absent).
DistributeLayoutAttr setupLoadGatherAnchorLayout(
LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize,
DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
@@ -291,6 +287,8 @@ DistributeLayoutAttr setupLoadMatrixAnchorLayout(
DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for a store scatter operation.
+/// `contigChunkSize` is the per-lane contiguous run the offsets allow, from
+/// the op's `contiguity` attribute (1 when it is absent).
/// `numSg` is only used for Subgroup-kind layouts.
DistributeLayoutAttr setupStoreScatterAnchorLayout(LayoutKind layoutKind,
VectorType vectorTy,
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 3a56d21240109..500e5d055e4ad 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -462,7 +462,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{},
@@ -497,7 +496,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{},
@@ -742,7 +740,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{},
@@ -777,7 +774,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/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index d782665951708..8659c7bf530c9 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -705,6 +705,23 @@ static Value addOffsetToBaseAddr(ConversionPatternRewriter &rewriter,
return newAddr;
}
+// Returns true when every element of `mask` carries the same bit, so gating a
+// whole contiguous block on element 0 is equivalent. Splat constants,
+// broadcasts of a scalar, and a `vector.from_elements` of one repeated value
+// qualify.
+static bool isUniformMask(Value mask) {
+ if (!isa<VectorType>(mask.getType()))
+ return true;
+ DenseElementsAttr splat;
+ if (matchPattern(mask, m_Constant(&splat)) && splat.isSplat())
+ return true;
+ if (auto broadcast = mask.getDefiningOp<vector::BroadcastOp>())
+ return !isa<VectorType>(broadcast.getSource().getType());
+ if (auto fromElements = mask.getDefiningOp<vector::FromElementsOp>())
+ return llvm::all_equal(fromElements.getElements());
+ return false;
+}
+
template <typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
OpType, xegpu::LoadGatherOp, xegpu::StoreScatterOp>::value>>
@@ -770,6 +787,33 @@ class LoadStoreToXeVMPattern : public OpConversionPattern<OpType> {
basePtrI64);
}
Value mask = adaptor.getMask();
+
+ // Coalesce a chunked access into one contiguous block access.
+ //
+ // Distribution gives every element its own offset and mask bit. A lane
+ // whose `lane_data` is wider than one element therefore arrives here with
+ // `vector<D>` offsets and mask beside a `vector<D>` value, and those D
+ // offsets are one ascending run: `lane_data` is by definition the number of
+ // neighbouring elements the lane takes, and `inst_data` blocks the op so
+ // that no lane is given more. A single element converts to a scalar offset
+ // instead, so a surviving vector offset always means a contiguous run.
+ //
+ // A block access needs one base offset and one mask bit, so take both from
+ // element 0. The values come from the converted operands, so no casts are
+ // needed. The mask is still checked: a tail mask can leave a lane with a
+ // partial run, and one bit cannot stand for D different bits. That case
+ // fails to match below, as before.
+ auto origOffsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ if (isa<VectorType>(offset.getType()) && origOffsetsTy && valOrResVecTy &&
+ origOffsetsTy.getNumElements() == valOrResVecTy.getNumElements() &&
+ isUniformMask(op.getMask())) {
+ offset = vector::ExtractOp::create(rewriter, loc, offset,
+ ArrayRef<int64_t>{0});
+ if (isa<VectorType>(mask.getType()))
+ mask = vector::ExtractOp::create(rewriter, loc, mask,
+ ArrayRef<int64_t>{0});
+ }
+
if (dyn_cast<VectorType>(offset.getType())) {
// Offset needs be scalar. Single element vector is converted to scalar
// by type converter.
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 34b667cfa1f7c..2474defc3453c 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -66,53 +66,36 @@ 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);
- 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.";
-
+ // One offset and one mask bit per element, so the value has the same shape as
+ // the mask and the offsets. A scalar and a size-1 vector both stand for a
+ // single element, so they are interchangeable.
+ int64_t maskSize = maskVecTy ? maskVecTy.getNumElements() : 1;
+ if (valueTy.getNumElements() != maskSize ||
+ (maskVecTy && getShapeOf(valueTy) != getShapeOf(maskTy)))
+ return emitError() << "Value shape must match mask shape.";
return success();
}
@@ -607,7 +590,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()))
@@ -622,14 +604,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();
@@ -638,15 +620,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) {
@@ -656,8 +638,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);
}
//===----------------------------------------------------------------------===//
@@ -677,7 +659,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()))
@@ -692,14 +673,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) {
@@ -710,15 +690,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());
@@ -726,8 +708,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 f696020b2de12..56aaf8883a821 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1023,17 +1023,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
@@ -1870,13 +1859,19 @@ 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.
+///
+/// A consumer layout that carries lane_layout and lane_data is required.
+/// `maxChunkSize` is not read yet: the consumer's lane_data already fixes the
+/// per-lane chunk.
///
-/// For Lane layout, lane_layout/lane_data are taken from the consumer when
-/// present; otherwise derived from the same default.
+/// TODO: derive lane_layout/lane_data here when the consumer has none, capped
+/// by `maxChunkSize`, the way `setupGenericStoreAnchorLayout` does via
+/// `computeScatterIOLaneLayoutAndData`. That path is missing today, so the
+/// assert below stands in for it.
static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize,
@@ -1920,6 +1915,8 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
ArrayRef<int64_t> resShape = resVecTy.getShape();
auto context = resVecTy.getContext();
+ // The per-lane chunk is bounded by what the offsets prove contiguous
+ // (`contigChunkSize`) and by what one lane can access in one instruction.
const auto *uArchInstruction = dyn_cast<xegpu::uArch::LoadGatherInstruction>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
int maxChunkSize =
@@ -2000,6 +1997,8 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
ArrayRef<int64_t> srcShape = srcVecTy.getShape();
auto context = srcVecTy.getContext();
+ // The per-lane chunk is bounded by what the offsets prove contiguous
+ // (`contigChunkSize`) and by what one lane can access in one instruction.
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::StoreScatterInstruction>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index ca1c6501a4856..63916762befef 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1241,7 +1241,9 @@ void LayoutInfoPropagation::visitLoadGatherOp(
if (!uArch)
return;
VectorType resVecTy = load.getValueType();
- int chunkSize = load.getChunkSize().value_or(1);
+ // The contiguity analysis reports how many neighbouring elements one lane may
+ // take in a single access. Absent, only one element per lane is safe.
+ int contigChunkSize = load.getContiguity().value_or(1);
LayoutInfo resLayoutInfo = results[0]->getValue();
if (!resLayoutInfo.isAssigned())
@@ -1275,14 +1277,12 @@ void LayoutInfoPropagation::visitLoadGatherOp(
return;
}
requiredAnchorLayoutAttr = xegpu::setupLoadGatherAnchorLayout(
- layoutKind, resVecTy, chunkSize, consumerLayoutAttr, uArch);
+ layoutKind, resVecTy, contigChunkSize, 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.
@@ -1306,7 +1306,9 @@ void LayoutInfoPropagation::visitStoreScatterOp(
if (!uArch)
return;
VectorType srcVecTy = storeScatter.getValueType();
- int chunkSize = storeScatter.getChunkSize().value_or(1);
+ // The contiguity analysis reports how many neighbouring elements one lane may
+ // write in a single access. Absent, only one element per lane is safe.
+ int contigChunkSize = storeScatter.getContiguity().value_or(1);
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
@@ -1338,7 +1340,7 @@ void LayoutInfoPropagation::visitStoreScatterOp(
if (failed(numSgOrErr))
return;
requiredAnchorLayoutAttr = xegpu::setupStoreScatterAnchorLayout(
- layoutKind, srcVecTy, chunkSize, numSgOrErr.value_or(0), uArch);
+ layoutKind, srcVecTy, contigChunkSize, numSgOrErr.value_or(0), uArch);
if (!requiredAnchorLayoutAttr) {
markFailure(storeScatter,
"Failed to determine required layout for store scatter.");
@@ -1348,10 +1350,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 812233047fb88..0c9f86afb880f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -499,7 +499,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>
@@ -511,15 +511,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>
@@ -545,8 +537,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),
@@ -583,8 +574,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)
@@ -1026,7 +1017,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>
@@ -1038,15 +1029,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>
@@ -1073,8 +1056,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; }))
@@ -1114,9 +1096,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 1854c639c4797..aaee71bb86182 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -568,7 +568,6 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
return failure();
SmallVector<int64_t> targetMaskShape(*targetShape);
- int64_t chunkSize = op.getChunkSize().value_or(1);
// Unroll mask and offsets with correct shape
VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
@@ -576,50 +575,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)
@@ -628,8 +592,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);
@@ -658,55 +621,19 @@ struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
if (!targetShape)
return failure();
- int64_t chunkSize = op.getChunkSize().value_or(1);
-
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);
@@ -720,7 +647,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 8169d2a282f38..9e2a0401900e0 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 bd2fe1e7f1a6b..9b05fa8df38b7 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -247,22 +247,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/Conversion/XeGPUToXeVM/failed_conversion.mlir b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
index dfaa9521786a9..65ce4f90f6500 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -31,6 +31,22 @@ gpu.module @test_kernel {
// -----
+// Vector offsets are coalesced into one block access, but the mask must be
+// uniform: one block access can only be gated on a single bit.
+
+gpu.module @test_kernel {
+ gpu.func @load_gather_non_uniform_mask(%src: i64, %base: index, %mask: vector<2xi1>) -> vector<2xf32> {
+ %c1 = arith.constant 1 : index
+ %o1 = arith.addi %base, %c1 : index
+ %offsets = vector.from_elements %base, %o1 : vector<2xindex>
+ // expected-error at +1 {{failed to legalize operation 'xegpu.load' that was explicitly marked illegal}}
+ %0 = xegpu.load %src[%offsets], %mask : i64, vector<2xindex>, vector<2xi1> -> vector<2xf32>
+ gpu.return %0 : vector<2xf32>
+ }
+}
+
+// -----
+
// Verify that xegpu.lane_shuffle of a sub-byte element type is rejected: the
// shuffle redistributes whole bytes between the lanes, so fp4 fragments cannot
// be shuffled.
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_gather_coalesce.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_gather_coalesce.mlir
new file mode 100644
index 0000000000000..45e57dfd11d8a
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_gather_coalesce.mlir
@@ -0,0 +1,102 @@
+// RUN: mlir-opt %s --split-input-file -convert-xegpu-to-xevm -canonicalize | FileCheck %s
+
+// Lane-level gather/scatter carry one offset and one mask bit per element. A
+// lane whose `lane_data` is wider than one element gets a contiguous run, so
+// vector offsets here are coalesced into a single block load/store based at
+// element 0. The mask must still be uniform: one block access is gated on one
+// bit.
+
+gpu.module @test {
+// CHECK-LABEL: @load_gather_contiguous_run
+// CHECK-SAME: %[[ARG0:.*]]: i64, %[[ARG1:.*]]: index, %[[ARG2:.*]]: i1
+gpu.func @load_gather_contiguous_run(%src: i64, %base: index, %pred: i1) -> vector<4xf32> {
+ %mask = vector.broadcast %pred : i1 to vector<4xi1>
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %o1 = arith.addi %base, %c1 : index
+ %o2 = arith.addi %base, %c2 : index
+ %o3 = arith.addi %base, %c3 : index
+ %offsets = vector.from_elements %base, %o1, %o2, %o3 : vector<4xindex>
+ // The run base is element 0 of the offsets, and the whole block is gated on
+ // one mask bit.
+ // CHECK: %[[C4:.*]] = arith.constant 4 : i64
+ // CHECK: %[[OFFSETS:.*]] = vector.from_elements %[[ARG1]],
+ // CHECK: %[[I64:.*]] = vector.bitcast %[[OFFSETS]] : vector<4xindex> to vector<4xi64>
+ // CHECK: %[[BASE:.*]] = vector.extract %[[I64]][0] : i64 from vector<4xi64>
+ // CHECK: %[[BYTES:.*]] = arith.muli %[[BASE]], %[[C4]] : i64
+ // CHECK: %[[ADDR:.*]] = arith.addi %[[ARG0]], %[[BYTES]] : i64
+ // CHECK: %[[PTR:.*]] = llvm.inttoptr %[[ADDR]] : i64 to !llvm.ptr<1>
+ // CHECK: scf.if %[[ARG2]] -> (vector<4xf32>) {
+ // CHECK: llvm.load %[[PTR]] {{.*}} : !llvm.ptr<1> -> vector<4xf32>
+ %0 = xegpu.load %src[%offsets], %mask <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : i64, vector<4xindex>, vector<4xi1> -> vector<4xf32>
+ gpu.return %0 : vector<4xf32>
+}
+}
+
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @store_scatter_contiguous_run
+// CHECK-SAME: %[[ARG0:.*]]: i64, %[[ARG1:.*]]: index, %[[ARG2:.*]]: vector<4xf32>
+gpu.func @store_scatter_contiguous_run(%dst: i64, %base: index, %val: vector<4xf32>) {
+ %mask = arith.constant dense<true> : vector<4xi1>
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %o1 = arith.addi %base, %c1 : index
+ %o2 = arith.addi %base, %c2 : index
+ %o3 = arith.addi %base, %c3 : index
+ %offsets = vector.from_elements %base, %o1, %o2, %o3 : vector<4xindex>
+ // An all-true mask leaves no scf.if at all.
+ // CHECK-NOT: scf.if
+ // CHECK: %[[BASE:.*]] = vector.extract %{{.*}}[0] : i64 from vector<4xi64>
+ // CHECK: %[[PTR:.*]] = llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<1>
+ // CHECK: llvm.store %[[ARG2]], %[[PTR]] {{.*}} : vector<4xf32>, !llvm.ptr<1>
+ xegpu.store %val, %dst[%offsets], %mask <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<4xf32>, i64, vector<4xindex>, vector<4xi1>
+ gpu.return
+}
+}
+
+// -----
+
+// vector.step offsets, the run based at 0 that distribution emits for lane 0.
+
+gpu.module @test {
+// CHECK-LABEL: @load_gather_step_offsets
+// CHECK-SAME: %[[ARG0:.*]]: i64, %[[ARG1:.*]]: i1
+gpu.func @load_gather_step_offsets(%src: i64, %pred: i1) -> vector<4xf32> {
+ %offsets = vector.step : vector<4xindex>
+ %mask = vector.broadcast %pred : i1 to vector<4xi1>
+ // CHECK: %[[STEP:.*]] = vector.step : vector<4xindex>
+ // CHECK: %[[I64:.*]] = vector.bitcast %[[STEP]] : vector<4xindex> to vector<4xi64>
+ // CHECK: %[[BASE:.*]] = vector.extract %[[I64]][0] : i64 from vector<4xi64>
+ // CHECK: %[[PTR:.*]] = llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<1>
+ // CHECK: scf.if %[[ARG1]] -> (vector<4xf32>) {
+ // CHECK: llvm.load %[[PTR]] {{.*}} : !llvm.ptr<1> -> vector<4xf32>
+ %0 = xegpu.load %src[%offsets], %mask <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : i64, vector<4xindex>, vector<4xi1> -> vector<4xf32>
+ gpu.return %0 : vector<4xf32>
+}
+}
+
+// -----
+
+// A mask built from one repeated value is uniform too.
+
+gpu.module @test {
+// CHECK-LABEL: @load_gather_from_elements_mask
+gpu.func @load_gather_from_elements_mask(%src: i64, %base: index, %pred: i1) -> vector<2xf32> {
+ %c1 = arith.constant 1 : index
+ %o1 = arith.addi %base, %c1 : index
+ %offsets = vector.from_elements %base, %o1 : vector<2xindex>
+ %mask = vector.from_elements %pred, %pred : vector<2xi1>
+ // CHECK: %[[BASE:.*]] = vector.extract %{{.*}}[0] : i64 from vector<2xi64>
+ // CHECK: llvm.load %{{.*}} : !llvm.ptr<1> -> vector<2xf32>
+ %0 = xegpu.load %src[%offsets], %mask <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : i64, vector<2xindex>, vector<2xi1> -> vector<2xf32>
+ gpu.return %0 : vector<2xf32>
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index f06fbac607474..651a8934aaeba 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,12 +254,34 @@ 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
}
+// -----
+// A scalar mask and offset stand for a single element, so they cannot carry a
+// wider value. This was the old chunk_size form.
+func.func @load_gather_simt_2(%src: memref<?xf32>) {
+ %0 = arith.constant 0 : index
+ %1 = arith.constant 1 : i1
+ // expected-error at +1 {{Value shape must match mask shape}}
+ %2 = xegpu.load %src[%0], %1 : memref<?xf32>, index, i1 -> vector<8xf32>
+ return
+}
+
+// -----
+// Same for a size-1 mask and offset.
+func.func @load_gather_simt_3(%src: memref<?xf32>) {
+ %0 = arith.constant dense<0> : vector<1xindex>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // expected-error at +1 {{Value shape must match mask shape}}
+ %2 = xegpu.load %src[%0], %1
+ : memref<?xf32>, vector<1xindex>, vector<1xi1> -> vector<8xf32>
+ return
+}
+
// -----
func.func @store_scatter_vc_2(%dst: memref<?xf32>) {
%0 = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
@@ -275,10 +297,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,12 +309,23 @@ 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
}
+// -----
+// A scalar mask and offset cannot carry a wider value on the store side either.
+func.func @store_scatter_simt_2(%dst: memref<?xf32>) {
+ %0 = arith.constant 0 : index
+ %1 = arith.constant 1 : i1
+ %2 = arith.constant dense<2.9>: vector<8xf32>
+ // expected-error at +1 {{Value shape must match mask shape}}
+ xegpu.store %2, %dst[%0], %1 : vector<8xf32>, memref<?xf32>, index, i1
+ return
+}
+
// -----
func.func @prefetch_offset_wi_1(%src: memref<4x4xf32>) {
%offsets = arith.constant dense<[0]> : vector<1xindex>
@@ -321,30 +354,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 +388,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 +397,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..e5da1e57d9312 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -373,10 +373,10 @@ gpu.func @prefetch_nd_3d(%src: memref<4x8x16xf16>) {
gpu.return
}
-// 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: gpu.func @simt_load_4(%[[arg0:.*]]: memref<256xf16>, %[[arg1:.*]]: vector<8xindex>, %[[arg2:.*]]: vector<8xi1>) {
+gpu.func @simt_load_4(%arg0: memref<256xf16>, %arg1: vector<8xindex>, %arg2: vector<8xi1>) {
+ // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] : memref<256xf16>, vector<8xindex>, vector<8xi1> -> vector<8xf16>
+ %0 = xegpu.load %arg0[%arg1], %arg2 : memref<256xf16>, vector<8xindex>, vector<8xi1> -> 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<1xf16>
+ %0 = xegpu.load %arg0[%arg1], %arg2 : memref<256xf16>, index, i1 -> vector<1xf16>
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>
+// CHECK: gpu.func @simt_store_4(%[[arg0:.*]]: vector<8xf16>, %[[arg1:.*]]: memref<256xf16>, %[[arg2:.*]]: vector<8xindex>, %[[arg3:.*]]: vector<8xi1>) {
+gpu.func @simt_store_4(%arg0: vector<8xf16>, %arg1: memref<256xf16>, %arg2: vector<8xindex>, %arg3: vector<8xi1>) {
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] : vector<8xf16>, memref<256xf16>, vector<8xindex>, vector<8xi1>
+ xegpu.store %arg0, %arg1[%arg2], %arg3 : vector<8xf16>, memref<256xf16>, vector<8xindex>, vector<8xi1>
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>
- 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: gpu.func @simt_store_5(%[[arg0:.*]]: vector<1xf16>, %[[arg1:.*]]: memref<256xf16>, %[[arg2:.*]]: index, %[[arg3:.*]]: i1) {
+gpu.func @simt_store_5(%arg0: vector<1xf16>, %arg1: memref<256xf16>, %arg2: index, %arg3: i1) {
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] : vector<1xf16>, memref<256xf16>, index, i1
+ xegpu.store %arg0, %arg1[%arg2], %arg3 : vector<1xf16>, 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-invalid.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup-invalid.mlir
index e0271728dc423..1f8c703d17eb6 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup-invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup-invalid.mlir
@@ -27,7 +27,7 @@ gpu.module @test {
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
// expected-error at +1 {{Failed to determine required layout for store scatter.}}
- 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-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index dca0c33012b21..a48d34072016f 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -595,9 +595,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 97701e738009e..1ab96746a65c5 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 415447351ca45..6eab8264f2a16 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -274,42 +274,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 85f19d070a4b2..0682163aad602 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
@@ -239,61 +239,63 @@ gpu.module @xevm_module{
}
// -----
+// A gather/scatter inside an scf.if that yields the loaded value. Each lane
+// takes 8 elements, spelled as a congruent 128-element access with
+// lane_data = [2] rather than as a chunk.
// 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-DAG: %[[CST:.*]] = arith.constant dense<1.200000e+01> : vector<8xf16>
+// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<8xi1>
+// CHECK-DAG: %[[OFFSET:.*]] = arith.constant dense<12> : vector<8xindex>
+// CHECK: %[[IF:.*]] = scf.if %[[PREDICATE]] -> (vector<8xf16>) {
+// CHECK-NEXT: %[[LD:.*]] = xegpu.load %{{.*}}[%[[OFFSET]]], %[[MASK]]
+// CHECK-SAME: : memref<256xf16>, vector<8xindex>, vector<8xi1> -> vector<8xf16>
+// CHECK-NEXT: scf.yield %[[LD]] : vector<8xf16>
// CHECK-NEXT: } else {
-// CHECK-NEXT: scf.yield %[[CST]] : vector<1x8xf16>
+// CHECK-NEXT: scf.yield %[[CST]] : vector<8xf16>
// 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>
+// CHECK-NEXT: xegpu.store %[[IF]], %{{.*}}[%[[OFFSET]]], %[[MASK]]
+// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<8xindex>, vector<8xi1>
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>
+ %1 = arith.constant dense<1>: vector<128xi1>
+ %offset = arith.constant dense<12> : vector<128xindex>
+ %loaded = scf.if %pred -> (vector<128xf16>) {
+ %3 = xegpu.load %src[%offset], %1 <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}>
+ : memref<256xf16>, vector<128xindex>, vector<128xi1> -> vector<128xf16>
+ scf.yield %3 : vector<128xf16>
} else {
- %3 = arith.constant dense<12.> : vector<16x8xf16>
- scf.yield %3 : vector<16x8xf16>
+ %3 = arith.constant dense<12.> : vector<128xf16>
+ scf.yield %3 : vector<128xf16>
}
- 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>
+ xegpu.store %loaded, %src[%offset], %1 <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}>
+ : vector<128xf16>, memref<256xf16>, vector<128xindex>, vector<128xi1>
gpu.return
}
}
// -----
+// Same, but the scf.if yields nothing.
// 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: %[[MASK:.*]] = arith.constant dense<true> : vector<8xi1>
+// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<8xindex>
// 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: %[[LOADED:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]]
+// CHECK-SAME: memref<256xf16>, vector<8xindex>, vector<8xi1> -> vector<8xf16>
+// CHECK-NEXT: xegpu.store %[[LOADED]], %arg0[%[[OFFSET]]], %[[MASK]]
+// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<8xindex>, vector<8xi1>
// 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>
+ %1 = arith.constant dense<1>: vector<128xi1>
+ %offset = arith.constant dense<12> : vector<128xindex>
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>
+ %3 = xegpu.load %src[%offset], %1 <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}>
+ : memref<256xf16>, vector<128xindex>, vector<128xi1> -> vector<128xf16>
+ xegpu.store %3, %src[%offset], %1 <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}>
+ : vector<128xf16>, memref<256xf16>, vector<128xindex>, vector<128xi1>
}
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index c7cbdf16a1171..ecdca7851cb70 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 c88a5a80674f8..49a3f02dca58e 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 1579eeb1e6552..6eaadae363cb5 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -304,11 +304,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
}
@@ -319,31 +319,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>) {
@@ -751,9 +737,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>
@@ -770,7 +756,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>
@@ -787,7 +773,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
<{
@@ -843,7 +829,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>
@@ -860,7 +846,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
<{
@@ -875,7 +861,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>
@@ -892,7 +878,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
<{
@@ -988,7 +974,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
@@ -1017,7 +1003,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
More information about the Mlir-commits
mailing list