[Mlir-commits] [mlir] [mlir][XeGPU] Distribute coalesced gather/scatter via lane_data (PR #201688)
Md Abdullah Shahneous Bari
llvmlistbot at llvm.org
Tue Jul 21 08:03:29 PDT 2026
https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/201688
>From 178ee5add97f929794f053dd56c5904e0b42ee49 Mon Sep 17 00:00:00 2001
From: Claude <noreply at anthropic.com>
Date: Mon, 22 Jun 2026 15:32:10 +0000
Subject: [PATCH 1/2] [mlir][xegpu] Remove inferred chunk_size; rename
contiguity attr to chunk_size
Two changes to the xegpu.load (LoadGatherOp) / xegpu.store (StoreScatterOp)
gather/scatter ops:
1. Remove the old optional `chunk_size` I64 attribute that duplicated
information already carried by the operand/result types. The payload chunk
size (contiguous elements per work item) is derived from the value and mask
types via a shared helper `xegpu::getGatherScatterPayloadChunk` in
XeGPUUtils; the op verifier keeps a file-local equivalent. Consumers in
propagate-layout, sg-to-lane, unroll, and the operand-layout util now call
the helper instead of an op accessor.
2. Rename the `contiguity` attribute (added by #201684) to `chunk_size`, so the
name `chunk_size` now denotes the offsets-contiguity hint: an optional I64
whose value must be >= 2 and must divide the innermost offsets dimension
(verified by `isValidChunkSize`). The AxisInfo-based analysis
(runContiguityAnalysis) stamps it via the generated setChunkSize accessor;
the test driver and lit tests are updated accordingly.
Because the two uses no longer share the `getChunkSize()` name, there is no
collision: `getChunkSize()` is exclusively the generated accessor for the
renamed attribute.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 37 ++++----
.../Dialect/XeGPU/Transforms/Transforms.h | 4 +-
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 12 +++
.../VectorToXeGPU/VectorToXeGPU.cpp | 4 -
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 86 ++++++++++++-------
.../Transforms/XeGPUContiguityAnalysis.cpp | 8 +-
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 6 +-
.../Transforms/XeGPUSgToLaneDistribute.cpp | 17 ++--
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 20 ++---
.../Transforms/XeGPUWgToSgDistribute.cpp | 16 ++--
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 20 ++++-
.../Dialect/XeGPU/contiguity-analysis.mlir | 38 ++++----
mlir/test/Dialect/XeGPU/invalid.mlir | 53 +++---------
mlir/test/Dialect/XeGPU/ops.mlir | 34 ++++----
.../XeGPU/propagate-layout-subgroup.mlir | 8 +-
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 4 +-
.../XeGPU/sg-to-lane-distribute-unit.mlir | 12 +--
.../Dialect/XeGPU/sg-to-lane-distribute.mlir | 16 ++--
.../test-xegpu-coalesce-gather-scatter.mlir | 2 +-
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 40 ++++-----
.../Dialect/XeGPU/xegpu-unroll-patterns.mlir | 16 ++--
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 32 +++----
.../lib/Dialect/XeGPU/TestXeGPUTransforms.cpp | 10 +--
23 files changed, 253 insertions(+), 242 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 7f8389a6acc47..ee551e1065109 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -731,13 +731,17 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
mask is a vector of size equal to the subgroup size, or 1 at lane level.
scalar mask is also valid for lane level.
- - `chunk_size`: [optional] represents contiguous number of elements to load from per work item.
-
- `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache.
- `layout`: [optional] Describes the expected layout of the `tensor_desc` operand or the result
of load. Only valid at workgroup and subgroup levels.
+ - `chunk_size`: [optional] An I64 attribute describing the contiguity of the
+ `offsets`: the innermost `offsets` dimension is contiguous in runs of
+ `chunk_size` elements (so `chunk_size` must be >= 2 and must divide that
+ dimension). It is a hint used by layout propagation / coalescing; it does
+ not change the loaded data. Only valid on the vector-offsets form.
+
Results:
- `res`: represents loaded data
@@ -756,8 +760,8 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
```
Example 2 (lane level):
- lane level only accepts the offsets variant. chunk_size can be inferred from result
- type. In this example, chunk_size is 8.
+ lane level only accepts the offsets variant. The number of contiguous elements
+ loaded per work item is inferred from the result type; in this example it is 8.
```mlir
%2 = xegpu.load %1[%2], %0 <{l1_hint = #xegpu.cache_hint<cached>,
l2_hint = #xegpu.cache_hint<uncached>,
@@ -769,12 +773,12 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
let arguments = (ins XeGPU_GatherScatterBaseAddrType:$source,
AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
- AnyTypeOf<[XeGPU_MaskType, I1]>:$mask, OptionalAttr<I64Attr>:$chunk_size,
+ AnyTypeOf<[XeGPU_MaskType, I1]>:$mask,
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
OptionalAttr<DistributeLayoutAttr>:$layout,
- OptionalAttr<I64Attr>:$contiguity);
+ OptionalAttr<I64Attr>:$chunk_size);
let results = (outs XeGPU_ValueOrScalarType:$value);
let extraClassDeclaration = extraBaseClassDeclaration # [{
@@ -815,13 +819,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,
@@ -862,13 +864,18 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
mask is a vector of size equal to the subgroup size, or 1 at lane level.
scalar mask is also valid for lane level.
- - `chunk_size`: [optional] represents contiguous number of elements to store to per work item.
-
- `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache.
- `layout`: [optional] Describes the expected layout of the `tensor_desc` operand or the value
to be stored. Only valid at workgroup and subgroup levels.
+ - `chunk_size`: [optional] An I64 attribute describing the contiguity of the
+ `offsets`, analogously to `load`: the innermost `offsets` dimension is
+ contiguous in runs of `chunk_size` elements (so `chunk_size` must be >= 2
+ and must divide that dimension). It is a hint used by layout propagation /
+ coalescing; it does not change the stored data. Only valid on the
+ vector-offsets form.
+
Example 1 (Subgroup level):
A variant accepts memref as base pointer and an offset.
@@ -886,8 +893,8 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
```
Example 2 (Lane level):
- Lane level IR only accepts the offsets variant. chunk_size can be inferred from value
- type. In this example, chunk_size is 8.
+ Lane level IR only accepts the offsets variant. The number of contiguous elements
+ stored per work item is inferred from the value type; in this example it is 8.
```mlir
xegpu.store %0, %1[%2], %3 <{l1_hint = #xegpu.cache_hint<uncached>,
l2_hint = #xegpu.cache_hint<write_back>,
@@ -900,12 +907,12 @@ 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,
OptionalAttr<DistributeLayoutAttr>:$layout,
- OptionalAttr<I64Attr>:$contiguity);
+ OptionalAttr<I64Attr>:$chunk_size);
let extraClassDeclaration = extraBaseClassDeclaration#[{
Type getDestType() {
@@ -945,13 +952,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/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index 388bd6145df21..7f707d67ebd2a 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -93,11 +93,11 @@ void populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
//===----------------------------------------------------------------------===//
/// Run the AxisInfo-based contiguity analysis over `root` and stamp a
-/// `contiguity` attribute on every `xegpu.load` / `xegpu.store` whose
+/// `chunk_size` attribute on every `xegpu.load` / `xegpu.store` whose
/// offsets are contiguous (in runs of >= 2) along the innermost dimension.
/// The stamped value is the inner-dim contiguity; it is a target-independent
/// property consumed downstream (e.g. to derive a `lane_data` split). Ops that
-/// already carry a `contiguity` attribute are left untouched.
+/// already carry a `chunk_size` attribute are left untouched.
void runContiguityAnalysis(Operation *root);
/// Collect a set of patterns to unroll xegpu operations to a smaller shapes.
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 0125dfc44196b..33dfc1cea4ea0 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -41,6 +41,18 @@ namespace xegpu {
/// Flatten a set of ValueRange into a single SmallVector<Value>
SmallVector<Value> flattenValues(ArrayRef<ValueRange> values);
+/// Infer the payload chunk size (number of contiguous elements per work item)
+/// of a gather/scatter op (`xegpu.load` / `xegpu.store`) from its value/result
+/// and mask types. The mask carries one element per lane, so the payload is the
+/// value with the chunk dimension removed: the chunk size is the trailing value
+/// dimension whenever the value has more elements than the mask, and 1
+/// otherwise.
+/// - scalar value -> 1
+/// - value vector<LxC...>, mask vector<L...> -> C (trailing dim)
+/// - 1D value vector<N...>, scalar/size-1 mask -> N (lane-level chunk)
+/// - value and mask with matching element counts -> 1 (one elem per lane)
+int64_t getGatherScatterPayloadChunk(VectorType valueTy, Type maskTy);
+
/// If tensor descriptor has a layout attribute it is used in SIMT mode.
/// In this mode, the distributed vector shape is determined as follows:
/// Definitions:
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 9a994e87697f6..feab0eb51ec27 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -502,7 +502,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{},
@@ -537,7 +536,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{},
@@ -793,7 +791,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{},
@@ -828,7 +825,6 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
xegpu::StoreScatterOp::create(rewriter, loc, scatterOp.getValueToStore(),
flatMemref, localOffsets, scatterOp.getMask(),
- /*chunk_size=*/IntegerAttr{},
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 2ffe883eb0d9a..c0178035b59d5 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -61,6 +61,30 @@ static bool isWriteHintOrNone(const CachePolicyAttr &attr) {
kind == CachePolicy::WRITE_BACK || kind == CachePolicy::WRITE_THROUGH;
}
+// Infer the chunk size (number of contiguous elements per work item) of a
+// gather/scatter op from its value/result and mask types. The chunk size is no
+// longer carried as an attribute; it is fully determined by the types. The
+// mask carries one element per lane, so it is the value with the chunk
+// dimension removed. The chunk size is therefore the trailing value dimension
+// whenever the value has more elements than the mask, and 1 otherwise:
+// - scalar value -> 1
+// - value vector<LxC...>, mask vector<L...> -> C (trailing dim)
+// - 1D value vector<N...>, scalar/size-1 mask -> N (lane-level chunk)
+// - value and mask with matching element counts -> 1 (one elem per lane)
+static int64_t inferGatherScatterChunkSize(VectorType valueTy, Type maskTy) {
+ if (!valueTy)
+ return 1;
+ auto maskVecTy = dyn_cast<VectorType>(maskTy);
+ int64_t maskSize = maskVecTy ? maskVecTy.getNumElements() : 1;
+ int64_t valueSize = valueTy.getNumElements();
+ if (valueTy.getRank() >= 2)
+ return maskSize == valueSize ? 1 : valueTy.getShape().back();
+ // 1D value: a size-1 (or scalar) mask denotes a single work item performing a
+ // chunked load/store, so the whole vector is the chunk; a wider mask denotes
+ // one element per lane (chunk size 1).
+ return maskSize == 1 ? valueSize : 1;
+}
+
static LogicalResult
isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
VectorType valueTy, int64_t chunkSize,
@@ -113,23 +137,23 @@ isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
return success();
}
-// Validates the `contiguity` attribute against the op's offsets type: the
+// Validates the `chunk_size` attribute against the op's offsets type: the
// innermost offsets dimension is contiguous in runs of `size`, so `size` must
// be >= 2 and must divide that dimension.
static LogicalResult
-isValidContiguity(std::optional<uint64_t> contiguity, Type offsetsTy,
- function_ref<InFlightDiagnostic()> emitError) {
- if (!contiguity)
+isValidChunkSize(std::optional<uint64_t> chunkSize, Type offsetsTy,
+ function_ref<InFlightDiagnostic()> emitError) {
+ if (!chunkSize)
return success();
auto offsetsVecTy = dyn_cast<VectorType>(offsetsTy);
if (!offsetsVecTy)
- return emitError() << "contiguity requires vector offsets (one per lane).";
- int64_t size = static_cast<int64_t>(*contiguity);
+ return emitError() << "chunk_size requires vector offsets (one per lane).";
+ int64_t size = static_cast<int64_t>(*chunkSize);
int64_t inner = offsetsVecTy.getShape().back();
if (size < 2)
- return emitError() << "contiguity = " << size << " (must be >= 2)";
+ return emitError() << "chunk_size = " << size << " (must be >= 2)";
if (inner % size != 0)
- return emitError() << "contiguity = " << size
+ return emitError() << "chunk_size = " << size
<< " (must divide the innermost offsets dim " << inner
<< ")";
return success();
@@ -592,7 +616,7 @@ LogicalResult LoadGatherOp::verify() {
return emitOpError("invalid l3_hint: ") << getL3HintAttr();
auto srcTy = getSourceType();
- uint64_t chunkSize = static_cast<int64_t>(getChunkSize().value_or(1));
+ int64_t chunkSize = inferGatherScatterChunkSize(valueTy, maskTy);
auto memTy = dyn_cast<MemRefType>(srcTy);
if (memTy && (getElementType() != memTy.getElementType()))
@@ -604,8 +628,8 @@ LogicalResult LoadGatherOp::verify() {
}
auto offsetsTy = getOffsets().getType();
- if (failed(isValidContiguity(getContiguity(), offsetsTy,
- [&]() { return emitOpError(); })))
+ if (failed(isValidChunkSize(getChunkSize(), offsetsTy,
+ [&]() { return emitOpError(); })))
return failure();
return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
[&]() { return emitOpError(); });
@@ -614,7 +638,7 @@ LogicalResult LoadGatherOp::verify() {
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();
@@ -623,15 +647,14 @@ 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,
- /*contiguity=*/nullptr);
+ build(builder, state, valueType, source, offset, mask, l1_hint, l2_hint,
+ l3_hint, /*anchor_layout=*/nullptr, /*chunk_size=*/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) {
@@ -641,8 +664,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, /*chunk_size=*/nullptr);
}
//===----------------------------------------------------------------------===//
@@ -662,7 +685,7 @@ LogicalResult StoreScatterOp::verify() {
return emitOpError("invalid l3_hint: ") << getL3HintAttr();
auto destTy = getDestType();
- uint64_t chunkSize = static_cast<int64_t>(getChunkSize().value_or(1));
+ int64_t chunkSize = inferGatherScatterChunkSize(valueTy, maskTy);
auto memTy = dyn_cast<MemRefType>(destTy);
if (memTy && (getElementType() != memTy.getElementType()))
@@ -674,8 +697,8 @@ LogicalResult StoreScatterOp::verify() {
}
auto offsetsTy = getOffsets().getType();
- if (failed(isValidContiguity(getContiguity(), offsetsTy,
- [&]() { return emitOpError(); })))
+ if (failed(isValidChunkSize(getChunkSize(), offsetsTy,
+ [&]() { return emitOpError(); })))
return failure();
return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
[&]() { return emitOpError(); });
@@ -684,7 +707,6 @@ LogicalResult StoreScatterOp::verify() {
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) {
@@ -695,15 +717,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, /*chunk_size=*/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());
@@ -711,8 +735,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, /*chunk_size=*/nullptr);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
index 3cc8152561681..64f1d6b992e4f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
@@ -872,7 +872,7 @@ using ::mlir::xegpu::detail::axis_dataflow::AxisInfoLattice;
// Analysis driver.
//===----------------------------------------------------------------------===//
-/// Stamp a `contiguity` attribute on `op` recording the inner-dim contiguity
+/// Stamp a `chunk_size` attribute on `op` recording the inner-dim contiguity
/// computed by the analysis. The contiguity is a target-independent property
/// of the offsets.
template <typename OpTy>
@@ -880,9 +880,9 @@ static void analyzeAndStampContiguity(OpTy op, DataFlowSolver &solver) {
auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
if (!offsetsTy || offsetsTy.getNumElements() <= 1)
return;
- // A pre-existing `contiguity` (user-authored, or stamped by an earlier run)
+ // A pre-existing `chunk_size` (user-authored, or stamped by an earlier run)
// takes precedence; leave it untouched so the analysis is idempotent.
- if (op.getContiguity())
+ if (op.getChunkSize())
return;
const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
if (!lat || !lat->getValue().isInitialized())
@@ -898,7 +898,7 @@ static void analyzeAndStampContiguity(OpTy op, DataFlowSolver &solver) {
--contiguity;
if (contiguity < 2)
return;
- op.setContiguity(contiguity);
+ op.setChunkSize(contiguity);
}
} // namespace
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 1681d295ae0ff..d1e27094564cb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1218,7 +1218,8 @@ void LayoutInfoPropagation::visitLoadGatherOp(
if (!uArch)
return;
VectorType resVecTy = load.getValueType();
- int chunkSize = load.getChunkSize().value_or(1);
+ int chunkSize =
+ xegpu::getGatherScatterPayloadChunk(resVecTy, load.getMaskType());
LayoutInfo resLayoutInfo = results[0]->getValue();
if (!resLayoutInfo.isAssigned())
@@ -1283,7 +1284,8 @@ void LayoutInfoPropagation::visitStoreScatterOp(
if (!uArch)
return;
VectorType srcVecTy = storeScatter.getValueType();
- int chunkSize = storeScatter.getChunkSize().value_or(1);
+ int chunkSize =
+ xegpu::getGatherScatterPayloadChunk(srcVecTy, storeScatter.getMaskType());
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 874487da10b30..2ac3d1ff2bba5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -543,7 +543,8 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
return failure();
// Check that leading dimensions are unit.
- int chunkSize = op.getChunkSize().value_or(1);
+ int chunkSize =
+ xegpu::getGatherScatterPayloadChunk(origResultTy, op.getMaskType());
int effectiveVecRank = (chunkSize == 1) ? 1 : 2;
ArrayRef<int64_t> shape = origResultTy.getShape();
if (llvm::any_of(
@@ -581,8 +582,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)
@@ -1071,7 +1072,8 @@ struct SgToLaneStoreScatter
return failure();
// Check that all leading dimensions are unit dimensions.
- int chunkSize = op.getChunkSize().value_or(1);
+ int chunkSize =
+ xegpu::getGatherScatterPayloadChunk(origValueTy, op.getMaskType());
int effectiveVecRank = (chunkSize == 1) ? 1 : 2;
ArrayRef<int64_t> shape = origValueTy.getShape();
if (llvm::any_of(shape.take_front(origValueTy.getRank() - effectiveVecRank),
@@ -1112,10 +1114,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,
- /*contiguity=*/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 74c358cef90df..3841ae3070490 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -699,11 +699,8 @@ 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();
- }
+ int64_t chunkSize =
+ xegpu::getGatherScatterPayloadChunk(valueTy, mask.getType());
// Unroll mask and offsets with correct shape
VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
@@ -721,7 +718,6 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
targetMaskShape.pop_back();
int64_t blockedChunkSize = targetShape->back();
int64_t numNewChunks = chunkSize / blockedChunkSize;
- chunkSize = blockedChunkSize;
convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
convertedOffsetTypes = getUnrolledTypes(offsetsTy, targetMaskShape);
@@ -763,8 +759,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);
@@ -793,11 +788,8 @@ 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();
- }
+ int64_t chunkSize =
+ xegpu::getGatherScatterPayloadChunk(valueTy, mask.getType());
SmallVector<int64_t> targetMaskShape(*targetShape);
VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
@@ -812,7 +804,6 @@ struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
targetMaskShape.pop_back();
int64_t blockedChunkSize = targetShape->back();
int64_t numNewChunks = chunkSize / blockedChunkSize;
- chunkSize = blockedChunkSize;
convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
convertedOffsetTypes = getUnrolledTypes(offsetsTy, targetMaskShape);
@@ -859,7 +850,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 9ca6b3c2b0272..8cd1b1ad0d3f1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -824,14 +824,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);
@@ -871,16 +869,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 9620e21f9bfdf..92759fc7a69a6 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -40,6 +40,20 @@ SmallVector<Value> xegpu::flattenValues(ArrayRef<ValueRange> values) {
return result;
}
+int64_t xegpu::getGatherScatterPayloadChunk(VectorType valueTy, Type maskTy) {
+ if (!valueTy)
+ return 1;
+ auto maskVecTy = dyn_cast<VectorType>(maskTy);
+ int64_t maskSize = maskVecTy ? maskVecTy.getNumElements() : 1;
+ int64_t valueSize = valueTy.getNumElements();
+ if (valueTy.getRank() >= 2)
+ return maskSize == valueSize ? 1 : valueTy.getShape().back();
+ // 1D value: a size-1 (or scalar) mask denotes a single work item performing a
+ // chunked load/store, so the whole vector is the chunk; a wider mask denotes
+ // one element per lane (chunk size 1).
+ return maskSize == 1 ? valueSize : 1;
+}
+
FailureOr<VectorType>
mlir::xegpu::getDistributedVectorType(xegpu::TensorDescType tdescTy) {
auto layout = llvm::dyn_cast_if_present<LayoutAttr>(tdescTy.getLayout());
@@ -237,7 +251,8 @@ xegpu::getDistributeLayoutAttr(const OpOperand &opr) {
if (isa<xegpu::StoreScatterOp>(op)) {
xegpu::StoreScatterOp store(op);
- int chunkSize = store.getChunkSize().value_or(1);
+ int chunkSize = xegpu::getGatherScatterPayloadChunk(store.getValueType(),
+ store.getMaskType());
if (layout && idx >= 2 && chunkSize > 1)
return layout.dropDims(llvm::to_vector(
llvm::seq<int64_t>(layout.getRank() - 1, layout.getRank())));
@@ -245,7 +260,8 @@ xegpu::getDistributeLayoutAttr(const OpOperand &opr) {
}
if (isa<xegpu::LoadGatherOp>(op)) {
xegpu::LoadGatherOp load(op);
- int chunkSize = load.getChunkSize().value_or(1);
+ int chunkSize = xegpu::getGatherScatterPayloadChunk(load.getValueType(),
+ load.getMaskType());
if (layout && idx >= 1 && chunkSize > 1)
return layout.dropDims(llvm::to_vector(
llvm::seq<int64_t>(layout.getRank() - 1, layout.getRank())));
diff --git a/mlir/test/Dialect/XeGPU/contiguity-analysis.mlir b/mlir/test/Dialect/XeGPU/contiguity-analysis.mlir
index f223b325c2f9e..df9f7fe69c2e2 100644
--- a/mlir/test/Dialect/XeGPU/contiguity-analysis.mlir
+++ b/mlir/test/Dialect/XeGPU/contiguity-analysis.mlir
@@ -1,7 +1,7 @@
// RUN: mlir-opt -split-input-file \
// RUN: -test-xegpu-coalesce-gather-scatter="analyze-only=true" %s | FileCheck %s
-// Contiguity analysis: stamps the `contiguity` attribute on gather/scatter ops
+// Contiguity analysis: stamps the `chunk_size` attribute on gather/scatter ops
// whose offsets are contiguous (runs of >= 2) along the innermost dimension.
// The stamped value is the inner-dim contiguity, rounded down to a divisor of
// the inner extent. The analysis is target-independent and mask-independent;
@@ -11,7 +11,7 @@
// 1-D vector.step -> stride-1, fully contiguous over the 32-element inner dim.
// CHECK-LABEL: func.func @load_step_offsets(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: <{chunk_size = 32 : i64}>
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
@@ -25,7 +25,7 @@ func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
// Dense stride-1 constant offsets -> contiguity 32.
// CHECK-LABEL: func.func @load_dense_ap_offsets(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: <{chunk_size = 32 : i64}>
func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
%offsets = arith.constant dense<[
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
@@ -41,7 +41,7 @@ func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
// Stride-4 offsets: not contiguous, no attribute stamped.
// CHECK-LABEL: func.func @load_stride4_no_attr(
// CHECK: xegpu.load
-// CHECK-NOT: contiguity
+// CHECK-NOT: chunk_size
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_stride4_no_attr(%ptr: i64) -> vector<32xf32> {
%c4 = arith.constant 4 : index
@@ -58,7 +58,7 @@ func.func @load_stride4_no_attr(%ptr: i64) -> vector<32xf32> {
// All-equal offsets: inner dim is constant, not contiguous -> no attribute.
// CHECK-LABEL: func.func @load_broadcast_offsets_no_attr(
// CHECK: xegpu.load
-// CHECK-NOT: contiguity
+// CHECK-NOT: chunk_size
func.func @load_broadcast_offsets_no_attr(%ptr: i64) -> vector<32xf32> {
%offsets = arith.constant dense<0> : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
@@ -72,7 +72,7 @@ func.func @load_broadcast_offsets_no_attr(%ptr: i64) -> vector<32xf32> {
// (the mask check is a consumer concern).
// CHECK-LABEL: func.func @load_partial_mask(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: <{chunk_size = 32 : i64}>
func.func @load_partial_mask(%ptr: i64, %mask: vector<32xi1>) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
%v = xegpu.load %ptr[%offsets], %mask
@@ -84,7 +84,7 @@ func.func @load_partial_mask(%ptr: i64, %mask: vector<32xi1>) -> vector<32xf32>
// Store with vector.step offsets -> contiguity 32 on the store.
// CHECK-LABEL: func.func @store_step_offsets(
// CHECK: xegpu.store
-// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: <{chunk_size = 32 : i64}>
func.func @store_step_offsets(%ptr: i64, %v: vector<32xf32>) {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
@@ -97,7 +97,7 @@ func.func @store_step_offsets(%ptr: i64, %v: vector<32xf32>) {
// 2-D leading-unit dim: contiguity measured on the inner dim -> 32.
// CHECK-LABEL: func.func @load_2d_leading_unit(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: <{chunk_size = 32 : i64}>
func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
%step = vector.step : vector<32xindex>
%offsets = vector.shape_cast %step : vector<32xindex> to vector<1x32xindex>
@@ -111,7 +111,7 @@ func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
// True 2-D dense AP: each row stride-1 over 16 -> contiguity 16.
// CHECK-LABEL: func.func @load_2d_dense_ap(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 16 : i64}>
+// CHECK-SAME: <{chunk_size = 16 : i64}>
func.func @load_2d_dense_ap(%ptr: i64) -> vector<2x16xf32> {
%offsets = arith.constant dense<[
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
@@ -127,7 +127,7 @@ func.func @load_2d_dense_ap(%ptr: i64) -> vector<2x16xf32> {
// 2-D dense values whose inner row is not a stride-1 AP: no attribute.
// CHECK-LABEL: func.func @load_2d_non_ap(
// CHECK: xegpu.load
-// CHECK-NOT: contiguity
+// CHECK-NOT: chunk_size
func.func @load_2d_non_ap(%ptr: i64) -> vector<2x16xf32> {
%offsets = arith.constant dense<[
[0, 1, 2, 3, 4, 5, 6, 7, 100, 9, 10, 11, 12, 13, 14, 15],
@@ -146,7 +146,7 @@ func.func @load_2d_non_ap(%ptr: i64) -> vector<2x16xf32> {
// exercised even if the solver later folds all-constant arith ops.
// CHECK-LABEL: func.func @load_divui_recovers(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: <{chunk_size = 32 : i64}>
func.func @load_divui_recovers(%ptr: i64) -> vector<32xf32> {
%step = vector.step : vector<32xindex>
%c2 = arith.constant dense<2> : vector<32xindex>
@@ -162,7 +162,7 @@ func.func @load_divui_recovers(%ptr: i64) -> vector<32xf32> {
// `divui` by a constant that does not divide the inner stride: not recovered.
// CHECK-LABEL: func.func @load_divui_non_divisor(
// CHECK: xegpu.load
-// CHECK-NOT: contiguity
+// CHECK-NOT: chunk_size
func.func @load_divui_non_divisor(%ptr: i64) -> vector<16xf32> {
%step = vector.step : vector<16xindex>
%c2 = arith.constant dense<2> : vector<16xindex>
@@ -180,7 +180,7 @@ func.func @load_divui_non_divisor(%ptr: i64) -> vector<16xf32> {
// not contiguous -> no attribute.
// CHECK-LABEL: func.func @load_remui_inner_uniform(
// CHECK: xegpu.load
-// CHECK-NOT: contiguity
+// CHECK-NOT: chunk_size
func.func @load_remui_inner_uniform(%ptr: i64) -> vector<16xf32> {
%step = vector.step : vector<16xindex>
%c2 = arith.constant dense<2> : vector<16xindex>
@@ -196,7 +196,7 @@ func.func @load_remui_inner_uniform(%ptr: i64) -> vector<16xf32> {
// `shli` then `shrui` cancel: (step << 1) >> 1 -> stride 1 -> contiguity 32.
// CHECK-LABEL: func.func @load_shli_then_shrui(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: <{chunk_size = 32 : i64}>
func.func @load_shli_then_shrui(%ptr: i64) -> vector<32xf32> {
%step = vector.step : vector<32xindex>
%k = arith.constant dense<1> : vector<32xindex>
@@ -212,7 +212,7 @@ func.func @load_shli_then_shrui(%ptr: i64) -> vector<32xf32> {
// `shli` alone scales the stride to 2: not contiguous -> no attribute.
// CHECK-LABEL: func.func @load_shli_unchanged(
// CHECK: xegpu.load
-// CHECK-NOT: contiguity
+// CHECK-NOT: chunk_size
func.func @load_shli_unchanged(%ptr: i64) -> vector<32xf32> {
%step = vector.step : vector<32xindex>
%k = arith.constant dense<1> : vector<32xindex>
@@ -228,7 +228,7 @@ func.func @load_shli_unchanged(%ptr: i64) -> vector<32xf32> {
// contiguity 32.
// CHECK-LABEL: func.func @load_select_two_aps(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: <{chunk_size = 32 : i64}>
func.func @load_select_two_aps(%ptr: i64, %cond: i1) -> vector<32xf32> {
%step = vector.step : vector<32xindex>
%a = arith.constant dense<0> : vector<32xindex>
@@ -243,14 +243,14 @@ func.func @load_select_two_aps(%ptr: i64, %cond: i1) -> vector<32xf32> {
}
// -----
-// A pre-existing `contiguity` takes precedence: the analysis leaves it alone.
+// A pre-existing `chunk_size` takes precedence: the analysis leaves it alone.
// CHECK-LABEL: func.func @user_attr_preserved(
// CHECK: xegpu.load
-// CHECK-SAME: <{contiguity = 2 : i64}>
+// CHECK-SAME: <{chunk_size = 2 : i64}>
func.func @user_attr_preserved(%ptr: i64) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
- %v = xegpu.load %ptr[%offsets], %mask <{contiguity = 2 : i64}>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 2 : i64}>
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index d5d4950fe7d7e..c610ded4a5b50 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -233,7 +233,7 @@ 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}>
+ %2 = xegpu.load %src[%offsets], %mask
: memref<?xf32>, vector<4xindex>, vector<8xi1> -> vector<4x2xf32>
return
}
@@ -242,8 +242,8 @@ func.func @load_gather_vc_3(%src: memref<?xf32>) {
func.func @load_gather_simt_1(%src: memref<?xf32>) {
%0 = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
%1 = arith.constant dense<1>: vector<4xi1>
- // expected-error at +1 {{value elements must match chunk size}}
- %2 = xegpu.load %src[%0], %1 <{chunk_size = 2}>
+ // expected-error at +1 {{Mask should match value except the chunk size dim}}
+ %2 = xegpu.load %src[%0], %1
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<6xf32>
return
}
@@ -265,7 +265,7 @@ func.func @store_scatter_vc_3(%dst: memref<?xf32>) {
%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}>
+ xegpu.store %2, %dst[%0], %1
: vector<4x2xf32>, memref<?xf32>, vector<4xindex>, vector<8xi1>
return
}
@@ -275,8 +275,8 @@ func.func @store_scatter_simt_1(%dst: memref<?xf32>) {
%0 = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
%1 = arith.constant dense<1>: vector<4xi1>
%2 = arith.constant dense<2.9>: vector<6xf32>
- // expected-error at +1 {{value elements must match chunk size}}
- xegpu.store %2, %dst[%0], %1 <{chunk_size = 2}>
+ // expected-error at +1 {{Mask should match value except the chunk size dim}}
+ xegpu.store %2, %dst[%0], %1
: vector<6xf32>, memref<?xf32>, vector<4xindex>, vector<4xi1>
return
}
@@ -316,26 +316,6 @@ func.func @load_gather_offset_sg(%src: memref<?xf16>) {
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>
- return
-}
-
// -----
func.func @store_scatter_offset_wi_2(%src: memref<4x4xf16>) {
%val = arith.constant dense<2.9>: vector<4xf16>
@@ -363,16 +343,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
}
@@ -381,7 +352,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
}
@@ -792,16 +763,16 @@ func.func @dpas_mx_scale_b_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b
// -----
func.func @contiguity_too_small(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
- // expected-error at +1 {{contiguity = 1 (must be >= 2)}}
- %val = xegpu.load %src[%offset], %mask <{contiguity = 1 : i64}>
+ // expected-error at +1 {{chunk_size = 1 (must be >= 2)}}
+ %val = xegpu.load %src[%offset], %mask <{chunk_size = 1 : i64}>
: i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
return
}
// -----
func.func @contiguity_does_not_divide(%src: i64, %offset: vector<6xindex>, %mask: vector<6xi1>) {
- // expected-error at +1 {{contiguity = 4 (must divide the innermost offsets dim 6)}}
- %val = xegpu.load %src[%offset], %mask <{contiguity = 4 : i64}>
+ // expected-error at +1 {{chunk_size = 4 (must divide the innermost offsets dim 6)}}
+ %val = xegpu.load %src[%offset], %mask <{chunk_size = 4 : i64}>
: i64, vector<6xindex>, vector<6xi1> -> vector<6xf32>
return
}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index 6cffa3eec369b..2b82d426eba39 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -377,8 +377,8 @@ gpu.func @prefetch_nd_3d(%src: memref<4x8x16xf16>) {
// CHECK: gpu.func @simt_load_4(%[[arg0:.*]]: memref<256xf16>, %[[arg1:.*]]: vector<1xindex>, %[[arg2:.*]]: vector<1xi1>) {
gpu.func @simt_load_4(%arg0: memref<256xf16>, %arg1: vector<1xindex>, %arg2: vector<1xi1>) {
- // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{chunk_size = 8 : i64}> : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
- %0 = xegpu.load %arg0[%arg1], %arg2 <{chunk_size = 8 : i64}> : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
+ // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
+ %0 = xegpu.load %arg0[%arg1], %arg2 : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
gpu.return
}
@@ -391,8 +391,8 @@ gpu.func @simt_load_5(%arg0: memref<256xf16>, %arg1: vector<1xindex>, %arg2: vec
// CHECK: gpu.func @simt_load_6(%[[arg0:.*]]: memref<256xf16>, %[[arg1:.*]]: index, %[[arg2:.*]]: i1) {
gpu.func @simt_load_6(%arg0: memref<256xf16>, %arg1: index, %arg2: i1) {
- // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{chunk_size = 8 : i64}> : memref<256xf16>, index, i1 -> vector<8xf16>
- %0 = xegpu.load %arg0[%arg1], %arg2 <{chunk_size = 8 : i64}> : memref<256xf16>, index, i1 -> vector<8xf16>
+ // CHECK: %0 = xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] : memref<256xf16>, index, i1 -> vector<8xf16>
+ %0 = xegpu.load %arg0[%arg1], %arg2 : memref<256xf16>, index, i1 -> vector<8xf16>
gpu.return
}
@@ -407,23 +407,23 @@ gpu.func @simt_load_7(%arg0: memref<256xf16>, %arg1: index, %arg2: i1) {
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>}>
+ //CHECK: %[[R1:.*]] = xegpu.load %arg0[%cst], %cst_0 <{l1_hint = #xegpu.cache_hint<cached>}> : memref<?xf16>, vector<4xindex>, vector<4xi1> -> vector<4x2xf16>
+ %val = xegpu.load %src[%offset], %mask <{l1_hint = #xegpu.cache_hint<cached>}>
: memref<?xf16>, vector<4xindex>, vector<4xi1> -> vector<4x2xf16>
gpu.return
}
// CHECK: gpu.func @simt_store_4(%[[arg0:.*]]: vector<8xf16>, %[[arg1:.*]]: memref<256xf16>, %[[arg2:.*]]: vector<1xindex>, %[[arg3:.*]]: vector<1xi1>) {
gpu.func @simt_store_4(%arg0: vector<8xf16>, %arg1: memref<256xf16>, %arg2: vector<1xindex>, %arg3: vector<1xi1>) {
- // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
- xegpu.store %arg0, %arg1[%arg2], %arg3 <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
+ xegpu.store %arg0, %arg1[%arg2], %arg3 : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
gpu.return
}
// CHECK: gpu.func @simt_store_5(%[[arg0:.*]]: vector<8xf16>, %[[arg1:.*]]: memref<256xf16>, %[[arg2:.*]]: index, %[[arg3:.*]]: i1) {
gpu.func @simt_store_5(%arg0: vector<8xf16>, %arg1: memref<256xf16>, %arg2: index, %arg3: i1) {
- // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, index, i1
- xegpu.store %arg0, %arg1[%arg2], %arg3 <{chunk_size = 8 : i64}> : vector<8xf16>, memref<256xf16>, index, i1
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] : vector<8xf16>, memref<256xf16>, index, i1
+ xegpu.store %arg0, %arg1[%arg2], %arg3 : vector<8xf16>, memref<256xf16>, index, i1
gpu.return
}
@@ -446,25 +446,25 @@ 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>}>
+ //CHECK: xegpu.store %[[R0:.*]], %arg0[%cst_0], %cst_1 <{l1_hint = #xegpu.cache_hint<cached>}> : vector<4x2xf16>, memref<?xf16>, vector<4xindex>, vector<4xi1>
+ xegpu.store %val, %dest[%offset], %mask <{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.
- // CHECK: xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{contiguity = 4 : i64}> : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
- %val = xegpu.load %src[%offset], %mask <{contiguity = 4 : i64}>
+ // A user-provided `chunk_size` round-trips through the optional op attribute.
+ // CHECK: xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{chunk_size = 4 : i64}> : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ %val = xegpu.load %src[%offset], %mask <{chunk_size = 4 : i64}>
: i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.return
}
// CHECK: gpu.func @store_contiguity(%[[arg0:.*]]: vector<16xf32>, %[[arg1:.*]]: i64, %[[arg2:.*]]: vector<16xindex>, %[[arg3:.*]]: vector<16xi1>) {
gpu.func @store_contiguity(%val: vector<16xf32>, %dest: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
- // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{contiguity = 4 : i64}> : vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
- xegpu.store %val, %dest[%offset], %mask <{contiguity = 4 : i64}>
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{chunk_size = 4 : i64}> : vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
+ xegpu.store %val, %dest[%offset], %mask <{chunk_size = 4 : i64}>
: vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index fa00c1d894d8f..402700fd4c5e3 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -615,9 +615,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
}
@@ -646,9 +646,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>}>
+ // CHECK: xegpu.store %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}} <{l1_hint = #xegpu.cache_hint<cached>}>
// 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 79a5d229263c5..b8b8d86cf8b10 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -642,8 +642,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 fe356e6af35c1..3ced1d7e9c6f5 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -260,14 +260,14 @@ gpu.func @prefetch_nd() {
// 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: %[[LOAD:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]]
// 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]>}>
+ <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
: memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
gpu.return
}
@@ -275,20 +275,20 @@ gpu.func @scatter_load_chunksize(%src: memref<256xf16>) {
// 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: %[[LOAD:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]]
// 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: xegpu.store %[[C2]], %arg0[%[[OFFSET]]], %[[MASK]]
// 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]>}>
+ <{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]>}>
+ <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
: vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
index fa9897770a08e..e8a41f6fc99f0 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
@@ -257,7 +257,7 @@ gpu.module @xevm_module{
// 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-NEXT: %[[LD:.*]] = xegpu.load %{{.*}}[%[[OFFSET]]], %[[MASK]]
// 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>
@@ -265,14 +265,14 @@ gpu.module @xevm_module{
// 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-NEXT: xegpu.store %[[IF_CAST]], %{{.*}}[%[[OFFSET]]], %[[MASK]]
// 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}> {
+ %3 = xegpu.load %src[%offset], %1 {
layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
} : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
scf.yield %3 : vector<16x8xf16>
@@ -280,7 +280,7 @@ gpu.module @xevm_module{
%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>
+ xegpu.store %loaded, %src[%offset], %1 {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
@@ -291,9 +291,9 @@ gpu.module @xevm_module{
// 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-NEXT: %[[LOADED:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]]
// CHECK-SAME: memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
-// CHECK-NEXT: xegpu.store %[[LOADED]], %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
+// CHECK-NEXT: xegpu.store %[[LOADED]], %arg0[%[[OFFSET]]], %[[MASK]]
// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
// CHECK-NEXT: }
gpu.module @xevm_module{
@@ -302,10 +302,10 @@ gpu.module @xevm_module{
%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}> {
+ %3 = xegpu.load %src[%offset], %1 {
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>
+ xegpu.store %3, %src[%offset], %1 {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
}
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir b/mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir
index fb32e0012679b..885026f5c85ae 100644
--- a/mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir
+++ b/mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir
@@ -99,7 +99,7 @@ gpu.module @kernel_chunk4 [#xevm.target<chip = "pvc">] {
// CHECK-LABEL: func.func @load_no_target_unchanged(
// CHECK: xegpu.load
// CHECK-NOT: lane_data
-// CHECK-NOT: contiguity
+// CHECK-NOT: chunk_size
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_no_target_unchanged(%ptr: i64) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index b0b16c2adba6e..3d769c86ee2cc 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
}
}
@@ -514,7 +514,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,
@@ -525,7 +525,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>
}
@@ -534,7 +534,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,
@@ -547,7 +547,7 @@ 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
}
@@ -561,7 +561,7 @@ gpu.module @test_kernel {
// 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>
+ // CHECK-COUNT-4: xegpu.load {{.*}}[{{.*}}], {{.*}} <{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,
@@ -572,7 +572,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 = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
+ %ld = xegpu.load %src[%cst], %mask {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>
}
}
@@ -585,7 +585,7 @@ gpu.module @test_kernel {
// 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>
+ // CHECK-COUNT-4: xegpu.store {{.*}}[{{.*}}], {{.*}} <{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,
@@ -598,7 +598,7 @@ gpu.module @test_kernel {
%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>
+ xegpu.store %st_vec, %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -611,8 +611,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> {
@@ -624,7 +624,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>
}
@@ -692,10 +692,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 49fdf1cbee174..75492d8e248d3 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,7 +160,7 @@ 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
}
@@ -173,7 +173,7 @@ gpu.module @test {
// 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>
+ // CHECK-COUNT-4: xegpu.load {{.*}}[{{.*}}], {{.*}} <{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,
@@ -184,7 +184,7 @@ gpu.module @test {
%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>
+ %ld = xegpu.load %src[%cst], %mask {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>
}
@@ -196,7 +196,7 @@ gpu.module @test {
// 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>
+ // CHECK-COUNT-4: xegpu.store {{.*}}[{{.*}}], {{.*}} <{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,
@@ -209,7 +209,7 @@ gpu.module @test {
%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>
+ xegpu.store %st_vec, %src[%cst], %mask {layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 8c5dd6c55b99d..2ca6e8d46ef16 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -317,11 +317,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
}
@@ -332,12 +332,12 @@ 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
@@ -348,11 +348,11 @@ gpu.module @test_distribution {
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: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{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>}
+ %load = xegpu.load %src[%offset], %mask {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
}
@@ -783,9 +783,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>
@@ -803,7 +803,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>
@@ -820,7 +820,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
<{
@@ -878,7 +878,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>
@@ -895,7 +895,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
<{
@@ -911,7 +911,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>
@@ -928,7 +928,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
<{
@@ -1026,7 +1026,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
@@ -1055,7 +1055,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
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index f6b0c50da91dd..a94d172b44801 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -456,7 +456,7 @@ struct TestXeGPUCoalesceGatherScatter
}
StringRef getDescription() const final {
- return "Test driver that turns the contiguity attribute into a lane_data "
+ return "Test driver that turns the chunk_size attribute into a lane_data "
"layout on gather/scatter ops.";
}
@@ -477,7 +477,7 @@ struct TestXeGPUCoalesceGatherScatter
Option<bool> analyzeOnly{
*this, "analyze-only",
- llvm::cl::desc("Only run the analysis (stamp the contiguity attribute); "
+ llvm::cl::desc("Only run the analysis (stamp the chunk_size attribute); "
"do not apply."),
llvm::cl::init(false)};
@@ -509,16 +509,16 @@ struct TestXeGPUCoalesceGatherScatter
return xegpu::LayoutAttr::get(ctx, instData, laneLayout, laneData);
}
- /// Minimal driver: read the `contiguity` attribute the analysis stamped and
+ /// Minimal driver: read the `chunk_size` attribute the analysis stamped and
/// turn it into a `lane_data` layout. This is only a stand-in for the real
/// consumer (layout propagation) so the analysis output can be checked
/// end-to-end; it handles just the simple power-of-two case.
template <typename OpTy>
static void applyContiguity(OpTy op, unsigned maxChunkSize) {
- std::optional<uint64_t> contiguity = op.getContiguity();
+ std::optional<uint64_t> contiguity = op.getChunkSize();
if (!contiguity)
return;
- op.removeContiguityAttr();
+ op.removeChunkSizeAttr();
auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
auto valueTy = op.getValueType();
>From 3049773883c222ed540b6fd7cbb539654d12fecb Mon Sep 17 00:00:00 2001
From: Claude <noreply at anthropic.com>
Date: Mon, 20 Jul 2026 17:37:37 +0000
Subject: [PATCH 2/2] [mlir][XeGPU] Distribute coalesced gather/scatter via
lane_data
Rework SgToLaneLoadGather / SgToLaneStoreScatter to lower a coalesced
gather/scatter to the chunked memory access purely from the layout's
lane_data, without relying on a chunk_size attribute.
When lane_data[FCD] = D > 1 covers the lane's entire per-lane fragment (one
round: lane_layout[FCD] * D == FCD extent, so the distributed per-lane vector
has exactly D elements), the lane owns a single contiguous run
{base, base+1, ..., base+D-1}. Distribution takes the first offset/mask of the
group (scalar base + scalar mask) and emits a value vector<D>; the chunk size
is implied by the value type, which the XeVM lowering already consumes. The
dropped per-lane offset computations become dead code removed during lowering.
The round-robin case (lane_data[FCD] == 1, multiple strided rounds) is excluded
by the lane_data > 1 guard.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 48 +++++++++++++++++++
.../Dialect/XeGPU/sg-to-lane-distribute.mlir | 26 ++++++++++
2 files changed, 74 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 2ac3d1ff2bba5..9a84f172921b5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -580,6 +580,35 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
castValueTo(rewriter, cast<TypedValue<VectorType>>(distMask), maskTy1D);
Value distSource = adaptor.getSource();
+
+ // Coalesced case: the layout assigns a genuine contiguous chunk on the FCD
+ // via `lane_data[FCD] = D > 1`, and that chunk is the lane's *entire*
+ // per-lane fragment (one round: `lane_layout[FCD] * D == FCD extent`, so
+ // the distributed per-lane vector has exactly D elements). Only then does
+ // the lane own a single run of D contiguous elements
+ // `{base, base+1, ..., base+D-1}`, which the XeVM lowering expects as a
+ // chunked access: a scalar base offset + scalar mask + a `vector<D>` value
+ // (the chunk size is implied by the value type). Take the first offset /
+ // mask of the contiguous group; the remaining offset computations become
+ // dead code and are removed during lowering.
+ //
+ // This must NOT fire for the round-robin case (`lane_data[FCD] = 1`,
+ // multiple rounds, e.g. a reduction source where lane l owns
+ // `{l, l+SG, l+2*SG, ...}`): there the per-lane elements are strided, not
+ // contiguous, so a chunked access would read the wrong elements. That case
+ // has `lane_data[FCD] == 1`, so the guard below excludes it.
+ int64_t laneElems = distResultTy1D.getNumElements();
+ int64_t innerLaneData = 1;
+ if (auto laneDataArr = layout.getEffectiveLaneDataAsInt();
+ !laneDataArr.empty())
+ innerLaneData = laneDataArr.back();
+ if (innerLaneData > 1 && laneElems == innerLaneData) {
+ distOffsets = vector::ExtractOp::create(
+ rewriter, op.getLoc(), distOffsets, ArrayRef<int64_t>{0});
+ distMask = vector::ExtractOp::create(rewriter, op.getLoc(), distMask,
+ ArrayRef<int64_t>{0});
+ }
+
auto newOp = xegpu::LoadGatherOp::create(
rewriter, op.getLoc(), distResultTy1D, distSource, distOffsets,
distMask, op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(),
@@ -1113,6 +1142,25 @@ struct SgToLaneStoreScatter
castValueTo(rewriter, cast<TypedValue<VectorType>>(distMask), maskTy1D);
Value distDest = adaptor.getDest();
+
+ // Coalesced case (mirrors SgToLaneLoadGather): a genuine contiguous chunk
+ // on the FCD, `lane_data[FCD] = D > 1` covering the lane's entire per-lane
+ // fragment (one round), stores as a chunked access: scalar base offset +
+ // scalar mask + a `vector<D>` value. Take the first offset / mask of the
+ // contiguous group; the dropped offset computations are DCE'd in lowering.
+ // Excludes the round-robin case (`lane_data[FCD] == 1`).
+ int64_t laneElems = distValueTy1D.getNumElements();
+ int64_t innerLaneData = 1;
+ if (auto laneDataArr = layout.getEffectiveLaneDataAsInt();
+ !laneDataArr.empty())
+ innerLaneData = laneDataArr.back();
+ if (innerLaneData > 1 && laneElems == innerLaneData) {
+ distOffsets = vector::ExtractOp::create(
+ rewriter, op.getLoc(), distOffsets, ArrayRef<int64_t>{0});
+ distMask = vector::ExtractOp::create(rewriter, op.getLoc(), distMask,
+ ArrayRef<int64_t>{0});
+ }
+
xegpu::StoreScatterOp::create(rewriter, op.getLoc(), distValue, distDest,
distOffsets, distMask, op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr(),
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
index e8a41f6fc99f0..19fc7fc87ff18 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
@@ -574,3 +574,29 @@ gpu.module @xevm_module {
gpu.return
}
}
+
+// -----
+// Coalesced gather/scatter: the load/store layout has lane_data[FCD] = 2, i.e.
+// each lane owns 2 *contiguous* elements (one round: lane_layout[FCD] * 2 == 32
+// == FCD extent). Distribution must emit the chunked form the XeVM lowering
+// accepts: a scalar base offset + scalar mask + a value vector<2xf32> (the
+// chunk size is implied by the value type), taking element 0 of the per-lane
+// offsets/mask as the base -- NOT a 2-wide offsets/mask vector. This is driven
+// entirely by lane_data; the dropped offset lanes are DCE'd during lowering.
+gpu.module @xevm_module {
+ // CHECK-LABEL: gpu.func @coalesced_load_store
+ // CHECK: %[[LD:.*]] = xegpu.load %{{.*}}[%[[BASE:.*]]], %{{.*}} : i64, index, i1 -> vector<2xf32>
+ // CHECK: %[[MUL:.*]] = arith.mulf %[[LD]], %{{.*}} : vector<2xf32>
+ // CHECK: xegpu.store %[[MUL]], %{{.*}}[%[[BASE]]], %{{.*}} : vector<2xf32>, i64, index, i1
+ gpu.func @coalesced_load_store(%src: i64, %dst: i64) {
+ %step = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %src[%step], %mask <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %c = arith.constant dense<2.0> : vector<32xf32>
+ %p = arith.mulf %v, %c {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} : vector<32xf32>
+ xegpu.store %p, %dst[%step], %mask <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}>
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ gpu.return
+ }
+}
More information about the Mlir-commits
mailing list