[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 14:01:55 PDT 2026
https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/205122
>From 50d3ded4e60c535f2b6958dc6bf3ff3304d8d2af 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 | 87 +++++++++++++
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 | 116 +++---------------
.../Transforms/XeGPUWgToSgDistribute.cpp | 16 +--
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 17 +--
.../XeGPUToXeVM/failed_conversion.mlir | 31 +++++
.../loadstore_gather_coalesce.mlir | 100 +++++++++++++++
mlir/test/Dialect/XeGPU/invalid.mlir | 92 +++++++-------
mlir/test/Dialect/XeGPU/ops.mlir | 49 +++-----
.../XeGPU/propagate-layout-subgroup.mlir | 4 +-
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 4 +-
.../XeGPU/sg-to-lane-distribute-unit.mlir | 36 ------
.../Dialect/XeGPU/sg-to-lane-distribute.mlir | 61 ---------
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 82 +++----------
.../Dialect/XeGPU/xegpu-unroll-patterns.mlir | 56 +--------
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 42 +++----
22 files changed, 478 insertions(+), 623 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 cb39e25142c46..b0fd97fc9529a 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 defc389006af5..930e09980129e 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -166,12 +166,6 @@ DistributeLayoutAttr inferExtractSourceLayout(DistributeLayoutAttr resLayout,
ArrayRef<int64_t> resShape,
ArrayRef<int64_t> srcShape);
-/// Infers the layout attribute for mask and offset operand for Chunked load
-/// and store, given the anchor layout attribute for the value being load/store.
-DistributeLayoutAttr
-inferMaskOffsetLayoutForScatterIO(DistributeLayoutAttr payloadLayout,
- int chunkSize);
-
/// Infers the source layout attribute for an operand using result layout
/// attribute
DistributeLayoutAttr
@@ -277,6 +271,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);
@@ -287,6 +283,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 9863206f14fe1..a17f955def4ec 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -461,7 +461,6 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
vectorShape);
auto gatherOp = xegpu::LoadGatherOp::create(
rewriter, loc, vectorType, flatMemref, localOffsets, mask,
- /*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
@@ -496,7 +495,6 @@ static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
vectorShape);
xegpu::StoreScatterOp::create(rewriter, loc, writeOp.getVector(), flatMemref,
localOffsets, mask,
- /*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
@@ -741,7 +739,6 @@ struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
auto xeGatherOp = xegpu::LoadGatherOp::create(
rewriter, loc, vectorType, flatMemref, localOffsets, gatherOp.getMask(),
- /*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
@@ -776,7 +773,6 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
xegpu::StoreScatterOp::create(rewriter, loc, scatterOp.getValueToStore(),
flatMemref, localOffsets, scatterOp.getMask(),
- /*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index f782be5502a8a..35dc693f7260a 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -705,6 +705,72 @@ static Value addOffsetToBaseAddr(ConversionPatternRewriter &rewriter,
return newAddr;
}
+// If `v` is `base + cst`, returns {base, cst}. A bare value counts as
+// `{v, 0}`. Only an `index`/integer add by a constant is matched, on either
+// operand.
+static std::pair<Value, int64_t> splitConstantAddend(Value v) {
+ auto add = v.getDefiningOp<arith::AddIOp>();
+ if (!add)
+ return {v, 0};
+ APInt cst;
+ if (matchPattern(add.getRhs(), m_ConstantInt(&cst)))
+ return {add.getLhs(), cst.getSExtValue()};
+ if (matchPattern(add.getLhs(), m_ConstantInt(&cst)))
+ return {add.getRhs(), cst.getSExtValue()};
+ return {v, 0};
+}
+
+// Returns true when `offsets` is one ascending run of consecutive elements,
+// i.e. `{base, base + 1, ..., base + n - 1}`. Element 0 is then the base of the
+// run.
+//
+// A lane that owns such a run reads or writes a contiguous block, which the
+// chunked XeVM access expects as a single base offset. Two forms are matched,
+// both of which the XeGPU distribution emits:
+// - `vector.from_elements %b, %b + 1, ...`
+// - `vector.step`
+// Elements are compared in row-major order, which is the order
+// `vector.from_elements` uses, so nD offsets work too.
+static bool isContiguousRun(Value offsets) {
+ auto vecTy = dyn_cast<VectorType>(offsets.getType());
+ if (!vecTy || vecTy.getNumElements() < 1)
+ return false;
+
+ if (offsets.getDefiningOp<vector::StepOp>())
+ return true;
+
+ auto fromElements = offsets.getDefiningOp<vector::FromElementsOp>();
+ if (!fromElements)
+ return false;
+ ValueRange elements = fromElements.getElements();
+
+ // Every element must be the same base value plus its own index.
+ auto [base, firstAddend] = splitConstantAddend(elements.front());
+ for (int64_t i = 1, e = elements.size(); i < e; ++i) {
+ auto [elemBase, addend] = splitConstantAddend(elements[i]);
+ if (elemBase != base || addend != firstAddend + i)
+ return false;
+ }
+ return true;
+}
+
+// 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 +836,27 @@ class LoadStoreToXeVMPattern : public OpConversionPattern<OpType> {
basePtrI64);
}
Value mask = adaptor.getMask();
+
+ // Coalesce a chunked access. Distribution gives every element its own
+ // offset and mask bit, so a lane that owns a contiguous run arrives here
+ // with `vector<D>` offsets and mask beside a `vector<D>` value. One
+ // contiguous block access needs a single base offset and a single mask bit,
+ // so take them from lane 0. The offsets are proven consecutive and the mask
+ // uniform on the original operands, whose defining ops are still intact;
+ // the values come from the converted operands, so no casts are needed. When
+ // the proof fails the access stays a gather and the match fails below, as
+ // before.
+ auto origOffsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ if (isa<VectorType>(offset.getType()) && origOffsetsTy && valOrResVecTy &&
+ origOffsetsTy.getNumElements() == valOrResVecTy.getNumElements() &&
+ isContiguousRun(op.getOffsets()) && 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 686e4342215b0..fbd2739f8a74c 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -64,53 +64,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();
}
@@ -605,7 +588,6 @@ LogicalResult LoadGatherOp::verify() {
return emitOpError("invalid l3_hint: ") << getL3HintAttr();
auto srcTy = getSourceType();
- uint64_t chunkSize = static_cast<int64_t>(getChunkSize().value_or(1));
auto memTy = dyn_cast<MemRefType>(srcTy);
if (memTy && (getElementType() != memTy.getElementType()))
@@ -620,14 +602,14 @@ LogicalResult LoadGatherOp::verify() {
if (failed(isValidContiguity(getContiguity(), offsetsTy,
[&]() { return emitOpError(); })))
return failure();
- return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
+ return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy,
[&]() { return emitOpError(); });
}
void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
Type valueType, Value source,
ArrayRef<OpFoldResult> offsets, Value mask,
- IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint,
+ xegpu::CachePolicyAttr l1_hint,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint) {
auto loc = source.getLoc();
@@ -636,15 +618,15 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto values = getValueOrCreateConstantIndexOp(builder, loc, offsets);
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
- build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
- l2_hint, l3_hint, /*anchor_layout=*/nullptr,
+ build(builder, state, valueType, source, offset, mask, l1_hint, l2_hint,
+ l3_hint, /*anchor_layout=*/nullptr,
/*contiguity=*/nullptr);
}
void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
Type valueType, Value source,
ArrayRef<OpFoldResult> offsets, Value mask,
- IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint,
+ xegpu::CachePolicyAttr l1_hint,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint,
DistributeLayoutAttr layout) {
@@ -654,8 +636,8 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto values = getValueOrCreateConstantIndexOp(builder, loc, offsets);
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
- build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
- l2_hint, l3_hint, layout, /*contiguity=*/nullptr);
+ build(builder, state, valueType, source, offset, mask, l1_hint, l2_hint,
+ l3_hint, layout, /*contiguity=*/nullptr);
}
//===----------------------------------------------------------------------===//
@@ -675,7 +657,6 @@ LogicalResult StoreScatterOp::verify() {
return emitOpError("invalid l3_hint: ") << getL3HintAttr();
auto destTy = getDestType();
- uint64_t chunkSize = static_cast<int64_t>(getChunkSize().value_or(1));
auto memTy = dyn_cast<MemRefType>(destTy);
if (memTy && (getElementType() != memTy.getElementType()))
@@ -690,14 +671,13 @@ LogicalResult StoreScatterOp::verify() {
if (failed(isValidContiguity(getContiguity(), offsetsTy,
[&]() { return emitOpError(); })))
return failure();
- return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
+ return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy,
[&]() { return emitOpError(); });
}
void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
Value value, Value dest,
ArrayRef<OpFoldResult> offsets, Value mask,
- IntegerAttr chunk_size,
xegpu::CachePolicyAttr l1_hint,
xegpu::CachePolicyAttr l2_hint,
xegpu::CachePolicyAttr l3_hint) {
@@ -708,15 +688,17 @@ void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
// Call the correct builder overload that does not expect result types.
- build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
- l3_hint, /*anchor_layout=*/nullptr, /*contiguity=*/nullptr);
+ build(builder, state, value, dest, offset, mask, l1_hint, l2_hint, l3_hint,
+ /*anchor_layout=*/nullptr, /*contiguity=*/nullptr);
}
-void StoreScatterOp::build(
- OpBuilder &builder, OperationState &state, Value value, Value dest,
- ArrayRef<OpFoldResult> offsets, Value mask, IntegerAttr chunk_size,
- xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint,
- xegpu::CachePolicyAttr l3_hint, DistributeLayoutAttr layout) {
+void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
+ Value value, Value dest,
+ ArrayRef<OpFoldResult> offsets, Value mask,
+ xegpu::CachePolicyAttr l1_hint,
+ xegpu::CachePolicyAttr l2_hint,
+ xegpu::CachePolicyAttr l3_hint,
+ DistributeLayoutAttr layout) {
auto loc = dest.getLoc();
int64_t size = static_cast<int64_t>(offsets.size());
auto type = VectorType::get(size, builder.getIndexType());
@@ -724,8 +706,8 @@ void StoreScatterOp::build(
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
// Call the correct builder overload that does not expect result types.
- build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
- l3_hint, layout, /*contiguity=*/nullptr);
+ build(builder, state, value, dest, offset, mask, l1_hint, l2_hint, l3_hint,
+ layout, /*contiguity=*/nullptr);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index b11f7ecd3df06..0554cb676bf19 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1002,17 +1002,6 @@ xegpu::DistributeLayoutAttr xegpu::inferResultLayoutFromSourceForNonAnchorOp(
return nullptr;
}
-/// Infers the layout attribute for mask and offset operand for Chunked load
-/// and store, given the anchor layout attribute for the value being load/store.
-xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
- xegpu::DistributeLayoutAttr payloadLayout, int chunkSize) {
- auto rank = payloadLayout.getRank();
- if (chunkSize > 1)
- return payloadLayout.dropDims(
- llvm::to_vector(llvm::seq<int64_t>(rank - 1, rank)));
- return payloadLayout;
-}
-
//===----------------------------------------------------------------------===//
// Layout derivation helpers: factorize sgCount into
// sg_layout candidates, then
@@ -1849,13 +1838,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,
@@ -1899,6 +1894,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 =
@@ -1979,6 +1976,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 19d02e02b02ba..3f65161b5e6c4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1234,7 +1234,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())
@@ -1268,14 +1270,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.
@@ -1299,7 +1299,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;
@@ -1331,7 +1333,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.");
@@ -1341,10 +1343,8 @@ void LayoutInfoPropagation::visitStoreScatterOp(
}
LayoutInfo srcLayoutInfo = makeLayoutInfo(requiredAnchorLayoutAttr);
- assert((chunkSize <= 1) || (layoutKind != xegpu::LayoutKind::Subgroup));
- auto maskLayoutAttr = xegpu::inferMaskOffsetLayoutForScatterIO(
- requiredAnchorLayoutAttr, chunkSize);
- LayoutInfo maskLayoutInfo = makeLayoutInfo(maskLayoutAttr);
+ // The mask and offset operands share the value's anchor layout.
+ LayoutInfo maskLayoutInfo = makeLayoutInfo(requiredAnchorLayoutAttr);
// Propagate the payload operand layout
propagateIfChanged(operands[0], operands[0]->meet(srcLayoutInfo));
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index bc85140f9f121..45e0bac55c275 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -497,7 +497,7 @@ struct SgToLanePrefetchNd : public OpConversionPattern<xegpu::PrefetchNdOp> {
/// Distributes a subgroup-level LoadGather (xegpu.load) op to lane-level.
///
-/// Example 1 (1D, no chunk size):
+/// Example 1 (1D):
/// layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
/// %mask = producer_op : vector<16xi1>
/// %offset = producer_op : vector<16xindex>
@@ -509,15 +509,7 @@ struct SgToLanePrefetchNd : public OpConversionPattern<xegpu::PrefetchNdOp> {
/// %0 = xegpu.load %src[%offset], %mask : memref<256xf16>,
/// vector<1xindex>, vector<1xi1> -> vector<1xf16>
///
-/// Example 2 (2D with chunk size, same mask & offset):
-/// layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
-/// %0 = xegpu.load %src[%offset], %mask <{chunk_size=8}> :
-/// memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
-/// Distributed to:
-/// %0 = xegpu.load %src[%offset], %mask <{chunk_size=8}> :
-/// memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-///
-/// Example 3 (3D with leading unit dims):
+/// Example 2 (3D with leading unit dims):
/// layout = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>
/// %mask = producer_op : vector<1x1x16xi1>
/// %offset = producer_op : vector<1x1x16xindex>
@@ -543,8 +535,7 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
return failure();
// Check that leading dimensions are unit.
- int chunkSize = op.getChunkSize().value_or(1);
- int effectiveVecRank = (chunkSize == 1) ? 1 : 2;
+ int effectiveVecRank = 1;
ArrayRef<int64_t> shape = origResultTy.getShape();
if (llvm::any_of(
shape.take_front(origResultTy.getRank() - effectiveVecRank),
@@ -581,8 +572,8 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
Value distSource = adaptor.getSource();
auto newOp = xegpu::LoadGatherOp::create(
rewriter, op.getLoc(), distResultTy1D, distSource, distOffsets,
- distMask, op.getChunkSizeAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/nullptr, /*contiguity=*/nullptr);
+ distMask, op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(),
+ /*layout=*/nullptr, /*contiguity=*/nullptr);
Value result = newOp->getResult(0);
if (distResultTy1D != distResultTy)
@@ -1024,7 +1015,7 @@ struct SgToLaneStoreMatrix : public OpConversionPattern<xegpu::StoreMatrixOp> {
/// Distributes a subgroup-level StoreScatter (xegpu.store) op to
/// lane-level.
///
-/// Example 1 (1D, no chunk size):
+/// Example 1 (1D):
/// layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
/// %mask = producer_op : vector<16xi1>
/// %offset = producer_op : vector<16xindex>
@@ -1036,15 +1027,7 @@ struct SgToLaneStoreMatrix : public OpConversionPattern<xegpu::StoreMatrixOp> {
/// xegpu.store %payload, %src[%offset], %mask : vector<1xf16>,
/// memref<256xf16>, vector<1xindex>, vector<1xi1>
///
-/// Example 2 (2D with chunk size, same mask & offset):
-/// layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
-/// xegpu.store %payload, %src[%offset], %mask <{chunk_size=8}> :
-/// vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
-/// Distributed to:
-/// xegpu.store %payload, %src[%offset], %mask <{chunk_size=8}> :
-/// vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
-///
-/// Example 3 (3D with leading unit dims):
+/// Example 2 (3D with leading unit dims):
/// layout = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>
/// %mask = producer_op : vector<1x1x16xi1>
/// %offset = producer_op : vector<1x1x16xindex>
@@ -1071,8 +1054,7 @@ struct SgToLaneStoreScatter
return failure();
// Check that all leading dimensions are unit dimensions.
- int chunkSize = op.getChunkSize().value_or(1);
- int effectiveVecRank = (chunkSize == 1) ? 1 : 2;
+ int effectiveVecRank = 1;
ArrayRef<int64_t> shape = origValueTy.getShape();
if (llvm::any_of(shape.take_front(origValueTy.getRank() - effectiveVecRank),
[](int64_t d) { return d != 1; }))
@@ -1112,9 +1094,9 @@ struct SgToLaneStoreScatter
Value distDest = adaptor.getDest();
xegpu::StoreScatterOp::create(rewriter, op.getLoc(), distValue, distDest,
- distOffsets, distMask, op.getChunkSizeAttr(),
- op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/nullptr,
+ distOffsets, distMask, op.getL1HintAttr(),
+ op.getL2HintAttr(), op.getL3HintAttr(),
+ /*layout=*/nullptr,
/*contiguity=*/nullptr);
rewriter.eraseOp(op);
return success();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 5914e23f1f11a..be3a1b66c8b7d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -563,11 +563,6 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
return failure();
SmallVector<int64_t> targetMaskShape(*targetShape);
- int64_t chunkSize = 1;
- if (auto chunkSizeAttr = op->getAttr("chunk_size")) {
- if (auto intAttr = llvm::dyn_cast<IntegerAttr>(chunkSizeAttr))
- chunkSize = intAttr.getInt();
- }
// Unroll mask and offsets with correct shape
VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
@@ -575,50 +570,15 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
Type elemTy = valueTy.getElementType();
VectorType newValueTy = VectorType::get(*targetShape, elemTy);
- SmallVector<Type> convertedMaskTypes;
- SmallVector<Value> convertedMasks;
- SmallVector<Type> convertedOffsetTypes;
- SmallVector<Value> convertedOffsets;
-
- if (chunkSize > 1) {
- // For chunked loads, mask and offsets have one less dimension
- targetMaskShape.pop_back();
- int64_t blockedChunkSize = targetShape->back();
- int64_t numNewChunks = chunkSize / blockedChunkSize;
- chunkSize = blockedChunkSize;
-
- convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
- convertedOffsetTypes = getUnrolledTypes(offsetsTy, targetMaskShape);
-
- SmallVector<Value> convertedMasksBase =
- pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
- SmallVector<Value> convertedOffsetsBase =
- pack(offsets, convertedOffsetTypes, targetMaskShape, loc, rewriter);
-
- for (auto maskVal : convertedMasksBase)
- convertedMasks.append(numNewChunks, maskVal);
-
- for (auto [baseOffset, offsetType] :
- llvm::zip(convertedOffsetsBase, convertedOffsetTypes)) {
- for (int64_t i = 0; i < numNewChunks; ++i) {
- Value inc = arith::ConstantIndexOp::create(rewriter, loc,
- i * blockedChunkSize);
- Value incVec =
- vector::BroadcastOp::create(rewriter, loc, offsetType, inc);
- Value offsetVal =
- arith::AddIOp::create(rewriter, loc, baseOffset, incVec);
- convertedOffsets.push_back(offsetVal);
- }
- }
- } else {
- convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
- convertedMasks =
- pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
+ SmallVector<Type> convertedMaskTypes =
+ getUnrolledTypes(maskTy, targetMaskShape);
+ SmallVector<Value> convertedMasks =
+ pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
- convertedOffsetTypes = getUnrolledTypes(offsetsTy, *targetShape);
- convertedOffsets =
- pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
- }
+ SmallVector<Type> convertedOffsetTypes =
+ getUnrolledTypes(offsetsTy, *targetShape);
+ SmallVector<Value> convertedOffsets =
+ pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
auto layout = op.getLayoutAttr();
if (layout)
@@ -627,8 +587,7 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
SmallVector<Value> newOps;
for (auto [o, m] : llvm::zip(convertedOffsets, convertedMasks)) {
auto newOp = xegpu::LoadGatherOp::create(
- rewriter, loc, newValueTy, op.getSource(), o, m,
- rewriter.getI64IntegerAttr(chunkSize), op.getL1HintAttr(),
+ rewriter, loc, newValueTy, op.getSource(), o, m, op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr(), layout,
/*contiguity=*/nullptr);
newOps.push_back(newOp);
@@ -657,59 +616,19 @@ struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
if (!targetShape)
return failure();
- int64_t chunkSize = 1;
- if (auto chunkSizeAttr = op->getAttr("chunk_size")) {
- if (auto intAttr = llvm::dyn_cast<IntegerAttr>(chunkSizeAttr))
- chunkSize = intAttr.getInt();
- }
-
SmallVector<int64_t> targetMaskShape(*targetShape);
VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
VectorType offsetsTy = llvm::dyn_cast<VectorType>(offsets.getType());
- SmallVector<Type> convertedMaskTypes;
- SmallVector<Value> convertedMasks;
- SmallVector<Type> convertedOffsetTypes;
- SmallVector<Value> convertedOffsets;
-
- if (chunkSize > 1) {
- targetMaskShape.pop_back();
- int64_t blockedChunkSize = targetShape->back();
- int64_t numNewChunks = chunkSize / blockedChunkSize;
- chunkSize = blockedChunkSize;
-
- convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
- convertedOffsetTypes = getUnrolledTypes(offsetsTy, targetMaskShape);
-
- SmallVector<Value> convertedMasksBase =
- pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
- SmallVector<Value> convertedOffsetsBase =
- pack(offsets, convertedOffsetTypes, targetMaskShape, loc, rewriter);
-
- for (auto maskVal : convertedMasksBase)
- convertedMasks.append(numNewChunks, maskVal);
-
- for (auto [baseOffset, offsetType] :
- llvm::zip(convertedOffsetsBase, convertedOffsetTypes)) {
- for (int64_t i = 0; i < numNewChunks; ++i) {
- Value inc = arith::ConstantIndexOp::create(rewriter, loc,
- i * blockedChunkSize);
- Value incVec =
- vector::BroadcastOp::create(rewriter, loc, offsetType, inc);
- Value offsetVal =
- arith::AddIOp::create(rewriter, loc, baseOffset, incVec);
- convertedOffsets.push_back(offsetVal);
- }
- }
- } else {
- convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
- convertedMasks =
- pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
+ SmallVector<Type> convertedMaskTypes =
+ getUnrolledTypes(maskTy, targetMaskShape);
+ SmallVector<Value> convertedMasks =
+ pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
- convertedOffsetTypes = getUnrolledTypes(offsetsTy, *targetShape);
- convertedOffsets =
- pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
- }
+ SmallVector<Type> convertedOffsetTypes =
+ getUnrolledTypes(offsetsTy, *targetShape);
+ SmallVector<Value> convertedOffsets =
+ pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
SmallVector<Type> convertedValTypes =
getUnrolledTypes(valueTy, *targetShape);
@@ -723,7 +642,6 @@ struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
for (auto [v, o, m] :
llvm::zip(convertedValues, convertedOffsets, convertedMasks)) {
xegpu::StoreScatterOp::create(rewriter, loc, v, op.getDest(), o, m,
- rewriter.getI64IntegerAttr(chunkSize),
op.getL1HintAttr(), op.getL2HintAttr(),
op.getL3HintAttr(), layout,
/*contiguity=*/nullptr);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 0e8a386fb08b6..07c067ce076f2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -828,14 +828,12 @@ struct WgToSgLoadGatherOp : public OpConversionPattern<xegpu::LoadGatherOp> {
}
SmallVector<Value> newLoadOps;
- auto chunkSizeAttr =
- rewriter.getI64IntegerAttr(op.getChunkSize().value_or(1));
VectorType newTy = VectorType::get(sgShape, resultType.getElementType());
for (auto [offsets, mask] :
llvm::zip(adaptor.getOffsets(), adaptor.getMask())) {
auto newLayout = layout.dropSgLayoutAndData();
auto newLoadOp = xegpu::LoadGatherOp::create(
- rewriter, loc, newTy, op.getSource(), offsets, mask, chunkSizeAttr,
+ rewriter, loc, newTy, op.getSource(), offsets, mask,
op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(), newLayout,
/*contiguity=*/nullptr);
newLoadOps.push_back(newLoadOp);
@@ -875,16 +873,12 @@ struct WgToSgStoreScatterOp
"offsets have not been distributed");
}
- auto chunkSizeOpt = op.getChunkSize();
- int64_t chunkSize = chunkSizeOpt ? static_cast<int64_t>(*chunkSizeOpt) : 1;
- auto chunkSizeAttr = rewriter.getI64IntegerAttr(chunkSize);
for (auto [val, offs, mask] : llvm::zip(
adaptor.getValue(), adaptor.getOffsets(), adaptor.getMask())) {
- xegpu::StoreScatterOp::create(rewriter, loc, val, op.getDest(), offs,
- mask, chunkSizeAttr, op.getL1HintAttr(),
- op.getL2HintAttr(), op.getL3HintAttr(),
- layout.dropSgLayoutAndData(),
- /*contiguity=*/nullptr);
+ xegpu::StoreScatterOp::create(
+ rewriter, loc, val, op.getDest(), offs, mask, op.getL1HintAttr(),
+ op.getL2HintAttr(), op.getL3HintAttr(), layout.dropSgLayoutAndData(),
+ /*contiguity=*/nullptr);
}
rewriter.eraseOp(op);
return success();
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 76269cf193d13..67097cfd3b6aa 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -235,22 +235,9 @@ xegpu::getDistributeLayoutAttr(const OpOperand &opr) {
if (isa<xegpu::StoreNdOp, xegpu::StoreMatrixOp>(op) && (idx < 2))
return layout;
- if (isa<xegpu::StoreScatterOp>(op)) {
- xegpu::StoreScatterOp store(op);
- int chunkSize = store.getChunkSize().value_or(1);
- if (layout && idx >= 2 && chunkSize > 1)
- return layout.dropDims(llvm::to_vector(
- llvm::seq<int64_t>(layout.getRank() - 1, layout.getRank())));
+ // For gather/scatter ops the mask and offsets share the value's layout.
+ if (isa<xegpu::StoreScatterOp, xegpu::LoadGatherOp>(op))
return layout;
- }
- if (isa<xegpu::LoadGatherOp>(op)) {
- xegpu::LoadGatherOp load(op);
- int chunkSize = load.getChunkSize().value_or(1);
- if (layout && idx >= 1 && chunkSize > 1)
- return layout.dropDims(llvm::to_vector(
- llvm::seq<int64_t>(layout.getRank() - 1, layout.getRank())));
- return layout;
- }
}
std::string layoutName = xegpu::getTemporaryLayoutName(opr);
diff --git a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
index dfaa9521786a9..edb972a48d732 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -31,6 +31,37 @@ gpu.module @test_kernel {
// -----
+// Vector offsets are only coalesced into one block access when they form a
+// contiguous run. A gather of unrelated offsets has no block form.
+
+gpu.module @test_kernel {
+ gpu.func @load_gather_non_contiguous_offsets(%src: i64, %a: index, %b: index) -> vector<2xf32> {
+ %offsets = vector.from_elements %a, %b : vector<2xindex>
+ %mask = arith.constant dense<true> : vector<2xi1>
+ // 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>
+ }
+}
+
+// -----
+
+// A contiguous run still needs a uniform mask: 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..50ef9d3a2ab4e
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_gather_coalesce.mlir
@@ -0,0 +1,100 @@
+// 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. When
+// the offsets are one contiguous run and the mask is uniform, the access is
+// coalesced into a single block load/store from the base of the run.
+
+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 is a contiguous run based at 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 932cf6ff9b201..2967f84d7184f 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.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index 02d59f80892e7..d7ecf59048b81 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -596,9 +596,9 @@ gpu.module @test {
%val = arith.constant dense<25.5> : vector<256xf16>
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
- // CHECK: xegpu.store %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
+ // CHECK: xegpu.store %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}} <{l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
// CHECK-SAME: : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
- xegpu.store %val, %dest[%offset], %mask {chunk_size = 1, l1_hint = #xegpu.cache_hint<cached>}
+ xegpu.store %val, %dest[%offset], %mask {l1_hint = #xegpu.cache_hint<cached>}
: vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index b4e69173c1144..5b48c8ce759c9 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -641,8 +641,8 @@ gpu.module @test{
%5 = vector.broadcast %3 : index to vector<1xindex>
%6 = arith.addi %4, %5 : vector<1xindex>
%7 = vector.broadcast %6 : vector<1xindex> to vector<1x1x1x16xindex>
- xegpu.store %cst, %0[%7], %cst_0 <{chunk_size = 1 : i64}> : vector<1x1x1x16xf32>, i64, vector<1x1x1x16xindex>, vector<1x1x1x16xi1>
- xegpu.store %cst, %0[%7], %cst_0 <{chunk_size = 1 : i64}> : vector<1x1x1x16xf32>, i64, vector<1x1x1x16xindex>, vector<1x1x1x16xi1>
+ xegpu.store %cst, %0[%7], %cst_0 : vector<1x1x1x16xf32>, i64, vector<1x1x1x16xindex>, vector<1x1x1x16xi1>
+ xegpu.store %cst, %0[%7], %cst_0 : vector<1x1x1x16xf32>, i64, vector<1x1x1x16xindex>, vector<1x1x1x16xi1>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index 0f3cea16797f1..d3ccfaa54433e 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -277,42 +277,6 @@ gpu.func @prefetch_nd() {
gpu.return
}
-// CHECK-LABEL: gpu.func @scatter_load_chunksize
-// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
-// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<8xf16> to vector<1x8xf16>
-gpu.func @scatter_load_chunksize(%src: memref<256xf16>) {
- %offset = arith.constant dense<12> : vector<16xindex>
- %mask = arith.constant dense<true> : vector<16xi1>
- %0 = xegpu.load %src[%offset], %mask
- <{chunk_size = 8, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- gpu.return
-}
-
-// CHECK-LABEL: gpu.func @scatter_store_chunksize
-// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
-// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK: %[[C1:.*]] = vector.shape_cast %[[LOAD]] : vector<8xf16> to vector<1x8xf16>
-// CHECK: %[[C2:.*]] = vector.shape_cast %[[C1]] : vector<1x8xf16> to vector<8xf16>
-// CHECK: xegpu.store %[[C2]], %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
-gpu.func @scatter_store_chunksize(%src: memref<256xf16>) {
- %offset = arith.constant dense<12> : vector<16xindex>
- %mask = arith.constant dense<true> : vector<16xi1>
- %0 = xegpu.load %src[%offset], %mask
- <{chunk_size = 8, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %0, %src[%offset], %mask
- <{chunk_size = 8, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
- : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
- gpu.return
-}
-
// CHECK-LABEL: gpu.func @scatter_load
// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
index 85b939b6d054a..ad046a7e432ed 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
@@ -250,67 +250,6 @@ gpu.module @xevm_module{
}
}
-// -----
-// CHECK-LABEL: gpu.func @scatter_ops_scf_yield
-// CHECK: (%{{.*}}: memref<256xf16>, %[[PREDICATE:[a-zA-Z0-9]+]]: i1) {
-// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.200000e+01> : vector<1x8xf16>
-// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK-DAG: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
-// CHECK: %[[IF:.*]] = scf.if %[[PREDICATE]] -> (vector<1x8xf16>) {
-// CHECK-NEXT: %[[LD:.*]] = xegpu.load %{{.*}}[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK-NEXT: %[[LD_CAST:.*]] = vector.shape_cast %[[LD]] : vector<8xf16> to vector<1x8xf16>
-// CHECK-NEXT: scf.yield %[[LD_CAST]] : vector<1x8xf16>
-// CHECK-NEXT: } else {
-// CHECK-NEXT: scf.yield %[[CST]] : vector<1x8xf16>
-// CHECK-NEXT: }
-// CHECK-NEXT: %[[IF_CAST:.*]] = vector.shape_cast %[[IF]] : vector<1x8xf16> to vector<8xf16>
-// CHECK-NEXT: xegpu.store %[[IF_CAST]], %{{.*}}[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
-gpu.module @xevm_module{
- gpu.func @scatter_ops_scf_yield(%src: memref<256xf16>, %pred : i1) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %loaded = scf.if %pred -> (vector<16x8xf16>) {
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {
- layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
- } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- scf.yield %3 : vector<16x8xf16>
- } else {
- %3 = arith.constant dense<12.> : vector<16x8xf16>
- scf.yield %3 : vector<16x8xf16>
- }
- xegpu.store %loaded, %src[%offset], %1 <{chunk_size=8}> {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
- gpu.return
- }
-}
-
-// -----
-// CHECK-LABEL: gpu.func @scatter_ops_scf_non_yield({{.*}}) {
-// CHECK: %[[PREDICATE:.*]] = llvm.mlir.poison : i1
-// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>
-// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>
-// CHECK: scf.if %[[PREDICATE]] {
-// CHECK-NEXT: %[[LOADED:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK-NEXT: xegpu.store %[[LOADED]], %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
-// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
-// CHECK-NEXT: }
-gpu.module @xevm_module{
- gpu.func @scatter_ops_scf_non_yield(%src: memref<256xf16>) {
- %pred = llvm.mlir.poison : i1
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- scf.if %pred {
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {
- layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
- } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=8}> {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
- }
- gpu.return
- }
-}
-
// -----
// CHECK-LABEL: gpu.func @mma_transpose_b(
// CHECK: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<16x8xi32>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index 490f2a622624e..3cc17e9b643e2 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -334,8 +334,8 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184,
192, 200, 208, 216, 224, 232, 240, 248
]> : vector<32xindex>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%cst], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%cst], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -353,8 +353,8 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184,
192, 200, 208, 216, 224, 232, 240, 248
]> : vector<32xindex>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%cst], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%cst], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -369,8 +369,8 @@ gpu.module @test_kernel {
//CHECK: arith.addi [[step]], [[cst]] : vector<16xindex>
%step = vector.step : vector<32xindex>
%mask = vector.create_mask %c16 : vector<32xi1>
- %ld = xegpu.load %src[%step], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%step], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%step], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%step], %mask {layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -557,7 +557,7 @@ gpu.module @test_kernel {
// -----
gpu.module @test_kernel {
// CHECK-LABEL: load_with_offsets
- // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.func @load_with_offsets(%src: ui64) -> vector<32xf32> {
%cst = arith.constant dense<[
0, 8, 16, 24, 32, 40, 48, 56,
@@ -568,7 +568,7 @@ gpu.module @test_kernel {
%c17 = arith.constant 17: index
%mask = vector.create_mask %c17 : vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %ld = xegpu.load %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %ld : vector<32xf32>
}
@@ -577,7 +577,7 @@ gpu.module @test_kernel {
// -----
gpu.module @test_kernel {
// CHECK-LABEL: store_with_offsets
- // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
+ // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
gpu.func @store_with_offsets(%src: ui64) {
%cst = arith.constant dense<[
0, 8, 16, 24, 32, 40, 48, 56,
@@ -590,62 +590,12 @@ gpu.module @test_kernel {
%mask = vector.create_mask %c17 : vector<32xi1>
%st_vec = arith.constant dense<1023.0>: vector<32xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ xegpu.store %st_vec, %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
-// -----
-gpu.module @test_kernel {
- // CHECK-LABEL: load_with_offsets_chunk
- // CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<32x4xf32>
- // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
- // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
- // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
- // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
- // CHECK-COUNT-4: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16x2xf32>
- gpu.func @load_with_offsets_chunk(%src: ui64) -> vector<32x4xf32> {
- %cst = arith.constant dense<[
- 0, 8, 16, 24, 32, 40, 48, 56,
- 64, 72, 80, 88, 96, 104, 112, 120,
- 128, 136, 144, 152, 160, 168, 176, 184,
- 192, 200, 208, 216, 224, 232, 240, 248
- ]> : vector<32xindex>
-
- %c17 = arith.constant 17: index
- %mask = vector.create_mask %c17 : vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
- gpu.return %ld : vector<32x4xf32>
- }
-}
-
-// -----
-gpu.module @test_kernel {
- // CHECK-LABEL: store_with_offsets_chunk
- // CHECK: [[cst:%.+]] = arith.constant dense<1.023000e+03> : vector<16x2xf32
- // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
- // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
- // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
- // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
- // CHECK-COUNT-4: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16x2xf32>, ui64, vector<16xindex>, vector<16xi1>
- gpu.func @store_with_offsets_chunk(%src: ui64) {
- %cst = arith.constant dense<[
- 0, 8, 16, 24, 32, 40, 48, 56,
- 64, 72, 80, 88, 96, 104, 112, 120,
- 128, 136, 144, 152, 160, 168, 176, 184,
- 192, 200, 208, 216, 224, 232, 240, 248
- ]> : vector<32xindex>
-
- %c17 = arith.constant 17: index
- %mask = vector.create_mask %c17 : vector<32xi1>
-
- %st_vec = arith.constant dense<1023.>: vector<32x4xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
- gpu.return
- }
-}
-
// -----
gpu.module @test_kernel {
// CHECK-LABEL: preserve_unit_dim_of_load_inst_data
@@ -654,8 +604,8 @@ gpu.module @test_kernel {
// CHECK: [[cst_0:%.+]] = arith.constant dense<true> : vector<1x1x16xi1>
// CHECK: [[cst_1:%.+]] = arith.constant dense<{{.*}}> : vector<1x1x16xindex>
// CHECK: [[cst_2:%.+]] = arith.constant dense<{{.*}}> : vector<1x1x16xindex>
- // CHECK: [[ld_0:%.+]] = xegpu.load [[arg0]][[[cst_1]]], [[cst_0]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
- // CHECK: [[ld_1:%.+]] = xegpu.load [[arg0]][[[cst_2]]], [[cst_0]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
+ // CHECK: [[ld_0:%.+]] = xegpu.load [[arg0]][[[cst_1]]], [[cst_0]] <{l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
+ // CHECK: [[ld_1:%.+]] = xegpu.load [[arg0]][[[cst_2]]], [[cst_0]] <{l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
// CHECK: [[ins_0:%.+]] = vector.insert_strided_slice [[ld_0]], [[cst]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x1x16xf32> into vector<1x1x32xf32>
// CHECK: [[ins_1:%.+]] = vector.insert_strided_slice [[ld_1]], [[ins_0]] offsets = [0, 0, 16], strides = [1, 1, 1] : vector<1x1x16xf32> into vector<1x1x32xf32>
gpu.func @preserve_unit_dim_of_load_inst_data(%src: ui64) -> vector<1x1x32xf32> {
@@ -667,7 +617,7 @@ gpu.module @test_kernel {
]]> : vector<1x1x32xindex>
%mask = arith.constant dense<true> : vector<1x1x32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [1, 1, 16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %ld = xegpu.load %src[%cst], %mask {layout = #xegpu.layout<inst_data = [1, 1, 16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
gpu.return %ld : vector<1x1x32xf32>
}
@@ -735,10 +685,10 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]]
]> : vector<1x1x32xindex>
%mask = arith.constant dense<true> : vector<1x1x32xi1>
- %a = xegpu.load %A[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
- %b = xegpu.load %B[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %a = xegpu.load %A[%cst], %mask {layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %b = xegpu.load %B[%cst], %mask {layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
%addf = arith.addf %a, %b : vector<1x1x32xf32>
- xegpu.store %addf, %C[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : vector<1x1x32xf32>, ui64, vector<1x1x32xindex>, vector<1x1x32xi1>
+ xegpu.store %addf, %C[%cst], %mask {layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : vector<1x1x32xf32>, ui64, vector<1x1x32xindex>, vector<1x1x32xi1>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
index f23ae46905652..9811a730e20f8 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
@@ -128,7 +128,7 @@ gpu.module @test {
//-----
// CHECK-LABEL: load_with_offsets
// CHECK-SAME: [[arg0:%.+]]: ui64
- // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.func @load_with_offsets(%src: ui64) -> vector<32xf32> {
%cst = arith.constant dense<[
0, 8, 16, 24, 32, 40, 48, 56,
@@ -139,7 +139,7 @@ gpu.module @test {
%c17 = arith.constant 17: index
%mask = vector.create_mask %c17: vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %ld = xegpu.load %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %ld : vector<32xf32>
}
@@ -147,7 +147,7 @@ gpu.module @test {
//-----
// CHECK-LABEL: store_with_offsets
// CHECK-SAME: [[arg0:%.+]]: ui64
- // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
+ // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
gpu.func @store_with_offsets(%src: ui64) {
%cst = arith.constant dense<[
0, 8, 16, 24, 32, 40, 48, 56,
@@ -160,59 +160,11 @@ gpu.module @test {
%mask = vector.create_mask %c17: vector<32xi1>
%st_vec = arith.constant dense<1023.0>: vector<32xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ xegpu.store %st_vec, %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
-//-----
- // CHECK-LABEL: load_with_offsets_chunk
- // CHECK-SAME: [[arg0:%.+]]: ui64
- // CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<32x4xf32>
- // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
- // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
- // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
- // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
- // CHECK-COUNT-4: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16x2xf32>
- gpu.func @load_with_offsets_chunk(%src: ui64) -> vector<32x4xf32> {
- %cst = arith.constant dense<[
- 0, 8, 16, 24, 32, 40, 48, 56,
- 64, 72, 80, 88, 96, 104, 112, 120,
- 128, 136, 144, 152, 160, 168, 176, 184,
- 192, 200, 208, 216, 224, 232, 240, 248
- ]> : vector<32xindex>
-
- %c17 = arith.constant 17: index
- %mask = vector.create_mask %c17: vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
- gpu.return %ld : vector<32x4xf32>
- }
-
-//-----
- // CHECK-LABEL: store_with_offsets_chunk
- // CHECK-SAME: [[arg0:%.+]]: ui64
- // CHECK: [[cst:%.+]] = arith.constant dense<1.023000e+03> : vector<16x2xf32
- // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
- // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
- // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
- // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
- // CHECK-COUNT-4: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16x2xf32>, ui64, vector<16xindex>, vector<16xi1>
- gpu.func @store_with_offsets_chunk(%src: ui64) {
- %cst = arith.constant dense<[
- 0, 8, 16, 24, 32, 40, 48, 56,
- 64, 72, 80, 88, 96, 104, 112, 120,
- 128, 136, 144, 152, 160, 168, 176, 184,
- 192, 200, 208, 216, 224, 232, 240, 248
- ]> : vector<32xindex>
-
- %c17 = arith.constant 17: index
- %mask = vector.create_mask %c17: vector<32xi1>
-
- %st_vec = arith.constant dense<1023.>: vector<32x4xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
- gpu.return
- }
-
//-----
// CHECK-LABEL: load_nd_store_nd
// CHECK-SAME: [[arg0:%.+]]: memref<256x318xf32>
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 9267f3bbbfbda..3b8626693cfb9 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -313,11 +313,11 @@ gpu.module @test_distribution {
gpu.func @load_gather(%src : memref<?xf16>) {
// CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<32x4xindex>
// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<32x4xi1>
- // CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}>
+ // CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{l1_hint = #xegpu.cache_hint<cached>}>
// CHECK-SAME: : memref<?xf16>, vector<32x4xindex>, vector<32x4xi1> -> vector<32x4xf16>
%offset = arith.constant dense<0> : vector<256x16xindex>
%mask = arith.constant dense<1> : vector<256x16xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>, l1_hint = #xegpu.cache_hint<cached>}
+ %load = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>, l1_hint = #xegpu.cache_hint<cached>}
: memref<?xf16>, vector<256x16xindex>, vector<256x16xi1> -> vector<256x16xf16>
gpu.return
}
@@ -328,31 +328,17 @@ gpu.module @test_distribution {
// CHECK: %[[VAL:.*]] = arith.constant dense<2.550000e+01> : vector<8xf16>
// CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<8xindex>
// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8xi1>
- // CHECK: xegpu.store %[[VAL]], %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<inst_data = [8]>}>
+ // CHECK: xegpu.store %[[VAL]], %[[ARG0]][%[[CST]]], %[[MASK]] <{l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<inst_data = [8]>}>
// CHECK-SAME: : vector<8xf16>, memref<256xf16>, vector<8xindex>, vector<8xi1>
%val = arith.constant dense<25.5> : vector<256xf16>
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
- xegpu.store %val, %dest[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>,
+ xegpu.store %val, %dest[%offset], %mask {layout = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>,
l1_hint = #xegpu.cache_hint<cached>}
: vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
gpu.return
}
- // CHECK-LABEL: @load_with_non_unit_chunk_size
- // CHECK-SAME: %[[ARG0:.*]]: memref<?xf16>
- gpu.func @load_with_non_unit_chunk_size(%src : memref<?xf16>) {
- // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<8xindex>
- // CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8xi1>
- // CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 4 : i64, l1_hint = #xegpu.cache_hint<cached>}>
- // CHECK-SAME: : memref<?xf16>, vector<8xindex>, vector<8xi1> -> vector<8x4xf16>
- %offset = arith.constant dense<0> : vector<256xindex>
- %mask = arith.constant dense<1> : vector<256xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 4, layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [8, 4]>, l1_hint = #xegpu.cache_hint<cached>}
- : memref<?xf16>, vector<256xindex>, vector<256xi1> -> vector<256x4xf16>
- gpu.return
- }
-
// CHECK-LABEL: distribute_load_matrix
// CHECK-SAME: [[arg0:%.+]]: memref<32768xi8, 3>
gpu.func @distribute_load_matrix(%arg0: memref<32768xi8, 3>) {
@@ -760,9 +746,9 @@ gpu.module @test_distribution {
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
- // CHECK: %[[LOAD:.*]] = xegpu.load {{.*}} <{chunk_size = 1 : i64, layout = #xegpu.slice<#xegpu.layout<inst_data = [8, 16]>, dims = [0]>}>
+ // CHECK: %[[LOAD:.*]] = xegpu.load {{.*}} <{layout = #xegpu.slice<#xegpu.layout<inst_data = [8, 16]>, dims = [0]>}>
// CHECK-SAME: memref<4096xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- %3 = xegpu.load %2[%offset], %mask {chunk_size = 1, layout = #xegpu.slice<#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>, dims = [0]> } : memref<4096xf32>, vector<256xindex>, vector<256xi1> -> vector<256xf32>
+ %3 = xegpu.load %2[%offset], %mask {layout = #xegpu.slice<#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>, dims = [0]> } : memref<4096xf32>, vector<256xindex>, vector<256xi1> -> vector<256xf32>
// CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[LOAD]] : vector<32xf32> to vector<32x32xf32>
%4 = vector.broadcast %3 : vector<256xf32> to vector<256x256xf32>
@@ -779,7 +765,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.000000e+00> : vector<1x32xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : vector<1x1x32xindex>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<1x1x32xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0:.*]][%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0:.*]][%[[CST_0]]], %[[CST_1]] : memref<?xf32>, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<0.000000e+00> : vector<1x32xf32>
// CHECK-DAG: %[[LOCAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_2]] [1] : vector<1x1x32xf32> to vector<1x32xf32>
// CHECK-DAG: %[[CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<1x32xf32> to vector<1x1x32xf32>
@@ -796,7 +782,7 @@ gpu.module @test_distribution {
%cst_3 = arith.constant dense<1.0> : vector<1x32xf32>
%offset = arith.constant dense<0> : vector<1x32x32xindex>
%mask = arith.constant dense<true> : vector<1x32x32xi1>
- %14 = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [1, 32, 1], sg_data = [1, 1, 32]>} : memref<?xf32>, vector<1x32x32xindex>, vector<1x32x32xi1> -> vector<1x32x32xf32>
+ %14 = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [1, 32, 1], sg_data = [1, 1, 32]>} : memref<?xf32>, vector<1x32x32xindex>, vector<1x32x32xi1> -> vector<1x32x32xf32>
%15 = vector.multi_reduction <add>, %14, %cst_3 [1] : vector<1x32x32xf32> to vector<1x32xf32>
%anchor = xegpu.convert_layout %15
<{
@@ -852,7 +838,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<1x1xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : vector<1x1x32x32xindex>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<1x1x32x32xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<1x1x32x32xindex>, vector<1x1x32x32xi1> -> vector<1x1x32x32xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST_0]]], %[[CST_1]] : memref<?xf32>, vector<1x1x32x32xindex>, vector<1x1x32x32xi1> -> vector<1x1x32x32xf32>
// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<0.000000e+00> : vector<1x1xf32>
// CHECK-DAG: %[[LOCAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_2]] [2, 3] : vector<1x1x32x32xf32> to vector<1x1xf32>
// CHECK-DAG: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<1x1xf32> to vector<1x1x1x1xf32>
@@ -869,7 +855,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<0.0> : vector<2x2xf32>
%offset = arith.constant dense<0> : vector<2x2x128x128xindex>
%mask = arith.constant dense<true> : vector<2x2x128x128xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [1, 1, 32, 32]>} : memref<?xf32>, vector<2x2x128x128xindex>, vector<2x2x128x128xi1> -> vector<2x2x128x128xf32>
+ %load = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [1, 1, 32, 32]>} : memref<?xf32>, vector<2x2x128x128xindex>, vector<2x2x128x128xi1> -> vector<2x2x128x128xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [2, 3] : vector<2x2x128x128xf32> to vector<2x2xf32>
%anchor = xegpu.convert_layout %reduce
<{
@@ -884,7 +870,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<16x16xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : vector<16x16x32x32xindex>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<16x16x32x32xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<16x16x32x32xindex>, vector<16x16x32x32xi1> -> vector<16x16x32x32xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST_0]]], %[[CST_1]] : memref<?xf32>, vector<16x16x32x32xindex>, vector<16x16x32x32xi1> -> vector<16x16x32x32xf32>
// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<0.000000e+00> : vector<16x16xf32>
// CHECK-DAG: %[[LOCAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_2]] [2, 3] : vector<16x16x32x32xf32> to vector<16x16xf32>
// CHECK-DAG: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<16x16xf32> to vector<16x16x1x1xf32>
@@ -901,7 +887,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<0.0> : vector<32x32xf32>
%offset = arith.constant dense<0> : vector<32x32x128x128xindex>
%mask = arith.constant dense<true> : vector<32x32x128x128xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [16, 16, 32, 32]>} : memref<?xf32>, vector<32x32x128x128xindex>, vector<32x32x128x128xi1> -> vector<32x32x128x128xf32>
+ %load = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [16, 16, 32, 32]>} : memref<?xf32>, vector<32x32x128x128xindex>, vector<32x32x128x128xi1> -> vector<32x32x128x128xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [2, 3] : vector<32x32x128x128xf32> to vector<32x32xf32>
%anchor = xegpu.convert_layout %reduce
<{
@@ -997,7 +983,7 @@ gpu.module @test_distribution {
gpu.func @convert_layout_3D(%arg0: memref<?xf32>) {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : vector<1x32x16xindex>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<true> : vector<1x32x16xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST]]], %[[CST_0]] <{chunk_size = 1 : i64, layout = #xegpu.layout<inst_data = [1, 16, 16]>}> : memref<?xf32>, vector<1x32x16xindex>, vector<1x32x16xi1> -> vector<1x32x16xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST]]], %[[CST_0]] <{layout = #xegpu.layout<inst_data = [1, 16, 16]>}> : memref<?xf32>, vector<1x32x16xindex>, vector<1x32x16xi1> -> vector<1x32x16xf32>
// CHECK-DAG: %[[ALLOCA:.*]] = memref.alloca() : memref<1048576xi8, 3>
// CHECK-DAG: %[[MDESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<1048576xi8, 3> -> !xegpu.mem_desc<8x128x256xf32>
// CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index
@@ -1026,7 +1012,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MDESC]][%[[LOAD_OFF_Z]], %[[LOAD_OFF_Y]], %[[LOAD_OFF_X]]] <{layout = #xegpu.layout<inst_data = [1, 16, 16]>}> : !xegpu.mem_desc<8x128x256xf32>, index, index, index -> vector<1x16x32xf32>
%offset = arith.constant dense<0> : vector<8x128x256xindex>
%mask = arith.constant dense<true> : vector<8x128x256xi1>
- %1 = xegpu.load %arg0[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>} : memref<?xf32>, vector<8x128x256xindex>, vector<8x128x256xi1> -> vector<8x128x256xf32>
+ %1 = xegpu.load %arg0[%offset], %mask {layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>} : memref<?xf32>, vector<8x128x256xindex>, vector<8x128x256xi1> -> vector<8x128x256xf32>
%2 = xegpu.convert_layout %1 <{input_layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>,
target_layout = #xegpu.layout<sg_layout = [8, 8, 8], sg_data = [1, 16, 32], inst_data = [1, 16, 16]>}> : vector<8x128x256xf32>
%anchor = xegpu.convert_layout %2
More information about the Mlir-commits
mailing list