[Mlir-commits] [mlir] [MLIR][XeVM] Lower xegpu.lane_shuffle to xevm.bitcast_shuffle (PR #215306)
Sang Ik Lee
llvmlistbot at llvm.org
Mon Aug 10 11:04:15 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/215306
>From e846b9fdb3e76c3f86a37be897b6f2c3d86ef646 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Thu, 6 Aug 2026 21:41:18 +0000
Subject: [PATCH 1/2] [MLIR][XeVM] Add xevm.bitcast_shuffle op and lowering to
LLVM
Add a new `xevm.bitcast_shuffle` op with the same semantics as the SPIR-V
`OpSubgroupBitcastShuffleINTEL` instruction from the
`SPV_INTEL_subgroup_bitcast_shuffle` extension.
The op performs a bit-preserving type conversion and shuffle of the source
data among the invocations in a subgroup. All invocations cooperate: the
components of the operand are concatenated across the subgroup in order of
the subgroup local invocation ID, and the resulting bit stream is
redistributed back to the invocations in chunks the size of a result
component. When the operand and the result have the same number of
components the behavior is that of a traditional bitcast. The operation is
reversible.
`src` and `res` must be a scalar or a 1D vector of numerical type, must have
different types and must have the same total bit width. The shuffle is
performed at byte granularity, so element types narrower than 8 bits are not
supported. Floating point formats without an LLVM representation, such as
the f8 formats, are expected to already have been converted to a same-width
integer type by the LLVM type converter.
The op is not lowered to the SPIR-V instruction. Instead,
`convert-xevm-to-llvm` emits a call to the IGC intrinsic
`llvm.genx.GenISA.SubgroupBitcastShuffle`, which is overloaded on both the
result and the operand type. The intrinsic has no bf16 overload, so bf16
operands and results are bitcast to i16 around the call.
---
mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td | 66 ++++++
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 129 ++++++++---
mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp | 16 ++
.../Conversion/XeVMToLLVM/xevm-to-llvm.mlir | 68 ++++++
mlir/test/Dialect/LLVMIR/invalid.mlir | 25 ++
mlir/test/Dialect/LLVMIR/xevm.mlir | 32 +++
.../XeVM/GPU/xevm_bitcast_shuffle.mlir | 219 ++++++++++++++++++
7 files changed, 529 insertions(+), 26 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/XeVM/GPU/xevm_bitcast_shuffle.mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
index 8c4e409e9c395..1ab08c04bd903 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
@@ -740,6 +740,72 @@ def XeVM_MMAMxOp
let hasVerifier = 1;
}
+def XeVM_BitcastShuffleElemType
+ : AnyTypeOf<[I8, I16, I32, I64, BF16, F16, F32, F64]>;
+
+def XeVM_BitcastShuffleOp
+ : XeVM_Op<"bitcast_shuffle">,
+ Results<(outs AnyTypeOf<
+ [XeVM_BitcastShuffleElemType,
+ FixedVectorOfRankAndType<[1],
+ [XeVM_BitcastShuffleElemType]>]>:$res)>,
+ Arguments<(ins AnyTypeOf<
+ [XeVM_BitcastShuffleElemType,
+ FixedVectorOfRankAndType<[1],
+ [XeVM_BitcastShuffleElemType]>]>:$src)> {
+
+ let summary = "Subgroup bit-preserving conversion and shuffle";
+
+ let description = [{
+ The `xevm.bitcast_shuffle` operation performs a bit-preserving type
+ conversion and shuffle of `src` among the invocations in a subgroup. It is
+ a cooperative operation: all invocations in the subgroup participate. This
+ operation may execute more efficiently than a traditional bitcast.
+
+ Both `src` and `res` must be a scalar or a 1D vector of numerical type, and
+ they must have different types. The total number of bits of `res` must equal
+ the total number of bits of `src`. The shuffle is performed at byte
+ granularity, so element types narrower than 8 bits, such as f4 and i4, are
+ not supported. Floating point formats that have no LLVM representation, such
+ as the f8 formats, are expected to already have been converted to a
+ same-width integer type.
+
+ If `res` has the same number of components as `src`, the behavior is the
+ same as a traditional bitcast.
+
+ Otherwise, the bitcast and shuffle is performed as follows. Take the first
+ component of `src` for each invocation in the subgroup and concatenate those
+ components together, ordered by subgroup local invocation ID. Then repeat
+ for the next component of `src` of each invocation, until all of `src` has
+ been concatenated. This forms an `M * N` bit number, where `M` is the size
+ in bits of `src` and `N` is the subgroup size.
+
+ Now take the first `C` bits of the concatenated number and assign them to
+ the first component of `res` of the first invocation in the subgroup, where
+ `C` is the size in bits of each component of `res` when it is a vector type,
+ or the size in bits of `res` when it is a scalar type. Assign the next `C`
+ bits to the first component of `res` of the next invocation, and so on,
+ ordered by subgroup local invocation ID. Once `C` bits have been assigned to
+ the first component of all invocations in the subgroup, repeat for the
+ second component of `res`, and so on, until all bits have been assigned.
+
+ Like a traditional bitcast, the operation is reversible: an
+ `xevm.bitcast_shuffle` to one type followed by another `xevm.bitcast_shuffle`
+ back to the original type yields the original source data.
+
+ Example:
+ ```mlir
+ %res = xevm.bitcast_shuffle %src : (vector<8xi16>) -> vector<4xi32>
+ ```
+ }];
+
+ let assemblyFormat = [{
+ $src attr-dict `:` functional-type(operands, results)
+ }];
+
+ let hasVerifier = 1;
+}
+
//===----------------------------------------------------------------------===//
// XeVM target attribute.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index c51da4d5d4d3d..f499733023dcd 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -99,6 +99,24 @@ std::string mangle(StringRef baseName, ArrayRef<Type> types,
return os.str();
}
+// Returns the LLVM intrinsic overload suffix of `ty`, i.e. the spelling used
+// when naming an overloaded intrinsic: `i32`, `f16`, `bf16`, `v8i16`, ...
+std::string getIntrinsicTypeSuffix(Type ty) {
+ return TypeSwitch<Type, std::string>(ty)
+ .Case([](VectorType ty) -> std::string {
+ return "v" + std::to_string(ty.getNumElements()) +
+ getIntrinsicTypeSuffix(ty.getElementType());
+ })
+ .Case([](IntegerType ty) -> std::string {
+ return "i" + std::to_string(ty.getWidth());
+ })
+ .Case([](BFloat16Type) -> std::string { return "bf16"; })
+ .Case([](FloatType ty) -> std::string {
+ return "f" + std::to_string(ty.getWidth());
+ })
+ .DefaultUnreachable("unhandled type for intrinsic suffix");
+}
+
std::string builtinElemType(ElemType elemType) {
switch (elemType) {
case ElemType::BF8:
@@ -1520,6 +1538,65 @@ class MMAMxToOCLPattern : public OpConversionPattern<MMAMxOp> {
}
};
+// Lowers `xevm.bitcast_shuffle` to a call to the IGC intrinsic
+// `llvm.genx.GenISA.SubgroupBitcastShuffle`, which is overloaded on both the
+// result and the operand type. E.g. a `vector<4xi8>` -> `vector<2xi16>` shuffle
+// becomes a call to
+// `llvm.genx.GenISA.SubgroupBitcastShuffle.v2i16.v4i8`.
+class BitcastShuffleToGenISAPattern
+ : public OpConversionPattern<BitcastShuffleOp> {
+ // The intrinsic is only available for integer and non-bfloat float overloads,
+ // so bfloat16 is passed through as int16. The shuffle is bit-preserving, so
+ // this is a no-op at the ISA level.
+ static Type encodeBF16AsI16(Type ty, Builder &builder) {
+ if (auto vecTy = dyn_cast<VectorType>(ty)) {
+ if (vecTy.getElementType().isBF16())
+ return VectorType::get(vecTy.getShape(), builder.getIntegerType(16));
+ return vecTy;
+ }
+ return ty.isBF16() ? builder.getIntegerType(16) : ty;
+ }
+
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(BitcastShuffleOp op, BitcastShuffleOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ Location loc = op.getLoc();
+ Type srcOrigTy = op.getSrc().getType();
+ Type resOrigTy = op.getRes().getType();
+ Type srcTy = encodeBF16AsI16(srcOrigTy, rewriter);
+ Type resTy = encodeBF16AsI16(resOrigTy, rewriter);
+
+ Value src = adaptor.getSrc();
+ if (srcTy != srcOrigTy)
+ src = LLVM::BitcastOp::create(rewriter, loc, srcTy, src);
+
+ Value result;
+ if (srcTy == resTy) {
+ // Once bf16 is encoded as i16 both sides may end up with the same type,
+ // e.g. `vector<2xbf16>` -> `vector<2xi16>`. The component count is
+ // unchanged, so this is a plain bitcast and the intrinsic (which requires
+ // distinct operand and result types) must not be called.
+ result = src;
+ } else {
+ std::string fnName = "llvm.genx.GenISA.SubgroupBitcastShuffle." +
+ getIntrinsicTypeSuffix(resTy) + "." +
+ getIntrinsicTypeSuffix(srcTy);
+
+ result = createDeviceFunctionCall(rewriter, fnName, resTy, {srcTy}, {src},
+ {}, convergentNoUnwindWillReturnAttrs,
+ op.getOperation())
+ ->getResult(0);
+ }
+
+ if (resTy != resOrigTy)
+ result = LLVM::BitcastOp::create(rewriter, loc, resOrigTy, result);
+
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+};
+
class AllocaToGlobalPattern : public OpConversionPattern<LLVM::AllocaOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
@@ -1787,30 +1864,30 @@ void ::mlir::populateXeVMToLLVMConversionPatterns(ConversionTarget &target,
return !op->hasAttr("cache_control");
});
target.addIllegalDialect<XeVMDialect>();
- patterns
- .add<LoadStorePrefetchToOCLPattern<BlockLoad2dOp>,
- LoadStorePrefetchToOCLPattern<BlockStore2dOp>,
- LoadStorePrefetchToOCLPattern<BlockPrefetch2dOp>, MMAToOCLPattern,
- MemfenceToOCLPattern, PrefetchToOCLPattern,
- LLVMLoadStoreToOCLPattern<LLVM::LoadOp>,
- LLVMLoadStoreToOCLPattern<LLVM::StoreOp>,
- BlockLoadStore1DToOCLPattern<BlockLoadOp>,
- BlockLoadStore1DToOCLPattern<BlockStoreOp>,
- LaunchConfigOpToOCLPattern<WorkitemIdXOp>,
- LaunchConfigOpToOCLPattern<WorkitemIdYOp>,
- LaunchConfigOpToOCLPattern<WorkitemIdZOp>,
- LaunchConfigOpToOCLPattern<WorkgroupDimXOp>,
- LaunchConfigOpToOCLPattern<WorkgroupDimYOp>,
- LaunchConfigOpToOCLPattern<WorkgroupDimZOp>,
- LaunchConfigOpToOCLPattern<WorkgroupIdXOp>,
- LaunchConfigOpToOCLPattern<WorkgroupIdYOp>,
- LaunchConfigOpToOCLPattern<WorkgroupIdZOp>,
- LaunchConfigOpToOCLPattern<GridDimXOp>,
- LaunchConfigOpToOCLPattern<GridDimYOp>,
- LaunchConfigOpToOCLPattern<GridDimZOp>,
- SubgroupOpWorkitemOpToOCLPattern<LaneIdOp>,
- SubgroupOpWorkitemOpToOCLPattern<SubgroupIdOp>,
- SubgroupOpWorkitemOpToOCLPattern<SubgroupSizeOp>, TruncfToOCLPattern,
- ExtfToOCLPattern, MMAMxToOCLPattern, AllocaToGlobalPattern>(
- patterns.getContext());
+ patterns.add<LoadStorePrefetchToOCLPattern<BlockLoad2dOp>,
+ LoadStorePrefetchToOCLPattern<BlockStore2dOp>,
+ LoadStorePrefetchToOCLPattern<BlockPrefetch2dOp>,
+ MMAToOCLPattern, MemfenceToOCLPattern, PrefetchToOCLPattern,
+ LLVMLoadStoreToOCLPattern<LLVM::LoadOp>,
+ LLVMLoadStoreToOCLPattern<LLVM::StoreOp>,
+ BlockLoadStore1DToOCLPattern<BlockLoadOp>,
+ BlockLoadStore1DToOCLPattern<BlockStoreOp>,
+ LaunchConfigOpToOCLPattern<WorkitemIdXOp>,
+ LaunchConfigOpToOCLPattern<WorkitemIdYOp>,
+ LaunchConfigOpToOCLPattern<WorkitemIdZOp>,
+ LaunchConfigOpToOCLPattern<WorkgroupDimXOp>,
+ LaunchConfigOpToOCLPattern<WorkgroupDimYOp>,
+ LaunchConfigOpToOCLPattern<WorkgroupDimZOp>,
+ LaunchConfigOpToOCLPattern<WorkgroupIdXOp>,
+ LaunchConfigOpToOCLPattern<WorkgroupIdYOp>,
+ LaunchConfigOpToOCLPattern<WorkgroupIdZOp>,
+ LaunchConfigOpToOCLPattern<GridDimXOp>,
+ LaunchConfigOpToOCLPattern<GridDimYOp>,
+ LaunchConfigOpToOCLPattern<GridDimZOp>,
+ SubgroupOpWorkitemOpToOCLPattern<LaneIdOp>,
+ SubgroupOpWorkitemOpToOCLPattern<SubgroupIdOp>,
+ SubgroupOpWorkitemOpToOCLPattern<SubgroupSizeOp>,
+ TruncfToOCLPattern, ExtfToOCLPattern, MMAMxToOCLPattern,
+ BitcastShuffleToGenISAPattern, AllocaToGlobalPattern>(
+ patterns.getContext());
}
diff --git a/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
index e8b3c7065f880..1249cc34866cc 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
@@ -384,6 +384,22 @@ LogicalResult ExtfOp::verify() {
return success();
}
+LogicalResult BitcastShuffleOp::verify() {
+ Type srcTy = getSrc().getType();
+ Type resTy = getRes().getType();
+ if (srcTy == resTy)
+ return emitOpError("src and res types must be different");
+
+ auto getTotalBitWidth = [](Type ty) -> unsigned {
+ if (auto vecTy = dyn_cast<VectorType>(ty))
+ return vecTy.getNumElements() * vecTy.getElementTypeBitWidth();
+ return ty.getIntOrFloatBitWidth();
+ };
+ if (getTotalBitWidth(srcTy) != getTotalBitWidth(resTy))
+ return emitOpError("src and res types must have the same total bit width");
+ return success();
+}
+
LogicalResult
XeVMTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError, int O,
StringRef triple, StringRef chip, DictionaryAttr flags,
diff --git a/mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir b/mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir
index 37064b1a9be9c..f85f76979e3c3 100644
--- a/mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir
+++ b/mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir
@@ -569,3 +569,71 @@ llvm.func @subgroup_id() -> i32 {
%1 = xevm.subgroup_id : i32
llvm.return %1 : i32
}
+
+// -----
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v4i32.v8i16(vector<8xi16>) -> vector<4xi32>
+// CHECK-SAME: attributes {convergent, no_unwind, will_return}
+llvm.func @bitcast_shuffle(%a: vector<8xi16>) -> vector<4xi32> {
+ // CHECK: %[[VAR0:.*]] = llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v4i32.v8i16(%[[ARG0:.*]]) {convergent,
+ // CHECK-SAME: function_type = !llvm.func<vector<4xi32> (vector<8xi16>)>, linkage = #llvm.linkage<external>,
+ // CHECK-SAME: no_unwind, sym_name = "llvm.genx.GenISA.SubgroupBitcastShuffle.v4i32.v8i16", visibility_ = 0 : i64, will_return}
+ // CHECK-SAME: : (vector<8xi16>) -> vector<4xi32>
+ %0 = xevm.bitcast_shuffle %a : (vector<8xi16>) -> vector<4xi32>
+ llvm.return %0 : vector<4xi32>
+}
+
+// -----
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2f16.f32(f32) -> vector<2xf16>
+llvm.func @bitcast_shuffle_scalar_src(%a: f32) -> vector<2xf16> {
+ // CHECK: llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2f16.f32({{.*}}) {{{.*}}} : (f32) -> vector<2xf16>
+ %0 = xevm.bitcast_shuffle %a : (f32) -> vector<2xf16>
+ llvm.return %0 : vector<2xf16>
+}
+
+// -----
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.i16.v2i8(vector<2xi8>) -> i16
+llvm.func @bitcast_shuffle_scalar_res(%a: vector<2xi8>) -> i16 {
+ // CHECK: llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.i16.v2i8({{.*}}) {{{.*}}} : (vector<2xi8>) -> i16
+ %0 = xevm.bitcast_shuffle %a : (vector<2xi8>) -> i16
+ llvm.return %0 : i16
+}
+
+// -----
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.i64.v4f16(vector<4xf16>) -> i64
+llvm.func @bitcast_shuffle_f16(%a: vector<4xf16>) -> i64 {
+ // CHECK: llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.i64.v4f16({{.*}}) {{{.*}}} : (vector<4xf16>) -> i64
+ %0 = xevm.bitcast_shuffle %a : (vector<4xf16>) -> i64
+ llvm.return %0 : i64
+}
+
+// -----
+// The intrinsic has no bf16 overload, bf16 is encoded as i16 around the call.
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i32.v4i16(vector<4xi16>) -> vector<2xi32>
+llvm.func @bitcast_shuffle_bf16_src(%a: vector<4xbf16>) -> vector<2xi32> {
+ // CHECK: %[[SRC:.*]] = llvm.bitcast %{{.*}} : vector<4xbf16> to vector<4xi16>
+ // CHECK: llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i32.v4i16(%[[SRC]]) {{{.*}}} : (vector<4xi16>) -> vector<2xi32>
+ %0 = xevm.bitcast_shuffle %a : (vector<4xbf16>) -> vector<2xi32>
+ llvm.return %0 : vector<2xi32>
+}
+
+// -----
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v4i16.v2i32(vector<2xi32>) -> vector<4xi16>
+llvm.func @bitcast_shuffle_bf16_res(%a: vector<2xi32>) -> vector<4xbf16> {
+ // CHECK: %[[RES:.*]] = llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v4i16.v2i32({{.*}}) {{{.*}}} : (vector<2xi32>) -> vector<4xi16>
+ // CHECK: llvm.bitcast %[[RES]] : vector<4xi16> to vector<4xbf16>
+ %0 = xevm.bitcast_shuffle %a : (vector<2xi32>) -> vector<4xbf16>
+ llvm.return %0 : vector<4xbf16>
+}
+
+// -----
+// Encoding bf16 as i16 makes both sides identical: this degenerates into a
+// plain bitcast and the intrinsic, which requires distinct types, is not
+// called.
+// CHECK-LABEL: llvm.func @bitcast_shuffle_bf16_same_width
+// CHECK-NOT: llvm.call
+llvm.func @bitcast_shuffle_bf16_same_width(%a: vector<2xbf16>) -> vector<2xi16> {
+ // CHECK: %[[RES:.*]] = llvm.bitcast %{{.*}} : vector<2xbf16> to vector<2xi16>
+ // CHECK: llvm.return %[[RES]]
+ %0 = xevm.bitcast_shuffle %a : (vector<2xbf16>) -> vector<2xi16>
+ llvm.return %0 : vector<2xi16>
+}
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index 6265f67e594d0..9207772757edd 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -2075,6 +2075,31 @@ llvm.func @invalid_xevm_extf_2(%arg0: i8) {
// -----
+llvm.func @invalid_xevm_bitcast_shuffle_1(%arg0: vector<8xi16>) {
+ // expected-error at +1 {{op src and res types must be different}}
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<8xi16>) -> vector<8xi16>
+ llvm.return
+}
+
+// -----
+
+llvm.func @invalid_xevm_bitcast_shuffle_2(%arg0: vector<8xi16>) {
+ // expected-error at +1 {{op src and res types must have the same total bit width}}
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<8xi16>) -> vector<8xi32>
+ llvm.return
+}
+
+// -----
+
+// The shuffle operates at byte granularity, sub-byte element types are invalid.
+llvm.func @invalid_xevm_bitcast_shuffle_3(%arg0: vector<8xi4>) {
+ // expected-error at +1 {{op operand #0 must be 8-bit signless integer or}}
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<8xi4>) -> vector<4xi8>
+ llvm.return
+}
+
+// -----
+
llvm.func @invalid_xevm_mma_mx(%loaded_c_casted: vector<4xf32>, %loaded_a: vector<8xi16>, %loaded_b_casted: vector<8xi32>, %scale_a: vector<2xi8>, %scale_b: vector<2xi8>) -> vector<8xf32> {
// expected-error at +1 {{op type of C operand must match result type}}
%c_result = xevm.mma_mx %loaded_a, %loaded_b_casted, %scale_a, %scale_b, %loaded_c_casted { shape=<m=8, n=16, k=64>,
diff --git a/mlir/test/Dialect/LLVMIR/xevm.mlir b/mlir/test/Dialect/LLVMIR/xevm.mlir
index fbce5355611d4..4bfa5420f33ba 100644
--- a/mlir/test/Dialect/LLVMIR/xevm.mlir
+++ b/mlir/test/Dialect/LLVMIR/xevm.mlir
@@ -144,6 +144,38 @@ func.func @extf_vector() -> vector<8xbf16> {
return %2 : vector<8xbf16>
}
+// -----
+// CHECK-LABEL: func.func @bitcast_shuffle_vector
+func.func @bitcast_shuffle_vector(%arg0: vector<8xi16>) -> vector<4xi32> {
+ // CHECK: xevm.bitcast_shuffle %{{.*}} : (vector<8xi16>) -> vector<4xi32>
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<8xi16>) -> vector<4xi32>
+ return %0 : vector<4xi32>
+}
+
+// -----
+// CHECK-LABEL: func.func @bitcast_shuffle_scalar_to_vector
+func.func @bitcast_shuffle_scalar_to_vector(%arg0: i32) -> vector<4xi8> {
+ // CHECK: xevm.bitcast_shuffle %{{.*}} : (i32) -> vector<4xi8>
+ %0 = xevm.bitcast_shuffle %arg0 : (i32) -> vector<4xi8>
+ return %0 : vector<4xi8>
+}
+
+// -----
+// CHECK-LABEL: func.func @bitcast_shuffle_vector_to_scalar
+func.func @bitcast_shuffle_vector_to_scalar(%arg0: vector<2xf16>) -> i32 {
+ // CHECK: xevm.bitcast_shuffle %{{.*}} : (vector<2xf16>) -> i32
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<2xf16>) -> i32
+ return %0 : i32
+}
+
+// -----
+// CHECK-LABEL: func.func @bitcast_shuffle_bf16
+func.func @bitcast_shuffle_bf16(%arg0: vector<4xbf16>) -> vector<2xi32> {
+ // CHECK: xevm.bitcast_shuffle %{{.*}} : (vector<4xbf16>) -> vector<2xi32>
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<4xbf16>) -> vector<2xi32>
+ return %0 : vector<2xi32>
+}
+
// -----
// CHECK-LABEL: func.func @memfence()
func.func @memfence() {
diff --git a/mlir/test/Integration/Dialect/XeVM/GPU/xevm_bitcast_shuffle.mlir b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_bitcast_shuffle.mlir
new file mode 100644
index 0000000000000..bb77a9662de93
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_bitcast_shuffle.mlir
@@ -0,0 +1,219 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=lane" \
+// RUN: | mlir-runner \
+// RUN: --shared-libs=%mlir_levelzero_runtime \
+// RUN: --shared-libs=%mlir_runner_utils \
+// RUN: --shared-libs=%mlir_c_runner_utils \
+// RUN: --entry-point-result=void \
+// RUN: | FileCheck %s
+
+// End-to-end test for `xevm.bitcast_shuffle`, which redistributes the bits of
+// the source data across the whole sub-group. Both kernels pin the sub-group
+// size to 16 with `intel_reqd_sub_group_size`, and are launched with 16
+// threads, so a single full sub-group cooperates on the shuffle.
+//
+// Both kernels read their source data from memory. The operation reinterprets
+// the SIMD register layout of the source, so a source that is uniform across
+// the sub-group, a splat constant in particular, is not a meaningful input: it
+// is held in a scalar register and there is no per-lane layout to reinterpret.
+module @bitcast_shuffle attributes {gpu.container_module} {
+
+ gpu.module @kernel {
+ // Reversibility check: a shuffle to another type followed by a shuffle back
+ // to the original type must reproduce the original data. This holds for any
+ // sub-group size, so no assumption is made about the shuffle pattern here.
+ // Lane L owns row L of a 16x4 i32 buffer.
+ gpu.func @shuffle_roundtrip(%ptr: !llvm.ptr<1>) kernel
+ attributes {llvm.intel_reqd_sub_group_size = 16 : i32} {
+ %lane = gpu.lane_id
+ %lane_i64 = arith.index_cast %lane : index to i64
+ %c4 = arith.constant 4 : i64
+ %offset = arith.muli %lane_i64, %c4 : i64
+ %lane_ptr = llvm.getelementptr %ptr[%offset]
+ : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, i32
+ %vec = llvm.load %lane_ptr : !llvm.ptr<1> -> vector<4xi32>
+ %shuffled = xevm.bitcast_shuffle %vec : (vector<4xi32>) -> vector<8xi16>
+ %restored = xevm.bitcast_shuffle %shuffled
+ : (vector<8xi16>) -> vector<4xi32>
+ llvm.store %restored, %lane_ptr : vector<4xi32>, !llvm.ptr<1>
+ gpu.return
+ }
+
+ // Value check of the shuffle pattern itself, with N = 16 lanes.
+ //
+ // Every 16-bit unit of the source is tagged with the position it starts
+ // out at: lane L holds `[c * 16 + L for c in 0..7]`, so a tag names the
+ // (component, lane) pair it comes from. The result is shuffled back to
+ // 16-bit units and widened to i32, so each printed value names the source
+ // unit that ended up there and the whole output is a permutation of
+ // 0..127.
+ //
+ // Number the 16-bit units of the concatenated source stream `u = c * 16 +
+ // L`, so a tag is just its own stream position. Lane L's result component
+ // d is result stream unit `j = d * 16 + L`, covering source units `2j` and
+ // `2j + 1`, hence source stream units `32d + 2L` and `32d + 2L + 1`. The
+ // low half of a result component holds the earlier of the two, as the
+ // concatenation is little endian.
+ gpu.func @shuffle_value(%src: !llvm.ptr<1>, %dst: !llvm.ptr<1>) kernel
+ attributes {llvm.intel_reqd_sub_group_size = 16 : i32} {
+ %lane = gpu.lane_id
+ %lane_i64 = arith.index_cast %lane : index to i64
+ %c8 = arith.constant 8 : i64
+ %offset = arith.muli %lane_i64, %c8 : i64
+ %src_ptr = llvm.getelementptr %src[%offset]
+ : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, i16
+ %vec = llvm.load %src_ptr : !llvm.ptr<1> -> vector<8xi16>
+ %res = xevm.bitcast_shuffle %vec : (vector<8xi16>) -> vector<4xi32>
+ // Split the result back into the 16-bit units it was assembled from and
+ // widen them, so that every tag can be read off the output.
+ %halves = llvm.bitcast %res : vector<4xi32> to vector<8xi16>
+ %wide = arith.extui %halves : vector<8xi16> to vector<8xi32>
+ %dst_ptr = llvm.getelementptr %dst[%offset]
+ : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, i32
+ llvm.store %wide, %dst_ptr : vector<8xi32>, !llvm.ptr<1>
+ gpu.return
+ }
+ }
+
+ func.func @test_roundtrip(%src: memref<16x4xi32>) -> memref<16x4xi32>
+ attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %memref_0 = gpu.alloc() : memref<16x4xi32>
+ gpu.memcpy %memref_0, %src : memref<16x4xi32>, memref<16x4xi32>
+ %0 = memref.extract_aligned_pointer_as_index %memref_0
+ : memref<16x4xi32> -> index
+ %1 = arith.index_cast %0 : index to i64
+ %2 = llvm.inttoptr %1 : i64 to !llvm.ptr
+ %casted = llvm.addrspacecast %2 : !llvm.ptr to !llvm.ptr<1>
+ gpu.launch_func @kernel::@shuffle_roundtrip
+ blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)
+ args(%casted : !llvm.ptr<1>)
+ %dst = memref.alloc() : memref<16x4xi32>
+ gpu.memcpy %dst, %memref_0 : memref<16x4xi32>, memref<16x4xi32>
+ gpu.dealloc %memref_0 : memref<16x4xi32>
+ return %dst : memref<16x4xi32>
+ }
+
+ func.func @test_shuffle(%src: memref<16x8xi16>) -> memref<16x8xi32>
+ attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %src_gpu = gpu.alloc() : memref<16x8xi16>
+ gpu.memcpy %src_gpu, %src : memref<16x8xi16>, memref<16x8xi16>
+ %dst_gpu = gpu.alloc() : memref<16x8xi32>
+ %0 = memref.extract_aligned_pointer_as_index %src_gpu
+ : memref<16x8xi16> -> index
+ %1 = arith.index_cast %0 : index to i64
+ %2 = llvm.inttoptr %1 : i64 to !llvm.ptr
+ %src_ptr = llvm.addrspacecast %2 : !llvm.ptr to !llvm.ptr<1>
+ %3 = memref.extract_aligned_pointer_as_index %dst_gpu
+ : memref<16x8xi32> -> index
+ %4 = arith.index_cast %3 : index to i64
+ %5 = llvm.inttoptr %4 : i64 to !llvm.ptr
+ %dst_ptr = llvm.addrspacecast %5 : !llvm.ptr to !llvm.ptr<1>
+ gpu.launch_func @kernel::@shuffle_value
+ blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)
+ args(%src_ptr : !llvm.ptr<1>, %dst_ptr : !llvm.ptr<1>)
+ %dst = memref.alloc() : memref<16x8xi32>
+ gpu.memcpy %dst, %dst_gpu : memref<16x8xi32>, memref<16x8xi32>
+ gpu.dealloc %src_gpu : memref<16x8xi16>
+ gpu.dealloc %dst_gpu : memref<16x8xi32>
+ return %dst : memref<16x8xi32>
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c1_i32 = arith.constant 1 : i32
+ %c4_i32 = arith.constant 4 : i32
+ %c16_i16 = arith.constant 16 : i16
+
+ // Fill the buffer with 1..64 in row-major order, so that every lane holds
+ // four distinct values and no value is repeated across the sub-group.
+ %A = memref.alloc() : memref<16x4xi32>
+ scf.for %i = %c0 to %c16 step %c1 {
+ scf.for %j = %c0 to %c4 step %c1 {
+ %i_i32 = arith.index_cast %i : index to i32
+ %j_i32 = arith.index_cast %j : index to i32
+ %row = arith.muli %i_i32, %c4_i32 : i32
+ %idx = arith.addi %row, %j_i32 : i32
+ %v = arith.addi %idx, %c1_i32 : i32
+ memref.store %v, %A[%i, %j] : memref<16x4xi32>
+ }
+ }
+
+ %B = call @test_roundtrip(%A) : (memref<16x4xi32>) -> memref<16x4xi32>
+ %B_cast = memref.cast %B : memref<16x4xi32> to memref<*xi32>
+ call @printMemrefI32(%B_cast) : (memref<*xi32>) -> ()
+
+ // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
+ // CHECK: [1, 2, 3, 4]
+ // CHECK: [5, 6, 7, 8]
+ // CHECK: [9, 10, 11, 12]
+ // CHECK: [13, 14, 15, 16]
+ // CHECK: [17, 18, 19, 20]
+ // CHECK: [21, 22, 23, 24]
+ // CHECK: [25, 26, 27, 28]
+ // CHECK: [29, 30, 31, 32]
+ // CHECK: [33, 34, 35, 36]
+ // CHECK: [37, 38, 39, 40]
+ // CHECK: [41, 42, 43, 44]
+ // CHECK: [45, 46, 47, 48]
+ // CHECK: [49, 50, 51, 52]
+ // CHECK: [53, 54, 55, 56]
+ // CHECK: [57, 58, 59, 60]
+ // CHECK: [61, 62, 63, 64]
+
+ // Tag the 16-bit unit held by lane L as component c with its own position
+ // in the concatenated source stream, `c * 16 + L`.
+ %C = memref.alloc() : memref<16x8xi16>
+ scf.for %l = %c0 to %c16 step %c1 {
+ scf.for %c = %c0 to %c8 step %c1 {
+ %l_i16 = arith.index_cast %l : index to i16
+ %c_i16 = arith.index_cast %c : index to i16
+ %col = arith.muli %c_i16, %c16_i16 : i16
+ %tag = arith.addi %col, %l_i16 : i16
+ memref.store %tag, %C[%l, %c] : memref<16x8xi16>
+ }
+ }
+
+ %D = call @test_shuffle(%C) : (memref<16x8xi16>) -> memref<16x8xi32>
+ %D_cast = memref.cast %D : memref<16x8xi32> to memref<*xi32>
+ call @printMemrefI32(%D_cast) : (memref<*xi32>) -> ()
+
+ // Lane L takes source stream units 32d + 2L and 32d + 2L + 1 for its
+ // result component d, so the lower half of the sub-group reads the even
+ // source components and the upper half the odd ones.
+ //
+ // A result that reproduces the source rows instead, that is row L reading
+ // [L, 16+L, 32+L, ...], means no data crossed lanes and the shuffle
+ // degenerated into a per-lane bitcast.
+ // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
+ // CHECK: [0, 1, 32, 33, 64, 65, 96, 97]
+ // CHECK: [2, 3, 34, 35, 66, 67, 98, 99]
+ // CHECK: [4, 5, 36, 37, 68, 69, 100, 101]
+ // CHECK: [6, 7, 38, 39, 70, 71, 102, 103]
+ // CHECK: [8, 9, 40, 41, 72, 73, 104, 105]
+ // CHECK: [10, 11, 42, 43, 74, 75, 106, 107]
+ // CHECK: [12, 13, 44, 45, 76, 77, 108, 109]
+ // CHECK: [14, 15, 46, 47, 78, 79, 110, 111]
+ // CHECK: [16, 17, 48, 49, 80, 81, 112, 113]
+ // CHECK: [18, 19, 50, 51, 82, 83, 114, 115]
+ // CHECK: [20, 21, 52, 53, 84, 85, 116, 117]
+ // CHECK: [22, 23, 54, 55, 86, 87, 118, 119]
+ // CHECK: [24, 25, 56, 57, 88, 89, 120, 121]
+ // CHECK: [26, 27, 58, 59, 90, 91, 122, 123]
+ // CHECK: [28, 29, 60, 61, 92, 93, 124, 125]
+ // CHECK: [30, 31, 62, 63, 94, 95, 126, 127]
+
+ memref.dealloc %A : memref<16x4xi32>
+ memref.dealloc %B : memref<16x4xi32>
+ memref.dealloc %C : memref<16x8xi16>
+ memref.dealloc %D : memref<16x8xi32>
+ return
+ }
+ func.func private @printMemrefI32(%ptr : memref<*xi32>) attributes { llvm.emit_c_interface }
+}
>From 86164d088992f143f50a41f89b33e23a5f3dc5aa Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 10 Aug 2026 15:24:51 +0000
Subject: [PATCH 2/2] [MLIR][XeGPU] Lower xegpu.lane_shuffle to
xevm.bitcast_shuffle
Add a `convert-xegpu-to-xevm` pattern that lowers `xegpu.lane_shuffle` onto
`xevm.bitcast_shuffle`.
`xevm.bitcast_shuffle` concatenates the components of its operand across the
subgroup, the first component of every lane first, and then hands chunks the
size of a result component back out to the lanes in order. Numbering the
elements of a `vector<NxT>` fragment held by lane `i` of a subgroup of size
`S` by their logical position, that concatenation is exactly the `pack` mode
input numbering `j * S + i`. Taking the result as a single `N * width(T)` bit
scalar then hands lane `i` the logical positions `i * N .. i * N + N - 1`,
which is the `pack` mode output numbering.
So `pack` becomes a vector-to-scalar `xevm.bitcast_shuffle` followed by a
bitcast back to the fragment type, and `unpack`, being its inverse, becomes a
bitcast to the scalar followed by a scalar-to-vector `xevm.bitcast_shuffle`.
The shuffle redistributes whole bytes between the lanes, so sub-byte element
types cannot be shuffled. The pattern reports a conversion failure for them,
which covers fp4 fragments. Fragments that are not 8, 16, 32 or 64 bits wide
are rejected the same way, as there is no integer type to pack them into.
---
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 52 +++++++++++++
.../XeGPUToXeVM/failed_conversion.mlir | 26 +++++++
.../Conversion/XeGPUToXeVM/lane_shuffle.mlir | 76 +++++++++++++++++++
3 files changed, 154 insertions(+)
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/lane_shuffle.mlir
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 78d99cf88b768..681dbd86234b0 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -1341,6 +1341,57 @@ class TruncfToXeVMPattern : public OpConversionPattern<arith::TruncFOp> {
}
};
+// Lowers `xegpu.lane_shuffle` to `xevm.bitcast_shuffle`.
+//
+// `xevm.bitcast_shuffle` concatenates the components of its operand across the
+// subgroup, the first component of every lane first, and then hands chunks the
+// size of a result component back out to the lanes in order. Numbering the
+// elements of a `vector<NxT>` fragment held by lane `i` of a subgroup of size
+// `S` by their logical position, that concatenation is exactly the `pack` mode
+// input numbering `j * S + i`. Taking the result as a single `N * width(T)` bit
+// scalar then hands lane `i` the logical positions `i * N .. i * N + N - 1`,
+// which is the `pack` mode output numbering.
+//
+// So `pack` is a vector-to-scalar `xevm.bitcast_shuffle` followed by a bitcast
+// back to the fragment type, and `unpack`, being its inverse, is a bitcast to
+// the scalar followed by a scalar-to-vector `xevm.bitcast_shuffle`.
+class LaneShuffleToXeVMPattern
+ : public OpConversionPattern<xegpu::LaneShuffleOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::LaneShuffleOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto vecTy = dyn_cast<VectorType>(adaptor.getSource().getType());
+ if (!vecTy)
+ return rewriter.notifyMatchFailure(op, "Expected a vector fragment.");
+ // The shuffle redistributes whole bytes between the lanes, so sub-byte
+ // element types, fp4 in particular, cannot be shuffled.
+ unsigned elemBits = vecTy.getElementTypeBitWidth();
+ if (elemBits < 8)
+ return rewriter.notifyMatchFailure(
+ op, "Sub-byte element types are not supported.");
+ int64_t fragmentBits = vecTy.getNumElements() * elemBits;
+ if (fragmentBits > 64 || !llvm::isPowerOf2_64(fragmentBits))
+ return rewriter.notifyMatchFailure(
+ op, "Expected a fragment of 8, 16, 32 or 64 bits.");
+
+ Location loc = op.getLoc();
+ Type packedTy = rewriter.getIntegerType(fragmentBits);
+ Value res;
+ if (op.getMode() == xegpu::LaneShuffleMode::Pack) {
+ res = xevm::BitcastShuffleOp::create(rewriter, loc, packedTy,
+ adaptor.getSource());
+ res = LLVM::BitcastOp::create(rewriter, loc, vecTy, res);
+ } else {
+ Value packed =
+ LLVM::BitcastOp::create(rewriter, loc, packedTy, adaptor.getSource());
+ res = xevm::BitcastShuffleOp::create(rewriter, loc, vecTy, packed);
+ }
+ rewriter.replaceOp(op, res);
+ return success();
+ }
+};
+
//===----------------------------------------------------------------------===//
// Pass Definition
//===----------------------------------------------------------------------===//
@@ -1664,4 +1715,5 @@ void mlir::populateXeGPUToXeVMConversionPatterns(
patterns.add<DpasMxToXeVMPattern>(typeConverter, patterns.getContext());
patterns.add<ExtfToXeVMPattern, TruncfToXeVMPattern>(typeConverter,
patterns.getContext());
+ patterns.add<LaneShuffleToXeVMPattern>(typeConverter, patterns.getContext());
}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
index cabc65aa0e41d..93bda1dfc0dca 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -28,3 +28,29 @@ gpu.module @test_kernel {
gpu.return
}
}
+
+// -----
+
+// Verify that xegpu.lane_shuffle of a sub-byte element type is rejected: the
+// shuffle redistributes whole bytes between the lanes, so fp4 fragments cannot
+// be shuffled.
+
+gpu.module @test_kernel {
+ gpu.func @lane_shuffle_f4(%a: vector<4xf4E2M1FN>) {
+ // expected-error at +1 {{failed to legalize operation 'xegpu.lane_shuffle' that was explicitly marked illegal}}
+ %0 = xegpu.lane_shuffle %a pack : vector<4xf4E2M1FN>
+ gpu.return
+ }
+}
+
+// -----
+
+// Verify that a xegpu.lane_shuffle fragment wider than 64 bits is rejected.
+
+gpu.module @test_kernel {
+ gpu.func @lane_shuffle_too_wide(%a: vector<4xi32>) {
+ // expected-error at +1 {{failed to legalize operation 'xegpu.lane_shuffle' that was explicitly marked illegal}}
+ %0 = xegpu.lane_shuffle %a pack : vector<4xi32>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle.mlir b/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle.mlir
new file mode 100644
index 0000000000000..e0ff25c34fb99
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle.mlir
@@ -0,0 +1,76 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm -split-input-file %s | FileCheck %s
+
+gpu.module @test {
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_i16
+// CHECK-SAME: %[[SRC:.*]]: vector<2xi16>
+gpu.func @lane_shuffle_pack_i16(%a: vector<2xi16>) -> vector<2xi16> {
+ // CHECK: %[[PACKED:.*]] = xevm.bitcast_shuffle %[[SRC]] : (vector<2xi16>) -> i32
+ // CHECK: %[[RES:.*]] = llvm.bitcast %[[PACKED]] : i32 to vector<2xi16>
+ // CHECK: gpu.return %[[RES]]
+ %0 = xegpu.lane_shuffle %a pack : vector<2xi16>
+ gpu.return %0 : vector<2xi16>
+}
+}
+
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: gpu.func @lane_shuffle_unpack_i16
+// CHECK-SAME: %[[SRC:.*]]: vector<2xi16>
+gpu.func @lane_shuffle_unpack_i16(%a: vector<2xi16>) -> vector<2xi16> {
+ // CHECK: %[[PACKED:.*]] = llvm.bitcast %[[SRC]] : vector<2xi16> to i32
+ // CHECK: %[[RES:.*]] = xevm.bitcast_shuffle %[[PACKED]] : (i32) -> vector<2xi16>
+ // CHECK: gpu.return %[[RES]]
+ %0 = xegpu.lane_shuffle %a unpack : vector<2xi16>
+ gpu.return %0 : vector<2xi16>
+}
+}
+
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_f16
+gpu.func @lane_shuffle_pack_f16(%a: vector<4xf16>) -> vector<4xf16> {
+ // CHECK: %[[PACKED:.*]] = xevm.bitcast_shuffle %{{.*}} : (vector<4xf16>) -> i64
+ // CHECK: llvm.bitcast %[[PACKED]] : i64 to vector<4xf16>
+ %0 = xegpu.lane_shuffle %a pack : vector<4xf16>
+ gpu.return %0 : vector<4xf16>
+}
+}
+
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_bf16
+gpu.func @lane_shuffle_pack_bf16(%a: vector<2xbf16>) -> vector<2xbf16> {
+ // CHECK: %[[PACKED:.*]] = xevm.bitcast_shuffle %{{.*}} : (vector<2xbf16>) -> i32
+ // CHECK: llvm.bitcast %[[PACKED]] : i32 to vector<2xbf16>
+ %0 = xegpu.lane_shuffle %a pack : vector<2xbf16>
+ gpu.return %0 : vector<2xbf16>
+}
+}
+
+// -----
+
+// The f8 element type is converted to a same-width integer before the shuffle.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @lane_shuffle_unpack_f8
+gpu.func @lane_shuffle_unpack_f8(%a: vector<4xf8E5M2>) -> vector<4xf8E5M2> {
+ // CHECK: %[[PACKED:.*]] = llvm.bitcast %{{.*}} : vector<4xi8> to i32
+ // CHECK: xevm.bitcast_shuffle %[[PACKED]] : (i32) -> vector<4xi8>
+ %0 = xegpu.lane_shuffle %a unpack : vector<4xf8E5M2>
+ gpu.return %0 : vector<4xf8E5M2>
+}
+}
+
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: gpu.func @lane_shuffle_pack_i8
+gpu.func @lane_shuffle_pack_i8(%a: vector<8xi8>) -> vector<8xi8> {
+ // CHECK: %[[PACKED:.*]] = xevm.bitcast_shuffle %{{.*}} : (vector<8xi8>) -> i64
+ // CHECK: llvm.bitcast %[[PACKED]] : i64 to vector<8xi8>
+ %0 = xegpu.lane_shuffle %a pack : vector<8xi8>
+ gpu.return %0 : vector<8xi8>
+}
+}
More information about the Mlir-commits
mailing list