[Mlir-commits] [mlir] [MLIR][XeVM] Add xevm.bitcast_shuffle op and lowering to LLVM (PR #215303)
Sang Ik Lee
llvmlistbot at llvm.org
Tue Aug 18 14:48:03 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/215303
>From f5c950f1796a76916b21ffe86e95e3355c0bf9c9 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] [MLIR][XeVM] Add xevm.bitcast_shuffle op and lowering to LLVM
Add `xevm.bitcast_shuffle`, matching the SPIR-V OpSubgroupBitcastShuffleINTEL
instruction from SPV_INTEL_subgroup_bitcast_shuffle. It performs a bit-preserving
conversion and shuffle of its operand across a subgroup: all invocations
cooperate, the components of the operand are concatenated in subgroup-local-id
order, and the resulting bit stream is handed back out in chunks the size of a
result component.
// pack: gather one value per lane out of the components spread across the subgroup
%scalar = xevm.bitcast_shuffle %vec : (vector<4xi16>) -> i64
// unpack: the exact inverse
%vec2 = xevm.bitcast_shuffle %scalar : (i64) -> vector<4xi16>
Exactly one side is a 1-D vector and the other a scalar, giving these two forms; a
vector-to-vector repack and a plain scalar-to-scalar bitcast are both rejected.
Both sides must have the same total bit width, and only i8/i16/i32/i64 are
accepted: the op is bit-preserving, so a producer holding floating point data
bitcasts it to a same-width integer first, which keeps the set of overloads small.
The op is not lowered to the SPIR-V instruction. `convert-xevm-to-llvm` emits a
call to the IGC intrinsic `llvm.genx.GenISA.SubgroupBitcastShuffle`, overloaded on
both operand and result type.
---
mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td | 85 +++++++
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 99 +++++---
mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp | 22 ++
.../Conversion/XeVMToLLVM/xevm-to-llvm.mlir | 48 ++++
mlir/test/Dialect/LLVMIR/invalid.mlir | 61 +++++
mlir/test/Dialect/LLVMIR/xevm.mlir | 40 ++++
.../XeVM/GPU/xevm_bitcast_shuffle.mlir | 219 ++++++++++++++++++
7 files changed, 548 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 c52a1c0c94402..2c4e43a20a0bf 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
@@ -741,6 +741,91 @@ def XeVM_MMAMxOp
let hasVerifier = 1;
}
+def XeVM_BitcastShuffleElemType : AnyTypeOf<[I8, I16, I32, I64]>;
+
+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.
+
+ Exactly one of `src` and `res` is a 1D vector and the other is a scalar, so
+ the operation comes in two forms:
+
+ * a *pack*, taking a vector and returning a scalar, which gathers the
+ components spread across the subgroup into one value per invocation;
+ * an *unpack*, taking a scalar and returning a vector, which is its exact
+ inverse.
+
+ The total number of bits of `res` must equal the total number of bits of
+ `src`. Only the integer types `i8`, `i16`, `i32` and `i64` are accepted, on
+ either side. The operation is bit-preserving and so does not depend on how
+ the bits are interpreted; a producer holding floating point data is expected
+ to bitcast it to a same-width integer type beforehand and to bitcast the
+ result back. This keeps the set of overloads the operation lowers to small,
+ and it is why sub-byte types such as f4 and i4 are excluded as well: the
+ shuffle is performed at byte granularity. Note that the packed side is a
+ single value of one of the supported scalar types, which bounds the vector
+ side at 64 bits in total.
+
+ 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. A scalar `src` counts as having a single component. 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.
+
+ A pack followed by an unpack back to the original type, or the other way
+ round, yields the original source data.
+
+ Example:
+ ```mlir
+ // pack
+ %scalar = xevm.bitcast_shuffle %vec : (vector<4xi16>) -> i64
+ // unpack
+ %vec2 = xevm.bitcast_shuffle %scalar : (i64) -> vector<4xi16>
+ ```
+
+ Shuffling a `vector<4xbf16>` fragment, for example, is expressed by
+ bitcasting to and from `vector<4xi16>` around the operation:
+
+ ```mlir
+ %bits = llvm.bitcast %frag : vector<4xbf16> to vector<4xi16>
+ %packed = xevm.bitcast_shuffle %bits : (vector<4xi16>) -> i64
+ %res = llvm.bitcast %packed : i64 to vector<4xbf16>
+ ```
+ }];
+
+ 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..27ae552e83670 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -99,6 +99,22 @@ std::string mangle(StringRef baseName, ArrayRef<Type> types,
return os.str();
}
+// Returns the mangling of `ty` used to name an overloaded `llvm.genx.GenISA.*`
+// intrinsic: `i32`, `v8i16`, ... Note that this is IGC's own scheme for its
+// intrinsics, not the Itanium mangling used for the SPIR-V friendly and OCL
+// builtins that `mangle` above produces.
+std::string getGenISATypeMangling(Type ty) {
+ return TypeSwitch<Type, std::string>(ty)
+ .Case([](VectorType ty) -> std::string {
+ return "v" + std::to_string(ty.getNumElements()) +
+ getGenISATypeMangling(ty.getElementType());
+ })
+ .Case([](IntegerType ty) -> std::string {
+ return "i" + std::to_string(ty.getWidth());
+ })
+ .DefaultUnreachable("unhandled type for GenISA mangling");
+}
+
std::string builtinElemType(ElemType elemType) {
switch (elemType) {
case ElemType::BF8:
@@ -1520,6 +1536,37 @@ 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`.
+//
+// Only integer types reach here: the op accepts nothing else, so a producer
+// holding floating point data bitcasts it to a same-width integer beforehand.
+class BitcastShuffleToGenISAPattern
+ : public OpConversionPattern<BitcastShuffleOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(BitcastShuffleOp op, BitcastShuffleOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ Type srcTy = op.getSrc().getType();
+ Type resTy = op.getRes().getType();
+
+ std::string fnName = "llvm.genx.GenISA.SubgroupBitcastShuffle." +
+ getGenISATypeMangling(resTy) + "." +
+ getGenISATypeMangling(srcTy);
+
+ Value result = createDeviceFunctionCall(
+ rewriter, fnName, resTy, {srcTy}, {adaptor.getSrc()}, {},
+ convergentNoUnwindWillReturnAttrs, op.getOperation())
+ ->getResult(0);
+
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+};
+
class AllocaToGlobalPattern : public OpConversionPattern<LLVM::AllocaOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
@@ -1787,30 +1834,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..1cced3c2c3d45 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
@@ -384,6 +384,28 @@ LogicalResult ExtfOp::verify() {
return success();
}
+LogicalResult BitcastShuffleOp::verify() {
+ Type srcTy = getSrc().getType();
+ Type resTy = getRes().getType();
+ auto srcVecTy = dyn_cast<VectorType>(srcTy);
+ auto resVecTy = dyn_cast<VectorType>(resTy);
+ // Only a pack (vector -> scalar) and an unpack (scalar -> vector) are
+ // supported, so exactly one side is a vector.
+ if (static_cast<bool>(srcVecTy) == static_cast<bool>(resVecTy))
+ return emitOpError("expected exactly one of src and res to be a vector: a "
+ "pack takes a vector and returns a scalar, an unpack "
+ "takes a scalar and returns a vector");
+
+ 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..a69b5970f9da4 100644
--- a/mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir
+++ b/mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir
@@ -569,3 +569,51 @@ llvm.func @subgroup_id() -> i32 {
%1 = xevm.subgroup_id : i32
llvm.return %1 : i32
}
+
+// -----
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.i64.v4i16(vector<4xi16>) -> i64
+// CHECK-SAME: attributes {convergent, no_unwind, will_return}
+llvm.func @bitcast_shuffle(%a: vector<4xi16>) -> i64 {
+ // CHECK: %[[VAR0:.*]] = llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.i64.v4i16(%[[ARG0:.*]]) {convergent,
+ // CHECK-SAME: function_type = !llvm.func<i64 (vector<4xi16>)>, linkage = #llvm.linkage<external>,
+ // CHECK-SAME: no_unwind, sym_name = "llvm.genx.GenISA.SubgroupBitcastShuffle.i64.v4i16", visibility_ = 0 : i64, will_return}
+ // CHECK-SAME: : (vector<4xi16>) -> i64
+ %0 = xevm.bitcast_shuffle %a : (vector<4xi16>) -> i64
+ llvm.return %0 : i64
+}
+
+// -----
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i16.i32(i32) -> vector<2xi16>
+llvm.func @bitcast_shuffle_scalar_src(%a: i32) -> vector<2xi16> {
+ // CHECK: llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i16.i32({{.*}}) {{{.*}}} : (i32) -> vector<2xi16>
+ %0 = xevm.bitcast_shuffle %a : (i32) -> vector<2xi16>
+ llvm.return %0 : vector<2xi16>
+}
+
+// -----
+// 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.v2i32(vector<2xi32>) -> i64
+llvm.func @bitcast_shuffle_i32(%a: vector<2xi32>) -> i64 {
+ // CHECK: llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.i64.v2i32({{.*}}) {{{.*}}} : (vector<2xi32>) -> i64
+ %0 = xevm.bitcast_shuffle %a : (vector<2xi32>) -> i64
+ llvm.return %0 : i64
+}
+
+// -----
+// A pack and an unpack of the same types round-trip through two intrinsic calls.
+// CHECK-LABEL: llvm.func @bitcast_shuffle_roundtrip
+llvm.func @bitcast_shuffle_roundtrip(%a: vector<2xi16>) -> vector<2xi16> {
+ // CHECK: %[[PACKED:.*]] = llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.i32.v2i16({{.*}}) {{{.*}}} : (vector<2xi16>) -> i32
+ // CHECK: %[[RES:.*]] = llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i16.i32(%[[PACKED]]) {{{.*}}} : (i32) -> vector<2xi16>
+ // CHECK: llvm.return %[[RES]]
+ %0 = xevm.bitcast_shuffle %a : (vector<2xi16>) -> i32
+ %1 = xevm.bitcast_shuffle %0 : (i32) -> vector<2xi16>
+ llvm.return %1 : vector<2xi16>
+}
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index cf91beb67a96c..3eaf7d2ae52b6 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -2089,6 +2089,67 @@ llvm.func @invalid_xevm_extf_2(%arg0: i8) {
// -----
+llvm.func @invalid_xevm_bitcast_shuffle_1(%arg0: vector<8xi16>) {
+ // expected-error at +1 {{op expected exactly one of src and res to be a vector}}
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<8xi16>) -> vector<8xi16>
+ llvm.return
+}
+
+// -----
+
+// Only a pack and an unpack are supported, a vector to vector repack is not.
+llvm.func @invalid_xevm_bitcast_shuffle_repack(%arg0: vector<8xi16>) {
+ // expected-error at +1 {{op expected exactly one of src and res to be a vector}}
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<8xi16>) -> vector<4xi32>
+ llvm.return
+}
+
+// -----
+
+// Neither is a plain scalar to scalar bitcast.
+llvm.func @invalid_xevm_bitcast_shuffle_scalar(%arg0: i32) {
+ // expected-error at +1 {{op expected exactly one of src and res to be a vector}}
+ %0 = xevm.bitcast_shuffle %arg0 : (i32) -> i32
+ llvm.return
+}
+
+// -----
+
+// The op is bit-preserving and only accepts integer types; a producer holding
+// floating point data bitcasts it to a same-width integer beforehand.
+llvm.func @invalid_xevm_bitcast_shuffle_float_res(%arg0: vector<2xi16>) {
+ // expected-error at +1 {{op result #0 must be 8-bit signless integer or}}
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<2xi16>) -> f32
+ llvm.return
+}
+
+// -----
+
+llvm.func @invalid_xevm_bitcast_shuffle_float_elem(%arg0: vector<4xbf16>) {
+ // expected-error at +1 {{op operand #0 must be 8-bit signless integer or}}
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<4xbf16>) -> i64
+ 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>) -> i64
+ 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..668801ceb3421 100644
--- a/mlir/test/Dialect/LLVMIR/xevm.mlir
+++ b/mlir/test/Dialect/LLVMIR/xevm.mlir
@@ -144,6 +144,46 @@ func.func @extf_vector() -> vector<8xbf16> {
return %2 : vector<8xbf16>
}
+// -----
+// CHECK-LABEL: func.func @bitcast_shuffle_pack
+func.func @bitcast_shuffle_pack(%arg0: vector<4xi16>) -> i64 {
+ // CHECK: xevm.bitcast_shuffle %{{.*}} : (vector<4xi16>) -> i64
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<4xi16>) -> i64
+ return %0 : i64
+}
+
+// -----
+// CHECK-LABEL: func.func @bitcast_shuffle_unpack
+func.func @bitcast_shuffle_unpack(%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_pack_i8
+func.func @bitcast_shuffle_pack_i8(%arg0: vector<2xi8>) -> i16 {
+ // CHECK: xevm.bitcast_shuffle %{{.*}} : (vector<2xi8>) -> i16
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<2xi8>) -> i16
+ return %0 : i16
+}
+
+// -----
+// CHECK-LABEL: func.func @bitcast_shuffle_pack_i32
+func.func @bitcast_shuffle_pack_i32(%arg0: vector<2xi32>) -> i64 {
+ // CHECK: xevm.bitcast_shuffle %{{.*}} : (vector<2xi32>) -> i64
+ %0 = xevm.bitcast_shuffle %arg0 : (vector<2xi32>) -> i64
+ return %0 : i64
+}
+
+// -----
+// CHECK-LABEL: func.func @bitcast_shuffle_unpack_i16
+func.func @bitcast_shuffle_unpack_i16(%arg0: i64) -> vector<4xi16> {
+ // CHECK: xevm.bitcast_shuffle %{{.*}} : (i64) -> vector<4xi16>
+ %0 = xevm.bitcast_shuffle %arg0 : (i64) -> vector<4xi16>
+ return %0 : vector<4xi16>
+}
+
// -----
// 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..447c0e995cba7
--- /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. The operation comes in two forms:
+// a pack, taking a vector and returning a scalar, and an unpack, which is its
+// inverse. 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 pack followed by an unpack 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 16x2 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
+ %c2 = arith.constant 2 : i64
+ %offset = arith.muli %lane_i64, %c2 : i64
+ %lane_ptr = llvm.getelementptr %ptr[%offset]
+ : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, i32
+ %vec = llvm.load %lane_ptr : !llvm.ptr<1> -> vector<2xi32>
+ %packed = xevm.bitcast_shuffle %vec : (vector<2xi32>) -> i64
+ %restored = xevm.bitcast_shuffle %packed : (i64) -> vector<2xi32>
+ llvm.store %restored, %lane_ptr : vector<2xi32>, !llvm.ptr<1>
+ gpu.return
+ }
+
+ // Value check of the pack 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..3]`, so a tag names the
+ // (component, lane) pair it comes from. The packed result is split back into
+ // 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..63.
+ //
+ // Number the 16-bit units of the concatenated source stream `u = c * 16 +
+ // L`, so a tag is just its own stream position. The result is a scalar, so
+ // it has a single component of C = 64 bits, and lane L receives result
+ // stream bits `[64L, 64L + 64)`, that is source stream units `4L` through
+ // `4L + 3`. The low half of the packed value holds the earliest of them, 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
+ %c4 = arith.constant 4 : i64
+ %offset = arith.muli %lane_i64, %c4 : i64
+ %src_ptr = llvm.getelementptr %src[%offset]
+ : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, i16
+ %vec = llvm.load %src_ptr : !llvm.ptr<1> -> vector<4xi16>
+ %res = xevm.bitcast_shuffle %vec : (vector<4xi16>) -> i64
+ // Split the packed 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 : i64 to vector<4xi16>
+ %wide = arith.extui %halves : vector<4xi16> to vector<4xi32>
+ %dst_ptr = llvm.getelementptr %dst[%offset]
+ : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, i32
+ llvm.store %wide, %dst_ptr : vector<4xi32>, !llvm.ptr<1>
+ gpu.return
+ }
+ }
+
+ func.func @test_roundtrip(%src: memref<16x2xi32>) -> memref<16x2xi32>
+ attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %memref_0 = gpu.alloc() : memref<16x2xi32>
+ gpu.memcpy %memref_0, %src : memref<16x2xi32>, memref<16x2xi32>
+ %0 = memref.extract_aligned_pointer_as_index %memref_0
+ : memref<16x2xi32> -> 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<16x2xi32>
+ gpu.memcpy %dst, %memref_0 : memref<16x2xi32>, memref<16x2xi32>
+ gpu.dealloc %memref_0 : memref<16x2xi32>
+ return %dst : memref<16x2xi32>
+ }
+
+ func.func @test_shuffle(%src: memref<16x4xi16>) -> memref<16x4xi32>
+ attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %src_gpu = gpu.alloc() : memref<16x4xi16>
+ gpu.memcpy %src_gpu, %src : memref<16x4xi16>, memref<16x4xi16>
+ %dst_gpu = gpu.alloc() : memref<16x4xi32>
+ %0 = memref.extract_aligned_pointer_as_index %src_gpu
+ : memref<16x4xi16> -> 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<16x4xi32> -> 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<16x4xi32>
+ gpu.memcpy %dst, %dst_gpu : memref<16x4xi32>, memref<16x4xi32>
+ gpu.dealloc %src_gpu : memref<16x4xi16>
+ gpu.dealloc %dst_gpu : memref<16x4xi32>
+ return %dst : memref<16x4xi32>
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c4 = arith.constant 4 : index
+ %c16 = arith.constant 16 : index
+ %c1_i32 = arith.constant 1 : i32
+ %c2_i32 = arith.constant 2 : i32
+ %c16_i16 = arith.constant 16 : i16
+
+ // Fill the buffer with 1..32 in row-major order, so that every lane holds
+ // two distinct values and no value is repeated across the sub-group.
+ %A = memref.alloc() : memref<16x2xi32>
+ scf.for %i = %c0 to %c16 step %c1 {
+ scf.for %j = %c0 to %c2 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, %c2_i32 : i32
+ %idx = arith.addi %row, %j_i32 : i32
+ %v = arith.addi %idx, %c1_i32 : i32
+ memref.store %v, %A[%i, %j] : memref<16x2xi32>
+ }
+ }
+
+ %B = call @test_roundtrip(%A) : (memref<16x2xi32>) -> memref<16x2xi32>
+ %B_cast = memref.cast %B : memref<16x2xi32> to memref<*xi32>
+ call @printMemrefI32(%B_cast) : (memref<*xi32>) -> ()
+
+ // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
+ // CHECK: [1, 2]
+ // CHECK: [3, 4]
+ // CHECK: [5, 6]
+ // CHECK: [7, 8]
+ // CHECK: [9, 10]
+ // CHECK: [11, 12]
+ // CHECK: [13, 14]
+ // CHECK: [15, 16]
+ // CHECK: [17, 18]
+ // CHECK: [19, 20]
+ // CHECK: [21, 22]
+ // CHECK: [23, 24]
+ // CHECK: [25, 26]
+ // CHECK: [27, 28]
+ // CHECK: [29, 30]
+ // CHECK: [31, 32]
+
+ // Tag the 16-bit unit held by lane L as component c with its own position
+ // in the concatenated source stream, `c * 16 + L`. Row L of the input is
+ // therefore [L, 16+L, 32+L, 48+L].
+ %C = memref.alloc() : memref<16x4xi16>
+ scf.for %l = %c0 to %c16 step %c1 {
+ scf.for %c = %c0 to %c4 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<16x4xi16>
+ }
+ }
+
+ %D = call @test_shuffle(%C) : (memref<16x4xi16>) -> memref<16x4xi32>
+ %D_cast = memref.cast %D : memref<16x4xi32> to memref<*xi32>
+ call @printMemrefI32(%D_cast) : (memref<*xi32>) -> ()
+
+ // The packed result of lane L holds source stream units 4L through 4L + 3,
+ // so the output is the source stream laid out contiguously per lane.
+ //
+ // A result that reproduces the input rows instead, that is row L reading
+ // [L, 16+L, 32+L, 48+L], means no data crossed lanes and the pack
+ // degenerated into a per-lane bitcast.
+ // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
+ // CHECK: [0, 1, 2, 3]
+ // CHECK: [4, 5, 6, 7]
+ // CHECK: [8, 9, 10, 11]
+ // CHECK: [12, 13, 14, 15]
+ // CHECK: [16, 17, 18, 19]
+ // CHECK: [20, 21, 22, 23]
+ // CHECK: [24, 25, 26, 27]
+ // CHECK: [28, 29, 30, 31]
+ // CHECK: [32, 33, 34, 35]
+ // CHECK: [36, 37, 38, 39]
+ // CHECK: [40, 41, 42, 43]
+ // CHECK: [44, 45, 46, 47]
+ // CHECK: [48, 49, 50, 51]
+ // CHECK: [52, 53, 54, 55]
+ // CHECK: [56, 57, 58, 59]
+ // CHECK: [60, 61, 62, 63]
+
+ memref.dealloc %A : memref<16x2xi32>
+ memref.dealloc %B : memref<16x2xi32>
+ memref.dealloc %C : memref<16x4xi16>
+ memref.dealloc %D : memref<16x4xi32>
+ return
+ }
+ func.func private @printMemrefI32(%ptr : memref<*xi32>) attributes { llvm.emit_c_interface }
+}
More information about the Mlir-commits
mailing list