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

[Sang Ik Lee]

================
@@ -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)
----------------
silee2 wrote:

Done

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