[Mlir-commits] [mlir] [MLIR][XeGPU] Return element count from uarch for scatter ops (PR #206704)

[Artem Kroviakov]

https://github.com/akroviakov created https://github.com/llvm/llvm-project/pull/206704

This PR fixes uArch to return element count for scatter ops, rather than the byte count, to better comply with the use case.

>From 21fec2835dc27fe9b15f0fd960c1ce355415dffb Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Tue, 30 Jun 2026 10:24:55 +0000
Subject: [PATCH] [MLIR][XeGPU] Return element count from uarch for scatter ops

---
 .../mlir/Dialect/XeGPU/uArch/uArchBase.h      | 10 +++++--
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 28 +++++++++++--------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index e1879bb3ffe1c..00637526b1fbe 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -208,17 +208,21 @@ struct ScatterIoInstructionInterface : public Instruction {
     return B->getInstructionKind() == Kind;
   }
 
-  virtual int32_t getMaxLaneAccessSizeBytes() const = 0;
+  virtual int32_t getMaxLaneAccessSizeElements(int32_t elemBitwidth) const = 0;
   virtual ~ScatterIoInstructionInterface() = default;
 };
 struct LoadGatherInstruction
     : public ScatterIoInstructionInterface<InstructionKind::LoadGather> {
-  int32_t getMaxLaneAccessSizeBytes() const override { return 16; }
+  int32_t getMaxLaneAccessSizeElements(int32_t elemBitwidth) const override {
+    return 16 / (elemBitwidth / 8);
+  }
 };
 
 struct StoreScatterInstruction
     : public ScatterIoInstructionInterface<InstructionKind::StoreScatter> {
-  int32_t getMaxLaneAccessSizeBytes() const override { return 16; }
+  int32_t getMaxLaneAccessSizeElements(int32_t elemBitwidth) const override {
+    return 16 / (elemBitwidth / 8);
+  }
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 61cd253508357..2a219fdf6348a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1713,8 +1713,9 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
 
   const auto *uArchInstruction = dyn_cast<xegpu::uArch::LoadGatherInstruction>(
       uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
-  int maxChunkSize =
-      std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
+  int maxChunkSize = std::min(uArchInstruction->getMaxLaneAccessSizeElements(
+                                  resVecTy.getElementTypeBitWidth()),
+                              contigChunkSize);
 
   return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
                                       maxChunkSize, resShape, subgroupSize);
@@ -1734,8 +1735,9 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
 
   const auto *uArchInstruction = dyn_cast<xegpu::uArch::LoadGatherInstruction>(
       uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
-  int maxChunkSize =
-      std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
+  int maxChunkSize = std::min(uArchInstruction->getMaxLaneAccessSizeElements(
+                                  resVecTy.getElementTypeBitWidth()),
+                              contigChunkSize);
   return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
                                       maxChunkSize, resShape, subgroupSize);
 }
@@ -1786,8 +1788,9 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
   const auto *uArchInstruction =
       dyn_cast<xegpu::uArch::StoreScatterInstruction>(
           uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
-  int maxChunkSize =
-      std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
+  int maxChunkSize = std::min(uArchInstruction->getMaxLaneAccessSizeElements(
+                                  srcVecTy.getElementTypeBitWidth()),
+                              contigChunkSize);
   return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
                                        srcShape, subgroupSize);
 }
@@ -1805,8 +1808,9 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
   const auto *uArchInstruction =
       dyn_cast<xegpu::uArch::StoreScatterInstruction>(
           uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
-  int maxChunkSize =
-      std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
+  int maxChunkSize = std::min(uArchInstruction->getMaxLaneAccessSizeElements(
+                                  srcVecTy.getElementTypeBitWidth()),
+                              contigChunkSize);
 
   return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
                                        srcShape, subgroupSize);
@@ -1825,7 +1829,7 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
 ///   - Otherwise a standard scatter-style factorization is computed via
 ///     `computeScatterIOLaneLayoutAndData`, bounded by `maxChunkSize` — the
 ///     per-lane load width reported by the uArch's LoadGather instruction
-///     (`getMaxLaneAccessSizeBytes`).
+///     (`getMaxLaneAccessSizeElements`).
 ///
 std::optional<xegpu::DistributeLayoutAttr>
 xegpu::completeScatterLoadLaneLayoutFromInstData(
@@ -1845,7 +1849,8 @@ xegpu::completeScatterLoadLaneLayoutFromInstData(
 
   // Reuse the load-side setup with inst_data as the destination shape.
   auto *context = specifiedLayout.getContext();
-  int maxChunkSize = uArchInstruction->getMaxLaneAccessSizeBytes();
+  int maxChunkSize = uArchInstruction->getMaxLaneAccessSizeElements(
+      getElementTypeOrSelf(elemTy).getIntOrFloatBitWidth());
   if (consumerLayout) {
     auto consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
     auto consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
@@ -1883,7 +1888,8 @@ xegpu::completeScatterStoreLaneLayoutFromInstData(
 
   // Reuse the store-side setup with inst_data as the source shape.
   auto *context = specifiedLayout.getContext();
-  int maxChunkSize = uArchInstruction->getMaxLaneAccessSizeBytes();
+  int maxChunkSize = uArchInstruction->getMaxLaneAccessSizeElements(
+      getElementTypeOrSelf(elemTy).getIntOrFloatBitWidth());
   auto [defLaneLayout, defLaneData] = computeScatterIOLaneLayoutAndData(
       specifiedInstData, subgroupSize, maxChunkSize);
   if (!isValidLaneLayout(specifiedInstData, defLaneLayout, defLaneData))