[Mlir-commits] [mlir] [MLIR][XeGPU] Fix expandDim sg_data cap for replicated layouts (PR #203299)

[Jianhui Li]

================
@@ -752,10 +758,13 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
     splice(sgLayout, expSgLayout);
   }
   if (hasSgData) {
+    // Cap by the full extent targetShape[i], not the per-sg share
+    // targetShape[i] / sgLayout[i]: a replicated dim (sg_layout[dim] *
+    // sg_data[dim] > extent) has sg_data up to the full extent, which would not
+    // fit within the per-sg share. For evenly-distributed layouts the
+    // inner-first fill lands on the per-sg share regardless, so this is
+    // equivalent.
     SmallVector<int64_t> dimSizeCap(targetShape.begin(), targetShape.end());
-    if (hasSgLayout)
----------------
Jianhui-Li wrote:

We need to keep this logic. To address the crash issue, simply add one more check 
  bool sgDataReplicated =
      hasSgData && origSgDataDim == computeProduct(targetShape);
and then conduct the dimSizeCap division loop under 
 if (hasSgLayout && !sgDataReplicated)

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