[Mlir-commits] [mlir] [mlir][xegpu] Lower lane_data repack convert_layout to bitcast_shuffle (PR #210837)

[Artem Kroviakov]

================
@@ -1721,6 +1830,47 @@ struct SgToLaneConvertLayout
       }
     }
 
+    // Handle a pure `lane_data` repack: `lane_layout` and `order` are unchanged
+    // and exactly one dimension's `lane_data` differs, switching between
+    // round-robin (lane_data 1) and contiguous (lane_data == run length). The
+    // number of elements per lane is unchanged, but their assignment to lanes
+    // is not, so the data is moved across lanes with `xegpu.bitcast_shuffle`.
+    //
+    // The changed dimension must be one of the two innermost ones: sg-to-lane
+    // distribution is 2D, with any higher (leading) dimensions being unit, so a
+    // repack of a 3rd-innermost-or-higher dimension is not supported.
+    if (inputLayout.getEffectiveOrderAsInt() ==
+            targetLayout.getEffectiveOrderAsInt() &&
+        inputLayout.getEffectiveLaneLayoutAsInt() ==
+            targetLayout.getEffectiveLaneLayoutAsInt()) {
+      // lane_layout and lane_data always share the same rank (enforced by the
+      // LayoutAttr verifier), so the two lane_data arrays have equal size here
+      // since their lane_layouts are equal.
+      auto laneData = inputLayout.getEffectiveLaneDataAsInt();
+      auto targetLaneData = targetLayout.getEffectiveLaneDataAsInt();
+      // Find the single dimension whose lane_data changed.
----------------
akroviakov wrote:

The comment says *single*, but the code finds *all* dims.
If it is a single dim, we do not need a vector.

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