[Mlir-commits] [mlir] [MLIR][XeVM] Add xevm.bitcast_shuffle op and lowering to LLVM (PR #215303)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 10 13:51:19 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Sang Ik Lee (silee2)
<details>
<summary>Changes</summary>
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.
---
Patch is 28.54 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/215303.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td (+66)
- (modified) mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp (+80-26)
- (modified) mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp (+16)
- (modified) mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir (+67)
- (modified) mlir/test/Dialect/LLVMIR/invalid.mlir (+25)
- (modified) mlir/test/Dialect/LLVMIR/xevm.mlir (+32)
- (added) mlir/test/Integration/Dialect/XeVM/GPU/xevm_bitcast_shuffle.mlir (+219)
``````````diff
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..9757ed4cef658 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -99,6 +99,27 @@ 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`, `f16`, `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());
+ })
+ // Floats of the same width share a spelling, so both bfloat16 and half
+ // mangle as `f16`.
+ .Case([](FloatType ty) -> std::string {
+ return "f" + std::to_string(ty.getWidth());
+ })
+ .DefaultUnreachable("unhandled type for GenISA mangling");
+}
+
std::string builtinElemType(ElemType elemType) {
switch (elemType) {
case ElemType::BF8:
@@ -1520,6 +1541,39 @@ 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`.
+//
+// Note that the overload suffix spells every float of a given width the same
+// way, so a bf16 operand is passed through unchanged but named `f16`, as IGC
+// expects. The suffix therefore does not identify the overload on its own; the
+// declaration carries the actual types.
+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 +1841,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..2d61f38ea0163 100644
--- a/mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir
+++ b/mlir/test/Conversion/XeVMToLLVM/xevm-to-llvm.mlir
@@ -569,3 +569,70 @@ 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
+}
+
+// -----
+// bf16 is passed through unchanged, but the intrinsic name spells it `f16`, as
+// the overload suffix does not distinguish floats of the same width.
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i32.v4f16(vector<4xbf16>) -> vector<2xi32>
+llvm.func @bitcast_shuffle_bf16_src(%a: vector<4xbf16>) -> vector<2xi32> {
+ // CHECK-NOT: llvm.bitcast
+ // CHECK: llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i32.v4f16({{.*}}) {{{.*}}} : (vector<4xbf16>) -> 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.v4f16.v2i32(vector<2xi32>) -> vector<4xbf16>
+llvm.func @bitcast_shuffle_bf16_res(%a: vector<2xi32>) -> vector<4xbf16> {
+ // CHECK: %[[RES:.*]] = llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v4f16.v2i32({{.*}}) {{{.*}}} : (vector<2xi32>) -> vector<4xbf16>
+ // CHECK: llvm.return %[[RES]]
+ %0 = xevm.bitcast_shuffle %a : (vector<2xi32>) -> vector<4xbf16>
+ llvm.return %0 : vector<4xbf16>
+}
+
+// -----
+// A shuffle between types of the same component count still goes through the
+// intrinsic, which sees distinct operand and result types.
+// CHECK-LABEL: llvm.func spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i16.v2f16(vector<2xbf16>) -> vector<2xi16>
+llvm.func @bitcast_shuffle_bf16_same_width(%a: vector<2xbf16>) -> vector<2xi16> {
+ // CHECK: %[[RES:.*]] = llvm.call spir_funccc @llvm.genx.GenISA.SubgroupBitcastShuffle.v2i16.v2f16({{.*}}) {{{.*}}} : (vector<2xbf16>) -> 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 = xev...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/215303
More information about the Mlir-commits
mailing list