[Mlir-commits] [mlir] [mlir][xegpu] Fill split-group inner dims on shape_cast result layout (PR #227560)

[Jianhui Li]

================
@@ -2765,27 +2794,59 @@ xegpu::DistributeLayoutAttr xegpu::setupShapeCastResultLayout(
                                      splitDimGroups))
     return consumerLayout;
 
-  SmallVector<int64_t> laneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+  SmallVector<int64_t> laneLayout =
+      consumerLayout.getEffectiveLaneLayoutAsInt();
   SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
   if (laneLayout.empty())
     return consumerLayout;
 
-  xegpu::DistributeLayoutAttr resLayout = consumerLayout;
-  for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
-    size_t i = 0;
-    while (i < dimGroup.size() && laneLayout[dimGroup[i]] == 1)
-      ++i;
-
-    for (++i; i < dimGroup.size(); ++i) {
-      int64_t dim = dimGroup[i];
-      int64_t dimSize = resShape[dim];
-      int64_t lanes = laneLayout[dim];
-      if (lanes <= 0 || dimSize % lanes != 0)
-        return consumerLayout;
-      resLayout = resLayout.setDimData(dim, /*sgData=*/-1,
-                                       instData.empty() ? -1 : dimSize,
-                                       dimSize / lanes);
+  xegpu::DistributeLayoutAttr resLayout = [&]() -> DistributeLayoutAttr {
+    DistributeLayoutAttr layout = consumerLayout;
+    for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
+      size_t i = 0;
+      while (i < dimGroup.size() && laneLayout[dimGroup[i]] == 1)
+        ++i;
+
+      for (++i; i < dimGroup.size(); ++i) {
+        int64_t dim = dimGroup[i];
+        int64_t dimSize = resShape[dim];
+        int64_t lanes = laneLayout[dim];
+        if (lanes <= 0 || dimSize % lanes != 0)
+          return consumerLayout;
+        layout =
+            layout.setDimData(dim, /*sgData=*/-1,
+                              instData.empty() ? -1 : dimSize, dimSize / lanes);
+      }
     }
+    return layout;
+  }();
+
+  // Stretching does not make every layout collapsible; reject the rest.
+  SmallVector<int64_t> resInstData = resLayout.getEffectiveInstDataAsInt();
+  SmallVector<int64_t> resLaneLayout = resLayout.getEffectiveLaneLayoutAsInt();
+  SmallVector<int64_t> resLaneData = resLayout.getEffectiveLaneDataAsInt();
+  if (resLaneLayout.empty() || resLaneData.empty())
+    return resLayout;
+
+  for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
+    // One lane's elements must be a single contiguous run of the collapsed dim.
+    if (!isContiguousTiling(dimGroup, resShape, resLaneData, /*tileShape=*/{},
----------------
Jianhui-Li wrote:

Yes. The order must be adjacent, and inferShapeCastSourceLayout does check that.  It worths a new PR to refactor the code to move the order requirement here upfront. 

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