[mlir][xegpu] Add xegpu.lane_shuffle op (#210777)

Add a lane-level XeGPU operation that re-distributes a subgroup's
fragments across its lanes without changing the element type.

The op takes a 1D vector — the fragment held by one lane — and returns a
fragment of the same type. A pack/unpack mode selects the direction:
viewing the subgroup as an S x N element grid (S = subgroup size, N =
elements per lane), pack moves element j of lane i from logical position
j*S + i to i*N + j, so a lane's elements end up at consecutive
positions; unpack is the reverse.
  
This implements xegpu.convert_layout semantics at the lane level when
lane_layout is unchanged but lane_data differs.
  
assisted-by-claude

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

@@ -132,6 +132,27 @@ def XeGPU_FenceScopeAttr:
     let assemblyFormat = "$value";
 }
 
+//===----------------------------------------------------------------------===//
+// XeGPU Lane Shuffle Enums.
+//===----------------------------------------------------------------------===//
+def XeGPU_LaneShufflePack: I32EnumAttrCase<"Pack", 0, "pack">;
+def XeGPU_LaneShuffleUnpack: I32EnumAttrCase<"Unpack", 1, "unpack">;
+def XeGPU_LaneShuffleMode: I32EnumAttr<"LaneShuffleMode",
+      "The direction of the lane shuffle performed by xegpu.lane_shuffle.",
+      [XeGPU_LaneShufflePack, XeGPU_LaneShuffleUnpack]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::xegpu";
+}
+
+def XeGPU_LaneShuffleModeAttr:
+  EnumAttr<XeGPU_Dialect, XeGPU_LaneShuffleMode, "lane_shuffle_mode"> {
+    let summary = [{Describes the direction of a lane shuffle.
+                    "pack" gathers a lane's elements into consecutive logical
+                    positions. "unpack" scatters them back out, strided by the
+                    subgroup size.}];
+    let assemblyFormat = "$value";
+}
+
 def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
   let cppNamespace = "::mlir::xegpu";
   let description = [{

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

@@ -1500,6 +1500,78 @@ def XeGPU_TruncfOp
   let hasVerifier = 1;
 }
 
+def XeGPU_LaneShuffleOp
+    : XeGPU_Op<"lane_shuffle", [Pure, AllTypesMatch<["source", "result"]>]> {
+  let summary = "Re-distributes a subgroup's fragments across its lanes.";
+
+  let description = [{
+    `xegpu.lane_shuffle` is a lane-level operation. The `source` operand is the
+    fragment (a 1D vector) held by one lane, and the `result` is the fragment
+    that lane holds afterwards. It re-distributes the elements across the lanes
+    of the subgroup without changing their type, so `source` and `result` have
+    the same type.
+
+    Let `S` be the subgroup size and `N` the number of elements per lane, so the
+    subgroup collectively holds `S * N` elements. Number these elements by their
+    logical position `0 .. S*N-1`. The `mode` selects which of the two
+    orientations of the `S x N` (lane, element) index grid the operation maps
+    between:
+
+    * `pack`: element `j` of lane `i` holds logical position `j * S + i` before
+      the operation and `i * N + j` after it. A lane's elements end up at
+      consecutive logical positions.
+    * `unpack`: the reverse mapping. A lane's elements end up strided by `S`.
+
+    The two modes are exact inverses, so shuffling one way and then the other
+    yields the original fragments.
+
+    For example, consider `vector<2xi16>` with a subgroup size of 8, denoting a
+    lane's fragment as `<a,b>`. In `pack` mode:
+
+    ```
+    lane:       0      1      2      3      4      5      6      7
+    source:   <0,8>  <1,9>  <2,10> <3,11> <4,12> <5,13> <6,14> <7,15>
+    result:   <0,1>  <2,3>  <4,5>  <6,7>  <8,9> <10,11> <12,13> <14,15>
+    ```
+
+    In general each lane's result elements come from `N` different lanes'
+    `source` fragments -- that is the cross-lane shuffle.
+
+    This operation implements the `xegpu.convert_layout` semantics at the lane
+    level, for the case where the `lane_layout` is the same but the `lane_data`
+    is different.
+
+    The operation assumes a full subgroup (the subgroup size equals the maximum
+    subgroup size) and that all lanes execute the same dynamic instance of the
+    operation; otherwise the behavior is undefined.
+
+    Arguments:
+    - `source`: The 1D vector fragment held by one lane.
+    - `mode`: The direction of the shuffle, either `pack` or `unpack`.
+  }];
+
+  let arguments = (ins XeGPU_LaneShuffleType:$source,
+                       XeGPU_LaneShuffleModeAttr:$mode);
+  let results = (outs XeGPU_LaneShuffleType:$result);
+
+  let extraClassDeclaration = [{
+    VectorType getSourceType() {
+      return llvm::cast<VectorType>(getSource().getType());
+    }
+
+    VectorType getResultType() {
+      return llvm::cast<VectorType>(getResult().getType());
+    }
+  }];
+
+  let assemblyFormat = [{
+    $source $mode attr-dict `:` type($source)
+  }];
+
+  let hasVerifier = 1;
+  let hasFolder = 1;
+}
+
 def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
                                           AnchorLayoutInterface]> {
   let summary = "It performs scaled mma computation";

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td

@@ -31,6 +31,9 @@ def XeGPU_VectorOrScalarType
     : AnyTypeOf<[VectorOfRankAndType<[1,2,3,4,5,6,7,8], [XeGPU_ScalarType, Index]>, XeGPU_ScalarType]>;
 def XeGPU_GatherScatterBaseAddrType
     : AnyTypeOf<[MemRefRankOf<[XeGPU_ScalarType], [1]>, XeGPU_PointerType]>;
+// Operand/result type of xegpu.lane_shuffle: a 1D vector of numerical type.
+def XeGPU_LaneShuffleType
+    : FixedVectorOfRankAndType<[1], [XeGPU_ScalarType]>;
 
 // common base class for types in XeGPU dialect
 class XeGPUTypeDef<string name, string typeMnemonic, list<Trait> traits = [],

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

@@ -946,6 +946,29 @@ LogicalResult TruncfOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// XeGPU_LaneShuffleOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult LaneShuffleOp::verify() {
+  // With a single element per lane there is nothing to re-distribute, so the
+  // operation would be a no-op.
+  if (getSourceType().getNumElements() < 2)
+    return emitOpError("requires a source vector with at least 2 elements.");
+
+  return success();
+}
+
+OpFoldResult LaneShuffleOp::fold(FoldAdaptor adaptor) {
+  // The two modes are exact inverses, so a pack feeding an unpack (or vice
+  // versa) restores the original fragments.
+  auto producer = getSource().getDefiningOp<LaneShuffleOp>();
+  if (producer && producer.getMode() != getMode())
+    return producer.getSource();
+
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // XeGPU_DpasMxOp
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/XeGPU/canonicalize.mlir

@@ -0,0 +1,36 @@
+// RUN: mlir-opt --canonicalize --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: func.func @fold_lane_shuffle_pack_unpack
+// CHECK-SAME: %[[ARG0:.+]]: vector<2xi16>
+// CHECK-NOT: xegpu.lane_shuffle
+// CHECK: return %[[ARG0]]
+func.func @fold_lane_shuffle_pack_unpack(%a: vector<2xi16>) -> vector<2xi16> {
+  %0 = xegpu.lane_shuffle %a pack : vector<2xi16>
+  %1 = xegpu.lane_shuffle %0 unpack : vector<2xi16>
+  return %1 : vector<2xi16>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @fold_lane_shuffle_unpack_pack
+// CHECK-SAME: %[[ARG0:.+]]: vector<4xf8E5M2>
+// CHECK-NOT: xegpu.lane_shuffle
+// CHECK: return %[[ARG0]]
+func.func @fold_lane_shuffle_unpack_pack(%a: vector<4xf8E5M2>) -> vector<4xf8E5M2> {
+  %0 = xegpu.lane_shuffle %a unpack : vector<4xf8E5M2>
+  %1 = xegpu.lane_shuffle %0 pack : vector<4xf8E5M2>
+  return %1 : vector<4xf8E5M2>
+}
+
+// -----
+
+// Two shuffles in the same direction are not inverses and must not fold.
+
+// CHECK-LABEL: func.func @no_fold_lane_shuffle_pack_pack
+// CHECK: xegpu.lane_shuffle {{.*}} pack
+// CHECK: xegpu.lane_shuffle {{.*}} pack
+func.func @no_fold_lane_shuffle_pack_pack(%a: vector<2xi16>) -> vector<2xi16> {
+  %0 = xegpu.lane_shuffle %a pack : vector<2xi16>
+  %1 = xegpu.lane_shuffle %0 pack : vector<2xi16>
+  return %1 : vector<2xi16>
+}

mlir/test/Dialect/XeGPU/invalid.mlir

@@ -693,6 +693,13 @@ func.func @truncf_invalid_result_size(%a: vector<8x16xf16>) {
   return
 }
 
+// -----
+func.func @lane_shuffle_single_element(%a: vector<1xi32>) {
+  // expected-error@+1 {{op requires a source vector with at least 2 elements}}
+  %1 = xegpu.lane_shuffle %a pack : vector<1xi32>
+  return
+}
+
 // -----
 func.func @dpas_mx_acc_result_type_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xbf16>) {
   // expected-error@+1 {{Accumulator type must match result type.}}

mlir/test/Dialect/XeGPU/ops.mlir

@@ -697,6 +697,48 @@ gpu.func @truncf(%a: vector<8x16xf16>) {
   gpu.return
 }
 
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_f4
+gpu.func @lane_shuffle_pack_f4(%a: vector<4xf4E2M1FN>) {
+  // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<4xf4E2M1FN>
+  %1 = xegpu.lane_shuffle %a pack : vector<4xf4E2M1FN>
+  gpu.return
+}
+
+// CHECK-LABEL: gpu.func @lane_shuffle_unpack_f4
+gpu.func @lane_shuffle_unpack_f4(%a: vector<4xf4E2M1FN>) {
+  // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<4xf4E2M1FN>
+  %1 = xegpu.lane_shuffle %a unpack : vector<4xf4E2M1FN>
+  gpu.return
+}
+
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_f8
+gpu.func @lane_shuffle_pack_f8(%a: vector<4xf8E5M2>) {
+  // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<4xf8E5M2>
+  %1 = xegpu.lane_shuffle %a pack : vector<4xf8E5M2>
+  gpu.return
+}
+
+// CHECK-LABEL: gpu.func @lane_shuffle_unpack_f8
+gpu.func @lane_shuffle_unpack_f8(%a: vector<4xf8E5M2>) {
+  // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<4xf8E5M2>
+  %1 = xegpu.lane_shuffle %a unpack : vector<4xf8E5M2>
+  gpu.return
+}
+
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_f16
+gpu.func @lane_shuffle_pack_f16(%a: vector<2xf16>) {
+  // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<2xf16>
+  %1 = xegpu.lane_shuffle %a pack : vector<2xf16>
+  gpu.return
+}
+
+// CHECK-LABEL: gpu.func @lane_shuffle_unpack_f16
+gpu.func @lane_shuffle_unpack_f16(%a: vector<2xf16>) {
+  // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<2xf16>
+  %1 = xegpu.lane_shuffle %a unpack : vector<2xf16>
+  gpu.return
+}
+
 // CHECK-LABEL: gpu.func @dpas_mx
 gpu.func @dpas_mx(%a : vector<8x32xf8E5M2>, %b: vector<32x16xf8E5M2>, %acc: vector<8x16xbf16>, %a_scale: vector<8x1xf8E8M0FNU>, %b_scale: vector<1x16xf8E8M0FNU>) {
   // CHECK: %{{.+}} = xegpu.dpas_mx %{{.+}}, %{{.+}}, %{{.+}} scale_a = %{{.+}} scale_b = %{{.+}} : (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>