[Mlir-commits] [mlir] [mlir][xegpu] Remove chunk_size attribute; infer from types (PR #205122)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Tue Jul 21 08:51:37 PDT 2026


llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Md Abdullah Shahneous Bari (mshahneo)

<details>
<summary>Changes</summary>

The chunk_size attribute on xegpu.load (LoadGatherOp) and xegpu.store (StoreScatterOp) was redundant with the operand/result types: the XeVM lowering never read it, VectorToXeGPU already built these ops with an empty chunk_size, and the op docs noted it could be inferred from the type.

Remove the attribute and infer the chunk size from the value/result and mask types via a computed getChunkSize() op method. The mask carries one element per lane, so the chunk size is the trailing value dimension when the value has more elements than the mask, and 1 otherwise. This keeps the existing op.getChunkSize() call sites working while respecting library layering (no XeGPUUtils dependency from the IR library).

Update the gather/scatter builders, transforms (propagate-layout, sg-to-lane, wg-to-sg, unroll), VectorToXeGPU lowerings, and lit tests. Two invalid.mlir cases that were only invalid because a stated chunk_size contradicted the type are removed; two others are reclassified to the type/shape diagnostic.

---

Patch is 95.88 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/205122.diff


23 Files Affected:

- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td (+21-16) 
- (modified) mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h (+2-2) 
- (modified) mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h (+12) 
- (modified) mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp (-4) 
- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp (+55-31) 
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp (+4-4) 
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp (+4-2) 
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp (+9-8) 
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp (+5-15) 
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp (+5-11) 
- (modified) mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp (+18-2) 
- (modified) mlir/test/Dialect/XeGPU/contiguity-analysis.mlir (+19-19) 
- (modified) mlir/test/Dialect/XeGPU/invalid.mlir (+12-41) 
- (modified) mlir/test/Dialect/XeGPU/ops.mlir (+17-17) 
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir (+4-4) 
- (modified) mlir/test/Dialect/XeGPU/propagate-layout.mlir (+2-2) 
- (modified) mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir (+6-6) 
- (modified) mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir (+8-8) 
- (modified) mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir (+1-1) 
- (modified) mlir/test/Dialect/XeGPU/xegpu-blocking.mlir (+20-20) 
- (modified) mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir (+8-8) 
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir (+16-16) 
- (modified) mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp (+5-5) 


``````````diff
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,
+       ...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/205122


More information about the Mlir-commits mailing list