[Mlir-commits] [mlir] [mlir][xegpu] Add xegpu.lane_shuffle op (PR #210777)
Jianhui Li
llvmlistbot at llvm.org
Thu Jul 30 14:47:02 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/210777
>From 90396243af4950884d9e936f4640720baf8dcd9a Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 20 Jul 2026 18:18:36 +0000
Subject: [PATCH 1/3] [mlir][xegpu] Add xegpu.bitcast_shuffle op
Add a new lane-level XeGPU operation that reinterprets a lane's fragment
as a different element type (possibly with a different number of elements
per lane) while re-distributing the data across the lanes of the subgroup.
The op takes a scalar or 1D vector of numerical type and produces a scalar
or 1D vector of a different type with the same total number of bits per
lane. A verifier enforces the different-type and equal-bit-size
constraints. Round-trip and invalid tests are included.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 66 +++++++++++++++++++
.../mlir/Dialect/XeGPU/IR/XeGPUTypes.td | 5 ++
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 29 ++++++++
mlir/test/Dialect/XeGPU/invalid.mlir | 14 ++++
mlir/test/Dialect/XeGPU/ops.mlir | 28 ++++++++
5 files changed, 142 insertions(+)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 7f8389a6acc47..0a6dc0b0240c5 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1491,6 +1491,72 @@ def XeGPU_TruncfOp
let hasVerifier = 1;
}
+def XeGPU_BitcastShuffleOp
+ : XeGPU_Op<"bitcast_shuffle", [Pure]> {
+ let summary = "Bit-preserving type conversion and re-distribution of a "
+ "subgroup's fragments across its lanes.";
+
+ let description = [{
+ `xegpu.bitcast_shuffle` is a lane-level operation. The `source` operand is
+ the fragment (a scalar or 1D vector) held by one lane, and the `result` is
+ the fragment that lane holds afterwards. It reinterprets the fragments as a
+ different element type, possibly with a different number of elements per
+ lane, and at the same time re-distributes the data across the lanes of the
+ subgroup.
+
+ Both `source` and `result` must be a scalar or a 1D vector of numerical
+ type, and they must have a different type. When they have the same number of
+ elements the operation degenerates into a plain bitcast. Otherwise the total
+ number of bits per lane must be equal on both sides.
+
+ For example, consider `vector<2xi16> -> i32` with a subgroup size of 8. Each
+ lane's `source` is a `vector<2xi16>` of two 16-bit half-words, so the
+ subgroup holds 8 * 2 = 16 half-words in total. Number these half-words 0..15
+ and denote a lane's `source` as `<a,b>`, so lane `i` starts with half-words
+ `<2*i, 2*i+1>`. The `result` `i32` holds two packed half-words, denoted
+ `a|b` (`a` in the low bits, `b` in the high bits):
+
+ ```
+ lane: 0 1 2 3 4 5 6 7
+ source: <0,1> <2,3> <4,5> <6,7> <8,9> <10,11> <12,13> <14,15>
+ result: 0|2 4|6 8|10 12|14 1|3 5|7 9|11 13|15
+ ```
+
+ The half-words are gathered one `source` element position at a time, in
+ lane order: first every lane's element 0 (half-words 0, 2, 4, ..., 14), then
+ every lane's element 1 (half-words 1, 3, 5, ..., 15). This sequence is then
+ handed back out two half-words per lane, in lane order, to form each lane's
+ `result`. In general each lane's two result half-words come from two
+ different lanes' `source` fragments -- that is the cross-lane shuffle.
+
+ The reverse pair `i32 -> vector<2xi16>` undoes this exactly, so bitcasting
+ to one type and back yields the original fragments.
+
+ 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.
+ }];
+
+ let arguments = (ins XeGPU_BitcastShuffleType:$source);
+ let results = (outs XeGPU_BitcastShuffleType:$result);
+
+ let extraClassDeclaration = [{
+ Type getSourceType() {
+ return getSource().getType();
+ }
+
+ Type getResultType() {
+ return getResult().getType();
+ }
+ }];
+
+ let assemblyFormat = [{
+ $source attr-dict `:` type($source) `->` type($result)
+ }];
+
+ let hasVerifier = 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 0423303c23493..dfaea842ba4cc 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
@@ -31,6 +31,11 @@ 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.bitcast_shuffle: a numerical scalar or a 1D
+// vector of numerical type.
+def XeGPU_BitcastShuffleType
+ : AnyTypeOf<[FixedVectorOfRankAndType<[1], [XeGPU_ScalarType]>,
+ 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 2ffe883eb0d9a..24cd67720696f 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -937,6 +937,35 @@ LogicalResult TruncfOp::verify() {
return success();
}
+//===----------------------------------------------------------------------===//
+// XeGPU_BitcastShuffleOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult BitcastShuffleOp::verify() {
+ Type srcType = getSourceType();
+ Type resType = getResultType();
+
+ // A bitcast_shuffle only makes sense between two different types; a
+ // same-type shuffle is a no-op.
+ if (srcType == resType)
+ return emitOpError("source and result must have different types.");
+
+ // Compute the total number of bits on both sides. The operation is
+ // bit-preserving, so the source and result must carry the same amount of
+ // data per lane.
+ auto sizeInBits = [](Type type) -> int64_t {
+ if (auto vecType = dyn_cast<VectorType>(type))
+ return vecType.getNumElements() * vecType.getElementTypeBitWidth();
+ return type.getIntOrFloatBitWidth();
+ };
+
+ if (sizeInBits(srcType) != sizeInBits(resType))
+ return emitOpError(
+ "source and result must have the same total size in bits.");
+
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// XeGPU_DpasMxOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index d5d4950fe7d7e..2ae6cacd0bf50 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -693,6 +693,20 @@ func.func @truncf_invalid_result_size(%a: vector<8x16xf16>) {
return
}
+// -----
+func.func @bitcast_shuffle_same_type(%a: vector<2xi16>) {
+ // expected-error at +1 {{op source and result must have different types}}
+ %1 = xegpu.bitcast_shuffle %a : vector<2xi16> -> vector<2xi16>
+ return
+}
+
+// -----
+func.func @bitcast_shuffle_size_mismatch(%a: i32) {
+ // expected-error at +1 {{op source and result must have the same total size in bits}}
+ %1 = xegpu.bitcast_shuffle %a : i32 -> vector<4xi16>
+ 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 6cffa3eec369b..3dbe570b3b1ad 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -688,6 +688,34 @@ gpu.func @truncf(%a: vector<8x16xf16>) {
gpu.return
}
+// CHECK-LABEL: gpu.func @bitcast_shuffle_scalar_to_vector_16
+gpu.func @bitcast_shuffle_scalar_to_vector_16(%a: i32) {
+ // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : i32 -> vector<2xi16>
+ %1 = xegpu.bitcast_shuffle %a : i32 -> vector<2xi16>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @bitcast_shuffle_vector_16_to_scalar
+gpu.func @bitcast_shuffle_vector_16_to_scalar(%a: vector<2xi16>) {
+ // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : vector<2xi16> -> i32
+ %1 = xegpu.bitcast_shuffle %a : vector<2xi16> -> i32
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @bitcast_shuffle_scalar_to_vector_8
+gpu.func @bitcast_shuffle_scalar_to_vector_8(%a: i32) {
+ // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : i32 -> vector<4xi8>
+ %1 = xegpu.bitcast_shuffle %a : i32 -> vector<4xi8>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @bitcast_shuffle_vector_8_to_scalar
+gpu.func @bitcast_shuffle_vector_8_to_scalar(%a: vector<4xi8>) {
+ // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : vector<4xi8> -> i32
+ %1 = xegpu.bitcast_shuffle %a : vector<4xi8> -> i32
+ 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>
>From b3170c34f65f44b3f3477e95ee258028eb7389ad Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 30 Jul 2026 21:27:34 +0000
Subject: [PATCH 2/3] [mlir][xegpu] Replace bitcast_shuffle with lane_shuffle
Refine the lane-level cross-lane redistribution op introduced in the
preceding commit. Instead of encoding the redistribution in a source ->
result type pair, `xegpu.lane_shuffle` keeps the element type unchanged
and takes a `pack`/`unpack` mode attribute selecting the direction of the
shuffle. This keeps the storage type (e.g. i32) out of the XeGPU IR, so
downstream consumers still see the real element type.
The subgroup collectively holds an S x N grid of elements, where S is the
subgroup size and N the number of elements per lane. `pack` maps element
j of lane i from logical position j*S + i to i*N + j, leaving each lane's
elements at consecutive logical positions; `unpack` is the reverse. The
two modes are exact inverses, which is folded away.
Adding `AllTypesMatch` removes the equal-bit-size check from the verifier,
leaving only a minimum element count. Round-trip, invalid and
canonicalization tests are updated accordingly.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 21 +++++
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 84 ++++++++++---------
.../mlir/Dialect/XeGPU/IR/XeGPUTypes.td | 8 +-
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 38 ++++-----
mlir/test/Dialect/XeGPU/canonicalize.mlir | 36 ++++++++
mlir/test/Dialect/XeGPU/invalid.mlir | 13 +--
mlir/test/Dialect/XeGPU/ops.mlir | 39 +++++----
7 files changed, 147 insertions(+), 92 deletions(-)
create mode 100644 mlir/test/Dialect/XeGPU/canonicalize.mlir
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 ac203f3ef259b..49b98922cee4c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1500,70 +1500,76 @@ def XeGPU_TruncfOp
let hasVerifier = 1;
}
-def XeGPU_BitcastShuffleOp
- : XeGPU_Op<"bitcast_shuffle", [Pure]> {
- let summary = "Bit-preserving type conversion and re-distribution of a "
- "subgroup's fragments across its lanes.";
+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.bitcast_shuffle` is a lane-level operation. The `source` operand is
- the fragment (a scalar or 1D vector) held by one lane, and the `result` is
- the fragment that lane holds afterwards. It reinterprets the fragments as a
- different element type, possibly with a different number of elements per
- lane, and at the same time re-distributes the data across the lanes of the
- subgroup.
-
- Both `source` and `result` must be a scalar or a 1D vector of numerical
- type, and they must have a different type. When they have the same number of
- elements the operation degenerates into a plain bitcast. Otherwise the total
- number of bits per lane must be equal on both sides.
-
- For example, consider `vector<2xi16> -> i32` with a subgroup size of 8. Each
- lane's `source` is a `vector<2xi16>` of two 16-bit half-words, so the
- subgroup holds 8 * 2 = 16 half-words in total. Number these half-words 0..15
- and denote a lane's `source` as `<a,b>`, so lane `i` starts with half-words
- `<2*i, 2*i+1>`. The `result` `i32` holds two packed half-words, denoted
- `a|b` (`a` in the low bits, `b` in the high bits):
+ `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,1> <2,3> <4,5> <6,7> <8,9> <10,11> <12,13> <14,15>
- result: 0|2 4|6 8|10 12|14 1|3 5|7 9|11 13|15
+ 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>
```
- The half-words are gathered one `source` element position at a time, in
- lane order: first every lane's element 0 (half-words 0, 2, 4, ..., 14), then
- every lane's element 1 (half-words 1, 3, 5, ..., 15). This sequence is then
- handed back out two half-words per lane, in lane order, to form each lane's
- `result`. In general each lane's two result half-words come from two
- different lanes' `source` fragments -- that is the cross-lane shuffle.
+ In general each lane's result elements come from `N` different lanes'
+ `source` fragments -- that is the cross-lane shuffle.
- The reverse pair `i32 -> vector<2xi16>` undoes this exactly, so bitcasting
- to one type and back yields the original fragments.
+ 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_BitcastShuffleType:$source);
- let results = (outs XeGPU_BitcastShuffleType:$result);
+ let arguments = (ins XeGPU_LaneShuffleType:$source,
+ XeGPU_LaneShuffleModeAttr:$mode);
+ let results = (outs XeGPU_LaneShuffleType:$result);
let extraClassDeclaration = [{
- Type getSourceType() {
- return getSource().getType();
+ VectorType getSourceType() {
+ return llvm::cast<VectorType>(getSource().getType());
}
- Type getResultType() {
- return getResult().getType();
+ VectorType getResultType() {
+ return llvm::cast<VectorType>(getResult().getType());
}
}];
let assemblyFormat = [{
- $source attr-dict `:` type($source) `->` type($result)
+ $source $mode attr-dict `:` type($source)
}];
let hasVerifier = 1;
+ let hasFolder = 1;
}
def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
index 6f1a7f8292329..ec0e083aa207c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
@@ -31,11 +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.bitcast_shuffle: a numerical scalar or a 1D
-// vector of numerical type.
-def XeGPU_BitcastShuffleType
- : AnyTypeOf<[FixedVectorOfRankAndType<[1], [XeGPU_ScalarType]>,
- XeGPU_ScalarType]>;
+// 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 83431955ae637..1bd5951c9f7f1 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -947,34 +947,28 @@ LogicalResult TruncfOp::verify() {
}
//===----------------------------------------------------------------------===//
-// XeGPU_BitcastShuffleOp
+// XeGPU_LaneShuffleOp
//===----------------------------------------------------------------------===//
-LogicalResult BitcastShuffleOp::verify() {
- Type srcType = getSourceType();
- Type resType = getResultType();
-
- // A bitcast_shuffle only makes sense between two different types; a
- // same-type shuffle is a no-op.
- if (srcType == resType)
- return emitOpError("source and result must have different types.");
-
- // Compute the total number of bits on both sides. The operation is
- // bit-preserving, so the source and result must carry the same amount of
- // data per lane.
- auto sizeInBits = [](Type type) -> int64_t {
- if (auto vecType = dyn_cast<VectorType>(type))
- return vecType.getNumElements() * vecType.getElementTypeBitWidth();
- return type.getIntOrFloatBitWidth();
- };
-
- if (sizeInBits(srcType) != sizeInBits(resType))
- return emitOpError(
- "source and result must have the same total size in bits.");
+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 fb07d1c35b4b9..dc68f5136b4a9 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -694,16 +694,9 @@ func.func @truncf_invalid_result_size(%a: vector<8x16xf16>) {
}
// -----
-func.func @bitcast_shuffle_same_type(%a: vector<2xi16>) {
- // expected-error at +1 {{op source and result must have different types}}
- %1 = xegpu.bitcast_shuffle %a : vector<2xi16> -> vector<2xi16>
- return
-}
-
-// -----
-func.func @bitcast_shuffle_size_mismatch(%a: i32) {
- // expected-error at +1 {{op source and result must have the same total size in bits}}
- %1 = xegpu.bitcast_shuffle %a : i32 -> vector<4xi16>
+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
}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index 98867d86ede16..1e62f9c68f527 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -697,31 +697,38 @@ gpu.func @truncf(%a: vector<8x16xf16>) {
gpu.return
}
-// CHECK-LABEL: gpu.func @bitcast_shuffle_scalar_to_vector_16
-gpu.func @bitcast_shuffle_scalar_to_vector_16(%a: i32) {
- // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : i32 -> vector<2xi16>
- %1 = xegpu.bitcast_shuffle %a : i32 -> vector<2xi16>
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_16
+gpu.func @lane_shuffle_pack_16(%a: vector<2xi16>) {
+ // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<2xi16>
+ %1 = xegpu.lane_shuffle %a pack : vector<2xi16>
gpu.return
}
-// CHECK-LABEL: gpu.func @bitcast_shuffle_vector_16_to_scalar
-gpu.func @bitcast_shuffle_vector_16_to_scalar(%a: vector<2xi16>) {
- // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : vector<2xi16> -> i32
- %1 = xegpu.bitcast_shuffle %a : vector<2xi16> -> i32
+// CHECK-LABEL: gpu.func @lane_shuffle_unpack_16
+gpu.func @lane_shuffle_unpack_16(%a: vector<2xi16>) {
+ // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<2xi16>
+ %1 = xegpu.lane_shuffle %a unpack : vector<2xi16>
gpu.return
}
-// CHECK-LABEL: gpu.func @bitcast_shuffle_scalar_to_vector_8
-gpu.func @bitcast_shuffle_scalar_to_vector_8(%a: i32) {
- // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : i32 -> vector<4xi8>
- %1 = xegpu.bitcast_shuffle %a : i32 -> vector<4xi8>
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_8
+gpu.func @lane_shuffle_pack_8(%a: vector<4xi8>) {
+ // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<4xi8>
+ %1 = xegpu.lane_shuffle %a pack : vector<4xi8>
gpu.return
}
-// CHECK-LABEL: gpu.func @bitcast_shuffle_vector_8_to_scalar
-gpu.func @bitcast_shuffle_vector_8_to_scalar(%a: vector<4xi8>) {
- // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : vector<4xi8> -> i32
- %1 = xegpu.bitcast_shuffle %a : vector<4xi8> -> i32
+// CHECK-LABEL: gpu.func @lane_shuffle_unpack_8
+gpu.func @lane_shuffle_unpack_8(%a: vector<4xi8>) {
+ // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<4xi8>
+ %1 = xegpu.lane_shuffle %a unpack : vector<4xi8>
+ 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
}
>From 00265a21176b12f1f7e107a7e7a07719bb94ced9 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 30 Jul 2026 21:45:13 +0000
Subject: [PATCH 3/3] [mlir][xegpu] Use float element types in lane_shuffle
round-trip tests
Cover the element types the operation is expected to be used with --
f4E2M1FN, f8E5M2 and f16 -- with a pack and an unpack case each, replacing
the integer-typed round-trip tests.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
mlir/test/Dialect/XeGPU/ops.mlir | 47 ++++++++++++++++++--------------
1 file changed, 27 insertions(+), 20 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index 1e62f9c68f527..f733491bfb7ee 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -697,38 +697,45 @@ gpu.func @truncf(%a: vector<8x16xf16>) {
gpu.return
}
-// CHECK-LABEL: gpu.func @lane_shuffle_pack_16
-gpu.func @lane_shuffle_pack_16(%a: vector<2xi16>) {
- // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<2xi16>
- %1 = xegpu.lane_shuffle %a pack : vector<2xi16>
+// 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_16
-gpu.func @lane_shuffle_unpack_16(%a: vector<2xi16>) {
- // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<2xi16>
- %1 = xegpu.lane_shuffle %a unpack : vector<2xi16>
+// 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_8
-gpu.func @lane_shuffle_pack_8(%a: vector<4xi8>) {
- // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<4xi8>
- %1 = xegpu.lane_shuffle %a pack : vector<4xi8>
+// 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_8
-gpu.func @lane_shuffle_unpack_8(%a: vector<4xi8>) {
- // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<4xi8>
- %1 = xegpu.lane_shuffle %a unpack : vector<4xi8>
+// 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_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>
+// 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
}
More information about the Mlir-commits
mailing list