[Mlir-commits] [mlir] [MLIR][XeGPU] Enable WG-level mxfp GEMM via generalized shape_cast collapse inference (PR #201496)

[Charitha Saumya]

================
@@ -637,6 +644,188 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
   return collapsedLayout;
 }
 
+// Derive a new layout by expanding a single dimension `dim` into multiple
+// adjacent dimensions whose extents are given by `targetShape`.
+//
+// Distribution policy on the expanded src dims (replacing `dim`):
+//   - sg_layout / lane_layout: spread outer-to-inner; each dim takes
+//     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).
----------------
charithaintc wrote:

what is sg_layout[i] here? target shape is the shape we expand into. so sg_layout[dim] should already have some value. shouldn't we just expand that?

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