[Mlir-commits] [mlir] [MLIR][XeGPU] Support partial subgroup lane distribution (PR #201667)

[Jianhui Li]

================
@@ -1531,7 +1538,99 @@ struct SgToLaneConvertLayout
       return rewriter.notifyMatchFailure(
           op, "lowering incompatible convert_layout not yet supported");
     }
+    // Handle special case where laneLayout involves partial subgroup
+    if (inputLayout.getEffectiveOrderAsInt() ==
+            targetLayout.getEffectiveOrderAsInt() &&
+        inputLayout.getRank() == 2 && targetLayout.getRank() == 2) {
+      auto laneLayout = inputLayout.getEffectiveLaneLayoutAsInt();
+      auto targetLaneLayout = targetLayout.getEffectiveLaneLayoutAsInt();
+      auto dataLayout = inputLayout.getEffectiveLaneDataAsInt();
+      auto targetDataLayout = targetLayout.getEffectiveLaneDataAsInt();
+      if (dataLayout == targetDataLayout && laneLayout != targetLaneLayout) {
+        int64_t factor = 1;
+        int64_t distDim = 1;
+        if (laneLayout[0] == 1) {
+          // distribute along inner dim
+          factor = laneLayout[1] / targetLaneLayout[1];
+        } else if (laneLayout[1] == 1) {
+          // distribute alone outer dim
+          factor = laneLayout[0] / targetLaneLayout[0];
+          distDim = 0;
+        } else {
+          return rewriter.notifyMatchFailure(
+              op, "distribution across multiple dimensions is not supported");
+        }
+        // Currently, partial subgroup support is limited to special cases
----------------
Jianhui-Li wrote:

Please consider extract the specific case handling as a proper function name. 
At high level, I see two cases. The func name could be something like:  
shuffleDataAsLaneLayoutChange(currentLaneNum, targetLaneNum, scalar_ele)
shuffleDataAsLaneDataChange(currentDataNum, targetDataNum, scalar_ele)


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