[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 13:47:29 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/215303
>From 8ad9a80e157f13048970e13cd034c86cc5785441 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 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. The operation is
reversible.
Exactly one of `src` and `res` is a 1D vector and the other is a scalar, so the
op comes in two forms: a pack, taking a vector and returning a scalar, and an
unpack, which is its exact inverse. This matches how `xegpu.lane_shuffle` is
lowered onto it. A vector-to-vector repack is not supported, and neither is a
plain scalar-to-scalar bitcast.
`src` and `res` must have the same total bit width. Only the integer types `i8`,
`i16`, `i32` and `i64` are accepted, on either side. The op 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 beforehand and to
bitcast the result back. That keeps the set of overloads the op 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 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 end-to-end test checks reversibility, by packing and unpacking back to the
original type, and the pack pattern itself, by tagging every 16-bit unit of the
source with its position in the concatenated stream and reading the tags off the
result.
---
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