[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).
+//   - lane_data: fill innermost-first, capped per dim by
+//     (targetShape[i] / sgLayout[i]) / laneLayout[i] (the per-lane share of
+//     the per-sg extent).
+//   - inst_data: seeded from laneLayout[i] * laneData[i] per dim, then the
+//     remaining factor is distributed innermost-first (capped per dim by
+//     the per-sg extent).
+//   - order: the original dim index is replaced by the expanded dim indices
+//     in innermost-fastest order; entries past `dim` shift up by
+//     `targetShape.size() - 1`.
----------------
charithaintc wrote:

this function is hard to understand. maybe some examples will help for furthur review. 

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