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

[Artem Kroviakov]

================
@@ -1665,6 +1665,115 @@ shuffleDataAsLaneLayoutChange(ConversionPatternRewriter &rewriter, Location loc,
   return res;
 }
 
+/// Redistributes `src` for a `convert_layout` that repacks the `lane_data`
+/// along dimension `repackDim` between round-robin and contiguous, keeping
+/// `lane_layout` and `order` unchanged. Each lane keeps the same `k` elements
+/// and total bits along `repackDim`, but their assignment to lanes changes, so
+/// a `xegpu.bitcast_shuffle` moves the data across lanes: it converts each run
+/// of `k` elements between the round-robin form `vector<k x iElem>` (each lane
+/// owns `k` strided elements) and the contiguous form `vector<1 x iWide>` (each
+/// lane owns one packed `k * elemBits` value), with a `vector.bitcast` to and
+/// from the original element type on either side.
+///
+/// `inputData`/`targetData` are the `repackDim` `lane_data` of the input and
+/// target layouts; exactly one must be 1 (round-robin) and the other `k`
+/// (contiguous). Returns failure if that does not hold or the packed run would
+/// exceed the widest representable integer (64 bits).
+static FailureOr<Value>
+repackLaneData(ConversionPatternRewriter &rewriter, Location loc, Value src,
+               int64_t repackDim, int64_t inputData, int64_t targetData) {
+  auto srcTy = dyn_cast<VectorType>(src.getType());
+  if (!srcTy)
+    return failure();
+  int64_t rank = srcTy.getRank();
+  Type elemTy = srcTy.getElementType();
+  int64_t elemBits = elemTy.getIntOrFloatBitWidth();
+  int64_t k = srcTy.getShape()[repackDim];
+
+  bool roundRobinToContig = inputData == 1 && targetData == k;
+  bool contigToRoundRobin = inputData == k && targetData == 1;
+  if (!roundRobinToContig && !contigToRoundRobin)
+    return failure();
+
+  int64_t wideBits = k * elemBits;
+  if (wideBits > 64)
+    return failure();
+
+  MLIRContext *ctx = rewriter.getContext();
+  VectorType roundRobinIntTy = VectorType::get({k}, IntegerType::get(ctx, elemBits));
+  VectorType contigIntTy = VectorType::get({1}, IntegerType::get(ctx, wideBits));
+  VectorType runTy = VectorType::get({k}, elemTy);
+
+  // Repacks one run (vector<k x elemTy>) into the target form, returning a
+  // vector<k x elemTy>.
+  auto repackRun = [&](Value run) -> Value {
+    if (roundRobinToContig) {
+      Value asInt =
+          vector::BitCastOp::create(rewriter, loc, roundRobinIntTy, run);
+      Value shuffled =
+          xegpu::BitcastShuffleOp::create(rewriter, loc, contigIntTy, asInt);
+      return vector::BitCastOp::create(rewriter, loc, runTy, shuffled);
+    }
+    Value asWide = vector::BitCastOp::create(rewriter, loc, contigIntTy, run);
----------------
akroviakov wrote:

```suggestion
    Value asWide = vector::BitCastOp::create(rewriter, loc, roundRobinToContig ? roundRobinIntTy : contigIntTy, run);
```
What about cutting the if statement above and reduce the op creation code?

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