[Mlir-commits] [mlir] 007c684 - [mlir][xegpu] Add xegpu.lane_shuffle op (#210777)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 1 11:40:19 PDT 2026
Author: Jianhui Li
Date: 2026-08-01T11:40:14-07:00
New Revision: 007c68420876af752e14e1c0e8bdab3a6a4058b0
URL: https://github.com/llvm/llvm-project/commit/007c68420876af752e14e1c0e8bdab3a6a4058b0
DIFF: https://github.com/llvm/llvm-project/commit/007c68420876af752e14e1c0e8bdab3a6a4058b0.diff
LOG: [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 at anthropic.com>
Added:
mlir/test/Dialect/XeGPU/canonicalize.mlir
Modified:
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
mlir/test/Dialect/XeGPU/invalid.mlir
mlir/test/Dialect/XeGPU/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index dabbfd4c79de2..2b977fc418553 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/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 = [{
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 65c046bb37ac4..49b98922cee4c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/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`
diff erent 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
diff erent.
+
+ 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";
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
index 78c29ae9e7f43..ec0e083aa207c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
+++ b/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 = [],
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 060da62302867..1bd5951c9f7f1 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/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
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/XeGPU/canonicalize.mlir b/mlir/test/Dialect/XeGPU/canonicalize.mlir
new file mode 100644
index 0000000000000..755d27ec42878
--- /dev/null
+++ b/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>
+}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index ac87045de7012..dc68f5136b4a9 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/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 at +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 at +1 {{Accumulator type must match result type.}}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index 5591eea00a2ea..f733491bfb7ee 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/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>
More information about the Mlir-commits
mailing list