[Mlir-commits] [mlir] 9f3ac76 - [MLIR][XeGPU] Fix issue with expandDim (#203299)

[llvmlistbot at llvm.org]

Author: Sang Ik Lee
Date: 2026-06-12T08:52:42-07:00
New Revision: 9f3ac763bd8c4be2db1a97cec13bb59c87a9595f

URL: https://github.com/llvm/llvm-project/commit/9f3ac763bd8c4be2db1a97cec13bb59c87a9595f
DIFF: https://github.com/llvm/llvm-project/commit/9f3ac763bd8c4be2db1a97cec13bb59c87a9595f.diff

LOG: [MLIR][XeGPU] Fix issue with expandDim (#203299)

In case, sgData is replicated or
bool sgDataReplicated = hasSgData && origSgDataDim ==
computeProduct(targetShape);
sgDataReplicate is evaluated to "true"
dimSizeCap and perSgShape shouldn't be divided by expSgLayout

Add a regression test (shape_cast_collapse_replicated) covering the
replicated sg_data collapse.

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
    mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index b780c66594eb0..075442d443658 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -661,7 +661,8 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
 //     min(remaining, targetShape[i]); leftover spills into the next inner
 //     dim.
 //   - sg_data: fill innermost-first, capped per dim by
-//     targetShape[i] / sgLayout[i] (the per-sg share of the extent).
+//     targetShape[i] / sgLayout[i] (the per-sg share of the extent),
+//     or targetShape[i] (if sg_data is replicated across all subgroups).
 //   - lane_data: fill innermost-first, capped per dim by
 //     (targetShape[i] / sgLayout[i]) / laneLayout[i] (the per-lane share of
 //     the per-sg extent).
@@ -751,9 +752,11 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
     expSgLayout = spread(origSgLayoutDim, targetShape, /*outerToInner=*/true);
     splice(sgLayout, expSgLayout);
   }
+  bool sgDataReplicated =
+      hasSgData && origSgDataDim == computeProduct(targetShape);
   if (hasSgData) {
     SmallVector<int64_t> dimSizeCap(targetShape.begin(), targetShape.end());
-    if (hasSgLayout)
+    if (hasSgLayout && !sgDataReplicated)
       for (int64_t i = 0; i < expCount; ++i)
         dimSizeCap[i] /= expSgLayout[i];
     SmallVector<int64_t> expSgData =
@@ -765,7 +768,7 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
   // targetShape[i] / sg_layout[i] when sg_layout is present, else
   // targetShape itself.
   SmallVector<int64_t> perSgShape(targetShape.begin(), targetShape.end());
-  if (hasSgLayout)
+  if (hasSgLayout && !sgDataReplicated)
     for (int64_t i = 0; i < expCount; ++i)
       perSgShape[i] /= expSgLayout[i];
 

diff  --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index 5021c8a746045..d44497d0bba34 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -575,3 +575,17 @@ gpu.module @test {
     gpu.return
   }
 }
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: gpu.func @shape_cast_collapse_replicated(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [2, 2, 1], sg_data = [8, 8, 8]>} dense<0.000000e+00> : vector<16x8x8xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<sg_layout = [2, 2], sg_data = [8, 64]>} : vector<16x8x8xf16> to vector<16x64xf16>
+  gpu.func @shape_cast_collapse_replicated(%dst: memref<16x64xf16>) kernel {
+    %cst = arith.constant dense<0.000000e+00> : vector<16x8x8xf16>
+    %0 = vector.shape_cast %cst : vector<16x8x8xf16> to vector<16x64xf16>
+    %tdesc = xegpu.create_nd_tdesc %dst : memref<16x64xf16> -> !xegpu.tensor_desc<16x64xf16>
+    xegpu.store_nd %0, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [8, 64]>}> : vector<16x64xf16>, !xegpu.tensor_desc<16x64xf16>
+    gpu.return
+  }
+}