[Mlir-commits] [mlir] [MLIR][XeVM] XeVM to LLVM: Update xevm.truncf handling (PR #194491)
Sang Ik Lee
llvmlistbot at llvm.org
Fri May 1 08:42:12 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/194491
>From 0be015b04c8c4d005ea38576bf6d660a08c37c8e Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 22 Apr 2026 21:21:03 +0000
Subject: [PATCH 01/10] Update comments.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 12 ++++++++++++
1 file changed, 12 insertions(+)
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index 7b3d4d3bab18b..4ace97c4cada6 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1129,6 +1129,18 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
// Supported source and result types are resticted for now.
auto srcEtype = op.getSrcEtype().getEtype();
auto dstEtype = op.getDstEtype().getEtype();
+ // TODO: support power of 2 number of elements
+ // batch_size =
+ // 16 if dst type == fp8
+ // 8 if dst type == fp4
+ // For num_elem > batch_size
+ // convert batch of batch_size
+ // cast batch to i32 elem type vector
+ // concat batches by shufflevector
+ // For num_elem = batch_size
+ // use API for conversion
+ // For num_elem < batch_size
+ // not supported for now
if (auto vecSrcTy = dyn_cast<VectorType>(op.getSrc().getType())) {
if (vecSrcTy.getNumElements() != 16)
return rewriter.notifyMatchFailure(
>From 3cae7cae58701a63b6b63f5d1e5dfece4c078aee Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 27 Apr 2026 10:38:06 -0700
Subject: [PATCH 02/10] Update code comment. Remove output vector size
retsriction.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 20 ++++++++++---------
1 file changed, 11 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index 4ace97c4cada6..93310a614afe8 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1129,7 +1129,14 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
// Supported source and result types are resticted for now.
auto srcEtype = op.getSrcEtype().getEtype();
auto dstEtype = op.getDstEtype().getEtype();
- // TODO: support power of 2 number of elements
+ // Currently only 16 input elements are supported as
+ // - Any vector beyond 16 elements not a valid OpenCL vector.
+ // - 2D block load can only load up to 16 16bit elements per lane.
+ // Widest load is 8x16xi32 with 16 lanes, which is 16 16bit
+ // elements per lane.
+ // - mma_mx A and B operands need more than 16 elements per lane
+ //
+ // Conversion is done in batches depending on the dst type.
// batch_size =
// 16 if dst type == fp8
// 8 if dst type == fp4
@@ -1139,8 +1146,7 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
// concat batches by shufflevector
// For num_elem = batch_size
// use API for conversion
- // For num_elem < batch_size
- // not supported for now
+ // Scalar case is not supported until usage case become clear.
if (auto vecSrcTy = dyn_cast<VectorType>(op.getSrc().getType())) {
if (vecSrcTy.getNumElements() != 16)
return rewriter.notifyMatchFailure(
@@ -1148,13 +1154,9 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
} else {
return rewriter.notifyMatchFailure(op, "Scalar src is not supported.");
}
- if (auto vecDstTy = dyn_cast<VectorType>(op.getDst().getType())) {
- if (vecDstTy.getNumElements() != 16)
- return rewriter.notifyMatchFailure(
- op, "Only vector dst of 16 elements is supported");
- } else {
+ auto vecDstTy = dyn_cast<VectorType>(op.getDst().getType());
+ if (!vecDstTy)
return rewriter.notifyMatchFailure(op, "Scalar dst is not supported.");
- }
if (srcEtype == TruncfSrcElemTypes::F16 &&
dstEtype == TruncfDstElemTypes::BF8) {
// BF8 is just F16 with lower 8 bits of mantessa discard.
>From 49ca371665241e9f99da8bfc2c7eebdc0d8594a9 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 27 Apr 2026 11:32:23 -0700
Subject: [PATCH 03/10] Add support for f16 to f8 truncation.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 81 +++++----
.../XeVM/GPU/xevm_block_scaled_dpas_f8.mlir | 158 ++++++++++++++++++
2 files changed, 208 insertions(+), 31 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_f8.mlir
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index 93310a614afe8..090f5f522e04c 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1157,37 +1157,56 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
auto vecDstTy = dyn_cast<VectorType>(op.getDst().getType());
if (!vecDstTy)
return rewriter.notifyMatchFailure(op, "Scalar dst is not supported.");
- if (srcEtype == TruncfSrcElemTypes::F16 &&
- dstEtype == TruncfDstElemTypes::BF8) {
- // BF8 is just F16 with lower 8 bits of mantessa discard.
- // Signbit Exponent Mantessa
- // BF8 1 5 2
- // F16 1 5 10
- // Xe arch is Little Endian so BF8 is just the second byte of the two
- // byte representation used for F16
- auto firstHalf =
- LLVM::ShuffleVectorOp::create(rewriter, op.getLoc(), op.getSrc(),
- op.getSrc(), {0, 1, 2, 3, 4, 5, 6, 7});
- auto secondHalf = LLVM::ShuffleVectorOp::create(
- rewriter, op.getLoc(), op.getSrc(), op.getSrc(),
- {8, 9, 10, 11, 12, 13, 14, 15});
- auto firstHalfCasted = LLVM::BitcastOp::create(
- rewriter, op.getLoc(), VectorType::get(16, rewriter.getI8Type()),
- firstHalf);
- auto secondHalfCasted = LLVM::BitcastOp::create(
- rewriter, op.getLoc(), VectorType::get(16, rewriter.getI8Type()),
- secondHalf);
- // Gather just the second bytes from every two byte F16 values
- auto resFirstHalf = LLVM::ShuffleVectorOp::create(
- rewriter, op.getLoc(), firstHalfCasted, firstHalfCasted,
- {1, 3, 5, 7, 9, 11, 13, 15});
- auto resSecondHalf = LLVM::ShuffleVectorOp::create(
- rewriter, op.getLoc(), secondHalfCasted, secondHalfCasted,
- {1, 3, 5, 7, 9, 11, 13, 15});
- auto res = LLVM::ShuffleVectorOp::create(
- rewriter, op.getLoc(), resFirstHalf, resSecondHalf,
- {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15});
- rewriter.replaceOp(op, res);
+ // BF16 type needs some preprocessing before conversion,
+ // First extended to F32 and then truncated to F16.
+ if (srcEtype == TruncfSrcElemTypes::BF16) {
+ // Step 1: Extend to F32
+ // Use float16 __builtin_IB_bftof_16(short16)
+ // Step 2: Truncf to F16
+ // Use half16 convert_half16(float16)
+ }
+ if (dstEtype == TruncfDstElemTypes::BF8) {
+ // Use char16 __builtin_IB_hftobf8_16(half16)
+ std::string fnName = "__builtin_IB_hftobf8_16";
+ SmallVector<Type> argTypes{op.getSrc().getType()};
+ SmallVector<Value> args{op.getSrc()};
+
+ auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
+ /*other=*/LLVM::ModRefInfo::NoModRef,
+ /*argMem=*/LLVM::ModRefInfo::NoModRef,
+ /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef,
+ /*errnoMem=*/LLVM::ModRefInfo::NoModRef,
+ /*targetMem0=*/LLVM::ModRefInfo::NoModRef,
+ /*targetMem1=*/LLVM::ModRefInfo::NoModRef);
+ auto funcAttrs = convergentNoUnwindWillReturnAttrs;
+ funcAttrs.memEffectsAttr = memAttr;
+ Value result =
+ createDeviceFunctionCall(rewriter, fnName, vecDstTy, argTypes, args,
+ {}, funcAttrs, op.getOperation())
+ ->getResult(0);
+
+ rewriter.replaceOp(op, result);
+ } else if (dstEtype == TruncfDstElemTypes::F8) {
+ // Use char16 __builtin_IB_hftohf8_16(half16)
+ std::string fnName = "__builtin_IB_hftohf8_16";
+ SmallVector<Type> argTypes{op.getSrc().getType()};
+ SmallVector<Value> args{op.getSrc()};
+
+ auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
+ /*other=*/LLVM::ModRefInfo::NoModRef,
+ /*argMem=*/LLVM::ModRefInfo::NoModRef,
+ /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef,
+ /*errnoMem=*/LLVM::ModRefInfo::NoModRef,
+ /*targetMem0=*/LLVM::ModRefInfo::NoModRef,
+ /*targetMem1=*/LLVM::ModRefInfo::NoModRef);
+ auto funcAttrs = convergentNoUnwindWillReturnAttrs;
+ funcAttrs.memEffectsAttr = memAttr;
+ Value result =
+ createDeviceFunctionCall(rewriter, fnName, vecDstTy, argTypes, args,
+ {}, funcAttrs, op.getOperation())
+ ->getResult(0);
+
+ rewriter.replaceOp(op, result);
} else {
return rewriter.notifyMatchFailure(
op, "Unsupported src, dst element type pair.");
diff --git a/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_f8.mlir b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_f8.mlir
new file mode 100644
index 0000000000000..f72992151c25c
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_f8.mlir
@@ -0,0 +1,158 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=lane zebin-chip=cri" \
+// 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
+
+// XFAIL: *
+module @gemm attributes {gpu.container_module} {
+ gpu.module @kernel {
+ gpu.func @block_scaled_dpas_f8(%a: !llvm.ptr<1>, %b: !llvm.ptr<1>, %c: !llvm.ptr<1>) kernel {
+ // TODO: some values are related can be derived from others like the following.
+ // %M = arith.constant 8 : i32
+ // %N = arith.constant 16 : i32
+ // %K = arith.constant 8 : i32
+ // %load_a_elem_bitwidth = arith.constant 32 : i32
+ // %a_elem_bitwidth = arith.constant 16 : i32
+ // %mx_elem_bitwidth = arith.constant 8 : i32
+ // %load_a_pack_ratio = arith.divsi %load_a_elem_bitwidth, %a_elem_bitwidth : i32
+ // %mx_pack_ratio = arith.divsi %load_a_elem_bitwidth, %mx_elem_bitwidth : i32
+ // %load_a_K = arith.muli %K, %load_a_pack_ratio : i32
+ // %load_b_K = arith.muli %K, %mx_pack_ratio : i32
+
+ %base_width_a = arith.constant 64 : i32
+ %base_height_a = arith.constant 8 : i32
+ %base_pitch_a = arith.constant 64 : i32
+ %x = arith.constant 0 : i32
+ %y = arith.constant 0 : i32
+ // A is loaded as fp16, but it will be truncated to f8 before MMA.
+ // The blockload2d op need to be configured to load with double the width
+ // in number of elements or double the element bitwidth.
+ // block load does not support width of 32 elements of 16 bit,
+ // but it supports width of 16 elements of 32 bit.
+ // So the configuration is set to load 8 elements of 32 bits per lane and then
+ // bitcast to 16 elements of fp16 element type.
+ %loaded_a = xevm.blockload2d %a, %base_width_a, %base_height_a, %base_pitch_a, %x, %y
+ <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32, v_blocks=1 : i32,
+ transpose=false, pack_register=false}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
+ %loaded_a_casted = vector.bitcast %loaded_a : vector<8xi32> to vector<16xf16>
+ %a_trunc = xevm.truncf %loaded_a_casted { src_etype = f16, dst_etype = f8 } : (vector<16xf16>) -> vector<16xi8>
+ %a_trunc_casted = vector.bitcast %a_trunc : vector<16xi8> to vector<8xi16>
+
+ %base_width_b = arith.constant 16 : i32
+ %base_height_b = arith.constant 32 : i32
+ %base_pitch_b = arith.constant 16 : i32
+ // B is already in f8, and it will be used as is for MMA.
+ // So the blockload2d op is configured to load normally with 8bit element bitwidth
+ // with pack_register request.
+ %loaded_b = xevm.blockload2d %b, %base_width_b, %base_height_b, %base_pitch_b, %x, %y
+ <{elem_size_in_bits=8 : i32, tile_width=16 : i32, tile_height=32 : i32, v_blocks=1 : i32,
+ transpose=false, pack_register=true}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
+
+ // Note: scale is not computed. Constant values are used for simplifying the example
+ %scale_a = arith.constant 1.0 : f8E8M0FNU
+ %scale_b = arith.constant 1.0 : f8E8M0FNU
+ %scale_a_casted = arith.bitcast %scale_a : f8E8M0FNU to i8
+ %scale_b_casted = arith.bitcast %scale_b : f8E8M0FNU to i8
+ // Note: c is not loaded. constant vector is used for simplifying the example
+ %loaded_c_casted = arith.constant dense<0.0> : vector<8xf32>
+
+ %c_result = xevm.mma_mx %a_trunc_casted, %loaded_b, %scale_a_casted, %scale_b_casted, %loaded_c_casted
+ {shape=<m=8, n=16, k=32>, types=<d=f32, a=f8, b=f8, c=f32>}
+ : (vector<8xi16>, vector<8xi32>, i8, i8, vector<8xf32>) -> vector<8xf32>
+ %c_result_casted = vector.bitcast %c_result : vector<8xf32> to vector<8xi32>
+
+ %base_width_c = arith.constant 64 : i32
+ %base_height_c = arith.constant 8 : i32
+ %base_pitch_c = arith.constant 64 : i32
+ xevm.blockstore2d %c, %base_width_c, %base_height_c, %base_pitch_c, %x, %y, %c_result_casted
+ <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32}>
+ : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)
+ gpu.return
+ }
+ }
+
+ func.func @test(%a : memref<8x32xf16>, %b : memref<32x16xf8E4M3FN>, %c : memref<8x16xf32>) -> memref<8x16xf32> attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+
+ %memref_a = gpu.alloc() : memref<8x32xf16>
+ gpu.memcpy %memref_a, %a : memref<8x32xf16>, memref<8x32xf16>
+ %a_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_a : memref<8x32xf16> -> index
+ %a_ptr_as_i64 = arith.index_cast %a_ptr_as_idx : index to i64
+ %a_ptr = llvm.inttoptr %a_ptr_as_i64 : i64 to !llvm.ptr
+ %a_ptr_casted = llvm.addrspacecast %a_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ %memref_b = gpu.alloc() : memref<32x16xf8E4M3FN>
+ gpu.memcpy %memref_b, %b : memref<32x16xf8E4M3FN>, memref<32x16xf8E4M3FN>
+ %b_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_b : memref<32x16xf8E4M3FN> -> index
+ %b_ptr_as_i64 = arith.index_cast %b_ptr_as_idx : index to i64
+ %b_ptr = llvm.inttoptr %b_ptr_as_i64 : i64 to !llvm.ptr
+ %b_ptr_casted = llvm.addrspacecast %b_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ %memref_c = gpu.alloc() : memref<8x16xf32>
+ gpu.memcpy %memref_c, %c : memref<8x16xf32>, memref<8x16xf32>
+ %c_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_c : memref<8x16xf32> -> index
+ %c_ptr_as_i64 = arith.index_cast %c_ptr_as_idx : index to i64
+ %c_ptr = llvm.inttoptr %c_ptr_as_i64 : i64 to !llvm.ptr
+ %c_ptr_casted = llvm.addrspacecast %c_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ gpu.launch_func @kernel::@block_scaled_dpas_f8 blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)
+ args(%a_ptr_casted : !llvm.ptr<1>, %b_ptr_casted : !llvm.ptr<1>, %c_ptr_casted : !llvm.ptr<1>)
+ gpu.dealloc %memref_a : memref<8x32xf16>
+ gpu.dealloc %memref_b : memref<32x16xf8E4M3FN>
+ %res = memref.alloc() : memref<8x16xf32>
+ gpu.memcpy %res, %memref_c : memref<8x16xf32>, memref<8x16xf32>
+ gpu.dealloc %memref_c : memref<8x16xf32>
+ return %res : memref<8x16xf32>
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c1f16 = arith.constant 1.0 : f16
+ %c1f8 = arith.constant 1.0 : f8E4M3FN
+ %c0f32 = arith.constant 0.0 : f32
+
+ %A = memref.alloc() : memref<8x32xf16>
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.for %j = %c0 to %c32 step %c1 {
+ memref.store %c1f16, %A[%i, %j] : memref<8x32xf16>
+ }
+ }
+
+ %B = memref.alloc() : memref<32x16xf8E4M3FN>
+ scf.for %i = %c0 to %c32 step %c1 {
+ scf.for %j = %c0 to %c16 step %c1 {
+ memref.store %c1f8, %B[%i, %j] : memref<32x16xf8E4M3FN>
+ }
+ }
+
+ %C = memref.alloc() : memref<8x16xf32>
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.for %j = %c0 to %c16 step %c1 {
+ memref.store %c0f32, %C[%i, %j] : memref<8x16xf32>
+ }
+ }
+
+ %C_res = call @test(%A, %B, %C) : (memref<8x32xf16>, memref<32x16xf8E4M3FN>, memref<8x16xf32>) -> memref<8x16xf32>
+ %C_cast = memref.cast %C_res : memref<8x16xf32> to memref<*xf32>
+ call @printMemrefF32(%C_cast) : (memref<*xf32>) -> ()
+
+ // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
+ // CHECK-COUNT-8: [32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32]
+ memref.dealloc %A : memref<8x32xf16>
+ memref.dealloc %B : memref<32x16xf8E4M3FN>
+ memref.dealloc %C : memref<8x16xf32>
+ memref.dealloc %C_res : memref<8x16xf32>
+ return
+ }
+ func.func private @printMemrefF32(%ptr : memref<*xf32>) attributes { llvm.emit_c_interface }
+
+}
>From 67a5b3eb2495f4c5443178b025be8a1819b2d25e Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 27 Apr 2026 12:04:20 -0700
Subject: [PATCH 04/10] Add intermediate bf16 to f16 conversion.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 67 ++++----
mlir/lib/Target/LLVM/XeVM/Target.cpp | 1 +
.../xevm_block_scaled_dpas_bf16_to_bf8.mlir | 158 ++++++++++++++++++
.../xevm_block_scaled_dpas_bf16_to_f8.mlir | 158 ++++++++++++++++++
4 files changed, 355 insertions(+), 29 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_bf16_to_bf8.mlir
create mode 100644 mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_bf16_to_f8.mlir
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index 090f5f522e04c..b52e42e0cfac6 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1147,39 +1147,58 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
// For num_elem = batch_size
// use API for conversion
// Scalar case is not supported until usage case become clear.
- if (auto vecSrcTy = dyn_cast<VectorType>(op.getSrc().getType())) {
- if (vecSrcTy.getNumElements() != 16)
- return rewriter.notifyMatchFailure(
- op, "Only vector src of 16 elements is supported");
- } else {
+ auto vecSrcTy = dyn_cast<VectorType>(op.getSrc().getType());
+ if (!vecSrcTy) {
return rewriter.notifyMatchFailure(op, "Scalar src is not supported.");
}
+ if (vecSrcTy.getNumElements() != 16)
+ return rewriter.notifyMatchFailure(
+ op, "Only vector src of 16 elements is supported");
auto vecDstTy = dyn_cast<VectorType>(op.getDst().getType());
if (!vecDstTy)
return rewriter.notifyMatchFailure(op, "Scalar dst is not supported.");
+ Value src = op.getSrc();
+ auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
+ /*other=*/LLVM::ModRefInfo::NoModRef,
+ /*argMem=*/LLVM::ModRefInfo::NoModRef,
+ /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef,
+ /*errnoMem=*/LLVM::ModRefInfo::NoModRef,
+ /*targetMem0=*/LLVM::ModRefInfo::NoModRef,
+ /*targetMem1=*/LLVM::ModRefInfo::NoModRef);
+ auto funcAttrs = convergentNoUnwindWillReturnAttrs;
+ funcAttrs.memEffectsAttr = memAttr;
// BF16 type needs some preprocessing before conversion,
// First extended to F32 and then truncated to F16.
if (srcEtype == TruncfSrcElemTypes::BF16) {
// Step 1: Extend to F32
// Use float16 __builtin_IB_bftof_16(short16)
+ src = LLVM::BitcastOp::create(
+ rewriter, op.getLoc(),
+ VectorType::get(vecSrcTy.getShape(), rewriter.getI16Type()), src);
+ std::string fnName = "__builtin_IB_bftof_16";
+ SmallVector<Type> argTypes{src.getType()};
+ SmallVector<Value> args{src};
+ Type resTy = VectorType::get(vecSrcTy.getShape(), rewriter.getF32Type());
+ src = createDeviceFunctionCall(rewriter, fnName, resTy, argTypes, args,
+ {}, funcAttrs, op.getOperation())
+ ->getResult(0);
// Step 2: Truncf to F16
// Use half16 convert_half16(float16)
+ std::string truncFnName = "convert_half16";
+ SmallVector<Type> truncArgTypes{src.getType()};
+ SmallVector<Value> truncArgs{src};
+ truncFnName = mangle(truncFnName, truncArgTypes);
+ resTy = VectorType::get(vecSrcTy.getShape(), rewriter.getF16Type());
+ src =
+ createDeviceFunctionCall(rewriter, truncFnName, resTy, truncArgTypes,
+ truncArgs, {}, funcAttrs, op.getOperation())
+ ->getResult(0);
}
if (dstEtype == TruncfDstElemTypes::BF8) {
// Use char16 __builtin_IB_hftobf8_16(half16)
std::string fnName = "__builtin_IB_hftobf8_16";
- SmallVector<Type> argTypes{op.getSrc().getType()};
- SmallVector<Value> args{op.getSrc()};
-
- auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
- /*other=*/LLVM::ModRefInfo::NoModRef,
- /*argMem=*/LLVM::ModRefInfo::NoModRef,
- /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef,
- /*errnoMem=*/LLVM::ModRefInfo::NoModRef,
- /*targetMem0=*/LLVM::ModRefInfo::NoModRef,
- /*targetMem1=*/LLVM::ModRefInfo::NoModRef);
- auto funcAttrs = convergentNoUnwindWillReturnAttrs;
- funcAttrs.memEffectsAttr = memAttr;
+ SmallVector<Type> argTypes{src.getType()};
+ SmallVector<Value> args{src};
Value result =
createDeviceFunctionCall(rewriter, fnName, vecDstTy, argTypes, args,
{}, funcAttrs, op.getOperation())
@@ -1189,18 +1208,8 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
} else if (dstEtype == TruncfDstElemTypes::F8) {
// Use char16 __builtin_IB_hftohf8_16(half16)
std::string fnName = "__builtin_IB_hftohf8_16";
- SmallVector<Type> argTypes{op.getSrc().getType()};
- SmallVector<Value> args{op.getSrc()};
-
- auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
- /*other=*/LLVM::ModRefInfo::NoModRef,
- /*argMem=*/LLVM::ModRefInfo::NoModRef,
- /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef,
- /*errnoMem=*/LLVM::ModRefInfo::NoModRef,
- /*targetMem0=*/LLVM::ModRefInfo::NoModRef,
- /*targetMem1=*/LLVM::ModRefInfo::NoModRef);
- auto funcAttrs = convergentNoUnwindWillReturnAttrs;
- funcAttrs.memEffectsAttr = memAttr;
+ SmallVector<Type> argTypes{src.getType()};
+ SmallVector<Value> args{src};
Value result =
createDeviceFunctionCall(rewriter, fnName, vecDstTy, argTypes, args,
{}, funcAttrs, op.getOperation())
diff --git a/mlir/lib/Target/LLVM/XeVM/Target.cpp b/mlir/lib/Target/LLVM/XeVM/Target.cpp
index 28c654f3d7c9a..a160e81321511 100644
--- a/mlir/lib/Target/LLVM/XeVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/XeVM/Target.cpp
@@ -462,6 +462,7 @@ void SPIRVSerializer::init() {
#if LLVM_HAS_SPIRV_TARGET
static const std::vector<std::string> getDefaultSPIRVExtensions() {
return {
+ "SPV_KHR_bfloat16",
"SPV_EXT_relaxed_printf_string_address_space",
"SPV_INTEL_cache_controls",
"SPV_INTEL_variable_length_array",
diff --git a/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_bf16_to_bf8.mlir b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_bf16_to_bf8.mlir
new file mode 100644
index 0000000000000..49b99e3418258
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_bf16_to_bf8.mlir
@@ -0,0 +1,158 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=lane zebin-chip=cri" \
+// 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
+
+// XFAIL: *
+module @gemm attributes {gpu.container_module} {
+ gpu.module @kernel {
+ gpu.func @block_scaled_dpas_bf8(%a: !llvm.ptr<1>, %b: !llvm.ptr<1>, %c: !llvm.ptr<1>) kernel {
+ // TODO: some values are related can be derived from others like the following.
+ // %M = arith.constant 8 : i32
+ // %N = arith.constant 16 : i32
+ // %K = arith.constant 8 : i32
+ // %load_a_elem_bitwidth = arith.constant 32 : i32
+ // %a_elem_bitwidth = arith.constant 16 : i32
+ // %mx_elem_bitwidth = arith.constant 8 : i32
+ // %load_a_pack_ratio = arith.divsi %load_a_elem_bitwidth, %a_elem_bitwidth : i32
+ // %mx_pack_ratio = arith.divsi %load_a_elem_bitwidth, %mx_elem_bitwidth : i32
+ // %load_a_K = arith.muli %K, %load_a_pack_ratio : i32
+ // %load_b_K = arith.muli %K, %mx_pack_ratio : i32
+
+ %base_width_a = arith.constant 64 : i32
+ %base_height_a = arith.constant 8 : i32
+ %base_pitch_a = arith.constant 64 : i32
+ %x = arith.constant 0 : i32
+ %y = arith.constant 0 : i32
+ // A is loaded as fp16, but it will be truncated to bf8 before MMA.
+ // The blockload2d op need to be configured to load with double the width
+ // in number of elements or double the element bitwidth.
+ // block load does not support width of 32 elements of 16 bit,
+ // but it supports width of 16 elements of 32 bit.
+ // So the configuration is set to load 8 elements of 32 bits per lane and then
+ // bitcast to 16 elements of fp16 element type.
+ %loaded_a = xevm.blockload2d %a, %base_width_a, %base_height_a, %base_pitch_a, %x, %y
+ <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32, v_blocks=1 : i32,
+ transpose=false, pack_register=false}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
+ %loaded_a_casted = vector.bitcast %loaded_a : vector<8xi32> to vector<16xbf16>
+ %a_trunc = xevm.truncf %loaded_a_casted { src_etype = bf16, dst_etype = bf8 } : (vector<16xbf16>) -> vector<16xi8>
+ %a_trunc_casted = vector.bitcast %a_trunc : vector<16xi8> to vector<8xi16>
+
+ %base_width_b = arith.constant 16 : i32
+ %base_height_b = arith.constant 32 : i32
+ %base_pitch_b = arith.constant 16 : i32
+ // B is already in bf8, and it will be used as is for MMA.
+ // So the blockload2d op is configured to load normally with 8bit element bitwidth
+ // with pack_register request.
+ %loaded_b = xevm.blockload2d %b, %base_width_b, %base_height_b, %base_pitch_b, %x, %y
+ <{elem_size_in_bits=8 : i32, tile_width=16 : i32, tile_height=32 : i32, v_blocks=1 : i32,
+ transpose=false, pack_register=true}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
+
+ // Note: scale is not computed. Constant values are used for simplifying the example
+ %scale_a = arith.constant 1.0 : f8E8M0FNU
+ %scale_b = arith.constant 1.0 : f8E8M0FNU
+ %scale_a_casted = arith.bitcast %scale_a : f8E8M0FNU to i8
+ %scale_b_casted = arith.bitcast %scale_b : f8E8M0FNU to i8
+ // Note: c is not loaded. constant vector is used for simplifying the example
+ %loaded_c_casted = arith.constant dense<0.0> : vector<8xf32>
+
+ %c_result = xevm.mma_mx %a_trunc_casted, %loaded_b, %scale_a_casted, %scale_b_casted, %loaded_c_casted
+ {shape=<m=8, n=16, k=32>, types=<d=f32, a=bf8, b=bf8, c=f32>}
+ : (vector<8xi16>, vector<8xi32>, i8, i8, vector<8xf32>) -> vector<8xf32>
+ %c_result_casted = vector.bitcast %c_result : vector<8xf32> to vector<8xi32>
+
+ %base_width_c = arith.constant 64 : i32
+ %base_height_c = arith.constant 8 : i32
+ %base_pitch_c = arith.constant 64 : i32
+ xevm.blockstore2d %c, %base_width_c, %base_height_c, %base_pitch_c, %x, %y, %c_result_casted
+ <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32}>
+ : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)
+ gpu.return
+ }
+ }
+
+ func.func @test(%a : memref<8x32xbf16>, %b : memref<32x16xf8E5M2>, %c : memref<8x16xf32>) -> memref<8x16xf32> attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+
+ %memref_a = gpu.alloc() : memref<8x32xbf16>
+ gpu.memcpy %memref_a, %a : memref<8x32xbf16>, memref<8x32xbf16>
+ %a_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_a : memref<8x32xbf16> -> index
+ %a_ptr_as_i64 = arith.index_cast %a_ptr_as_idx : index to i64
+ %a_ptr = llvm.inttoptr %a_ptr_as_i64 : i64 to !llvm.ptr
+ %a_ptr_casted = llvm.addrspacecast %a_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ %memref_b = gpu.alloc() : memref<32x16xf8E5M2>
+ gpu.memcpy %memref_b, %b : memref<32x16xf8E5M2>, memref<32x16xf8E5M2>
+ %b_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_b : memref<32x16xf8E5M2> -> index
+ %b_ptr_as_i64 = arith.index_cast %b_ptr_as_idx : index to i64
+ %b_ptr = llvm.inttoptr %b_ptr_as_i64 : i64 to !llvm.ptr
+ %b_ptr_casted = llvm.addrspacecast %b_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ %memref_c = gpu.alloc() : memref<8x16xf32>
+ gpu.memcpy %memref_c, %c : memref<8x16xf32>, memref<8x16xf32>
+ %c_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_c : memref<8x16xf32> -> index
+ %c_ptr_as_i64 = arith.index_cast %c_ptr_as_idx : index to i64
+ %c_ptr = llvm.inttoptr %c_ptr_as_i64 : i64 to !llvm.ptr
+ %c_ptr_casted = llvm.addrspacecast %c_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ gpu.launch_func @kernel::@block_scaled_dpas_bf8 blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)
+ args(%a_ptr_casted : !llvm.ptr<1>, %b_ptr_casted : !llvm.ptr<1>, %c_ptr_casted : !llvm.ptr<1>)
+ gpu.dealloc %memref_a : memref<8x32xbf16>
+ gpu.dealloc %memref_b : memref<32x16xf8E5M2>
+ %res = memref.alloc() : memref<8x16xf32>
+ gpu.memcpy %res, %memref_c : memref<8x16xf32>, memref<8x16xf32>
+ gpu.dealloc %memref_c : memref<8x16xf32>
+ return %res : memref<8x16xf32>
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c1bf16 = arith.constant 1.0 : bf16
+ %c1bf8 = arith.constant 1.0 : f8E5M2
+ %c0f32 = arith.constant 0.0 : f32
+
+ %A = memref.alloc() : memref<8x32xbf16>
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.for %j = %c0 to %c32 step %c1 {
+ memref.store %c1bf16, %A[%i, %j] : memref<8x32xbf16>
+ }
+ }
+
+ %B = memref.alloc() : memref<32x16xf8E5M2>
+ scf.for %i = %c0 to %c32 step %c1 {
+ scf.for %j = %c0 to %c16 step %c1 {
+ memref.store %c1bf8, %B[%i, %j] : memref<32x16xf8E5M2>
+ }
+ }
+
+ %C = memref.alloc() : memref<8x16xf32>
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.for %j = %c0 to %c16 step %c1 {
+ memref.store %c0f32, %C[%i, %j] : memref<8x16xf32>
+ }
+ }
+
+ %C_res = call @test(%A, %B, %C) : (memref<8x32xbf16>, memref<32x16xf8E5M2>, memref<8x16xf32>) -> memref<8x16xf32>
+ %C_cast = memref.cast %C_res : memref<8x16xf32> to memref<*xf32>
+ call @printMemrefF32(%C_cast) : (memref<*xf32>) -> ()
+
+ // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
+ // CHECK-COUNT-8: [32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32]
+ memref.dealloc %A : memref<8x32xbf16>
+ memref.dealloc %B : memref<32x16xf8E5M2>
+ memref.dealloc %C : memref<8x16xf32>
+ memref.dealloc %C_res : memref<8x16xf32>
+ return
+ }
+ func.func private @printMemrefF32(%ptr : memref<*xf32>) attributes { llvm.emit_c_interface }
+
+}
diff --git a/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_bf16_to_f8.mlir b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_bf16_to_f8.mlir
new file mode 100644
index 0000000000000..03778ef525399
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_bf16_to_f8.mlir
@@ -0,0 +1,158 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=lane zebin-chip=cri" \
+// 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
+
+// XFAIL: *
+module @gemm attributes {gpu.container_module} {
+ gpu.module @kernel {
+ gpu.func @block_scaled_dpas_f8(%a: !llvm.ptr<1>, %b: !llvm.ptr<1>, %c: !llvm.ptr<1>) kernel {
+ // TODO: some values are related can be derived from others like the following.
+ // %M = arith.constant 8 : i32
+ // %N = arith.constant 16 : i32
+ // %K = arith.constant 8 : i32
+ // %load_a_elem_bitwidth = arith.constant 32 : i32
+ // %a_elem_bitwidth = arith.constant 16 : i32
+ // %mx_elem_bitwidth = arith.constant 8 : i32
+ // %load_a_pack_ratio = arith.divsi %load_a_elem_bitwidth, %a_elem_bitwidth : i32
+ // %mx_pack_ratio = arith.divsi %load_a_elem_bitwidth, %mx_elem_bitwidth : i32
+ // %load_a_K = arith.muli %K, %load_a_pack_ratio : i32
+ // %load_b_K = arith.muli %K, %mx_pack_ratio : i32
+
+ %base_width_a = arith.constant 64 : i32
+ %base_height_a = arith.constant 8 : i32
+ %base_pitch_a = arith.constant 64 : i32
+ %x = arith.constant 0 : i32
+ %y = arith.constant 0 : i32
+ // A is loaded as fp16, but it will be truncated to f8 before MMA.
+ // The blockload2d op need to be configured to load with double the width
+ // in number of elements or double the element bitwidth.
+ // block load does not support width of 32 elements of 16 bit,
+ // but it supports width of 16 elements of 32 bit.
+ // So the configuration is set to load 8 elements of 32 bits per lane and then
+ // bitcast to 16 elements of fp16 element type.
+ %loaded_a = xevm.blockload2d %a, %base_width_a, %base_height_a, %base_pitch_a, %x, %y
+ <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32, v_blocks=1 : i32,
+ transpose=false, pack_register=false}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
+ %loaded_a_casted = vector.bitcast %loaded_a : vector<8xi32> to vector<16xbf16>
+ %a_trunc = xevm.truncf %loaded_a_casted { src_etype = bf16, dst_etype = f8 } : (vector<16xbf16>) -> vector<16xi8>
+ %a_trunc_casted = vector.bitcast %a_trunc : vector<16xi8> to vector<8xi16>
+
+ %base_width_b = arith.constant 16 : i32
+ %base_height_b = arith.constant 32 : i32
+ %base_pitch_b = arith.constant 16 : i32
+ // B is already in f8, and it will be used as is for MMA.
+ // So the blockload2d op is configured to load normally with 8bit element bitwidth
+ // with pack_register request.
+ %loaded_b = xevm.blockload2d %b, %base_width_b, %base_height_b, %base_pitch_b, %x, %y
+ <{elem_size_in_bits=8 : i32, tile_width=16 : i32, tile_height=32 : i32, v_blocks=1 : i32,
+ transpose=false, pack_register=true}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
+
+ // Note: scale is not computed. Constant values are used for simplifying the example
+ %scale_a = arith.constant 1.0 : f8E8M0FNU
+ %scale_b = arith.constant 1.0 : f8E8M0FNU
+ %scale_a_casted = arith.bitcast %scale_a : f8E8M0FNU to i8
+ %scale_b_casted = arith.bitcast %scale_b : f8E8M0FNU to i8
+ // Note: c is not loaded. constant vector is used for simplifying the example
+ %loaded_c_casted = arith.constant dense<0.0> : vector<8xf32>
+
+ %c_result = xevm.mma_mx %a_trunc_casted, %loaded_b, %scale_a_casted, %scale_b_casted, %loaded_c_casted
+ {shape=<m=8, n=16, k=32>, types=<d=f32, a=f8, b=f8, c=f32>}
+ : (vector<8xi16>, vector<8xi32>, i8, i8, vector<8xf32>) -> vector<8xf32>
+ %c_result_casted = vector.bitcast %c_result : vector<8xf32> to vector<8xi32>
+
+ %base_width_c = arith.constant 64 : i32
+ %base_height_c = arith.constant 8 : i32
+ %base_pitch_c = arith.constant 64 : i32
+ xevm.blockstore2d %c, %base_width_c, %base_height_c, %base_pitch_c, %x, %y, %c_result_casted
+ <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32}>
+ : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)
+ gpu.return
+ }
+ }
+
+ func.func @test(%a : memref<8x32xbf16>, %b : memref<32x16xf8E4M3FN>, %c : memref<8x16xf32>) -> memref<8x16xf32> attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+
+ %memref_a = gpu.alloc() : memref<8x32xbf16>
+ gpu.memcpy %memref_a, %a : memref<8x32xbf16>, memref<8x32xbf16>
+ %a_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_a : memref<8x32xbf16> -> index
+ %a_ptr_as_i64 = arith.index_cast %a_ptr_as_idx : index to i64
+ %a_ptr = llvm.inttoptr %a_ptr_as_i64 : i64 to !llvm.ptr
+ %a_ptr_casted = llvm.addrspacecast %a_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ %memref_b = gpu.alloc() : memref<32x16xf8E4M3FN>
+ gpu.memcpy %memref_b, %b : memref<32x16xf8E4M3FN>, memref<32x16xf8E4M3FN>
+ %b_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_b : memref<32x16xf8E4M3FN> -> index
+ %b_ptr_as_i64 = arith.index_cast %b_ptr_as_idx : index to i64
+ %b_ptr = llvm.inttoptr %b_ptr_as_i64 : i64 to !llvm.ptr
+ %b_ptr_casted = llvm.addrspacecast %b_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ %memref_c = gpu.alloc() : memref<8x16xf32>
+ gpu.memcpy %memref_c, %c : memref<8x16xf32>, memref<8x16xf32>
+ %c_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_c : memref<8x16xf32> -> index
+ %c_ptr_as_i64 = arith.index_cast %c_ptr_as_idx : index to i64
+ %c_ptr = llvm.inttoptr %c_ptr_as_i64 : i64 to !llvm.ptr
+ %c_ptr_casted = llvm.addrspacecast %c_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ gpu.launch_func @kernel::@block_scaled_dpas_f8 blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)
+ args(%a_ptr_casted : !llvm.ptr<1>, %b_ptr_casted : !llvm.ptr<1>, %c_ptr_casted : !llvm.ptr<1>)
+ gpu.dealloc %memref_a : memref<8x32xbf16>
+ gpu.dealloc %memref_b : memref<32x16xf8E4M3FN>
+ %res = memref.alloc() : memref<8x16xf32>
+ gpu.memcpy %res, %memref_c : memref<8x16xf32>, memref<8x16xf32>
+ gpu.dealloc %memref_c : memref<8x16xf32>
+ return %res : memref<8x16xf32>
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c1bf16 = arith.constant 1.0 : bf16
+ %c1f8 = arith.constant 1.0 : f8E4M3FN
+ %c0f32 = arith.constant 0.0 : f32
+
+ %A = memref.alloc() : memref<8x32xbf16>
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.for %j = %c0 to %c32 step %c1 {
+ memref.store %c1bf16, %A[%i, %j] : memref<8x32xbf16>
+ }
+ }
+
+ %B = memref.alloc() : memref<32x16xf8E4M3FN>
+ scf.for %i = %c0 to %c32 step %c1 {
+ scf.for %j = %c0 to %c16 step %c1 {
+ memref.store %c1f8, %B[%i, %j] : memref<32x16xf8E4M3FN>
+ }
+ }
+
+ %C = memref.alloc() : memref<8x16xf32>
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.for %j = %c0 to %c16 step %c1 {
+ memref.store %c0f32, %C[%i, %j] : memref<8x16xf32>
+ }
+ }
+
+ %C_res = call @test(%A, %B, %C) : (memref<8x32xbf16>, memref<32x16xf8E4M3FN>, memref<8x16xf32>) -> memref<8x16xf32>
+ %C_cast = memref.cast %C_res : memref<8x16xf32> to memref<*xf32>
+ call @printMemrefF32(%C_cast) : (memref<*xf32>) -> ()
+
+ // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
+ // CHECK-COUNT-8: [32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32]
+ memref.dealloc %A : memref<8x32xbf16>
+ memref.dealloc %B : memref<32x16xf8E4M3FN>
+ memref.dealloc %C : memref<8x16xf32>
+ memref.dealloc %C_res : memref<8x16xf32>
+ return
+ }
+ func.func private @printMemrefF32(%ptr : memref<*xf32>) attributes { llvm.emit_c_interface }
+
+}
>From 3410ba89ad14c0ad9d00d1406fe13fc205beb400 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 27 Apr 2026 14:03:27 -0700
Subject: [PATCH 05/10] Add support for e2m1 dst type.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 64 +++++++++++++++++++
1 file changed, 64 insertions(+)
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index b52e42e0cfac6..4ff1f4c750f64 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1216,6 +1216,70 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
->getResult(0);
rewriter.replaceOp(op, result);
+ } else if (dstEtype == TruncfDstElemTypes::E2M1) {
+ // Convert 8 elements at a time.
+ // To convert 8 elements, vector<8xf16>:
+ // Use:
+ // uint __builtin_IB_dnscl_hf16(uint, uint, 1, 0)
+ // uint __builtin_IB_dnscl_hf16(uint, uint, 1, 3)
+ // llvm.or
+ Value cast = LLVM::BitcastOp::create(
+ rewriter, op.getLoc(), VectorType::get(8, rewriter.getI32Type()),
+ src);
+
+ auto genDnscl = [&](Value input, Value idx0, Value idx1, Value dstTy,
+ Value mode) -> Value {
+ std::string fnName = "__builtin_IB_dnscl_hf16";
+ Value arg1 =
+ LLVM::ExtractElementOp::create(rewriter, op.getLoc(), input, idx0)
+ ->getResult(0);
+ Value arg2 =
+ LLVM::ExtractElementOp::create(rewriter, op.getLoc(), input, idx1)
+ ->getResult(0);
+ SmallVector<Type> argTypes{arg1.getType(), arg2.getType()};
+ SmallVector<Value> args{arg1, arg2};
+ Value dnscl = createDeviceFunctionCall(
+ rewriter, fnName, rewriter.getI32Type(), argTypes,
+ args, {}, funcAttrs, op.getOperation())
+ ->getResult(0);
+ return dnscl;
+ };
+
+ Value zero = LLVM::ConstantOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), 0);
+ Value one = LLVM::ConstantOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), 1);
+ Value two = LLVM::ConstantOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), 2);
+ Value three = LLVM::ConstantOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), 3);
+ Value even = genDnscl(cast, zero, two, one, zero);
+ Value odd = genDnscl(cast, one, three, one, three);
+ Value firstHalf = LLVM::OrOp::create(rewriter, op.getLoc(), even, odd);
+ Value four = LLVM::ConstantOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), 4);
+ Value five = LLVM::ConstantOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), 5);
+ Value six = LLVM::ConstantOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), 6);
+ Value seven = LLVM::ConstantOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), 7);
+ even = genDnscl(cast, four, six, one, zero);
+ odd = genDnscl(cast, five, seven, one, three);
+ Value secondHalf = LLVM::OrOp::create(rewriter, op.getLoc(), even, odd);
+ // Create vector<2xi32> from two i32 values and then bitcast to
+ // vector<8xi8> to match the dst type.
+ Value combined = LLVM::UndefOp::create(
+ rewriter, op.getLoc(), VectorType::get(2, rewriter.getI32Type()));
+ combined = LLVM::InsertElementOp::create(rewriter, op.getLoc(), combined,
+ firstHalf, zero)
+ ->getResult(0);
+ combined = LLVM::InsertElementOp::create(rewriter, op.getLoc(), combined,
+ secondHalf, one)
+ ->getResult(0);
+ Value result =
+ LLVM::BitcastOp::create(rewriter, op.getLoc(), vecDstTy, combined);
+ rewriter.replaceOp(op, result);
} else {
return rewriter.notifyMatchFailure(
op, "Unsupported src, dst element type pair.");
>From e32d61bfbd29f2e8eab6e43cd0db4ec16b3fe80e Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 27 Apr 2026 15:44:51 -0700
Subject: [PATCH 06/10] Add test case and fix issues.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 5 +-
mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp | 3 -
mlir/test/Dialect/LLVMIR/invalid.mlir | 8 -
.../XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir | 147 ++++++++++++++++++
4 files changed, 150 insertions(+), 13 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index 4ff1f4c750f64..3e3d67256c5f0 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1236,8 +1236,9 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
Value arg2 =
LLVM::ExtractElementOp::create(rewriter, op.getLoc(), input, idx1)
->getResult(0);
- SmallVector<Type> argTypes{arg1.getType(), arg2.getType()};
- SmallVector<Value> args{arg1, arg2};
+ SmallVector<Type> argTypes{arg1.getType(), arg2.getType(),
+ dstTy.getType(), mode.getType()};
+ SmallVector<Value> args{arg1, arg2, dstTy, mode};
Value dnscl = createDeviceFunctionCall(
rewriter, fnName, rewriter.getI32Type(), argTypes,
args, {}, funcAttrs, op.getOperation())
diff --git a/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
index fe116dca73e73..e14c8253baa3f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
@@ -366,9 +366,6 @@ LogicalResult TruncfOp::verify() {
if (isa<VectorType>(srcTy)) {
VectorType srcVecTy = dyn_cast<VectorType>(srcTy);
VectorType dstVecTy = dyn_cast<VectorType>(dstTy);
- if (srcVecTy.getNumElements() != dstVecTy.getNumElements())
- return emitOpError(
- "src and dst vector types should have the same number of elements");
if (srcVecTy.getElementTypeBitWidth() <= dstVecTy.getElementTypeBitWidth())
return emitError(
"dst element bitwidth should be less than src element bitwidth");
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index e849b59b846f7..355497599af12 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -2027,14 +2027,6 @@ llvm.func @invalid_xevm_truncf_1(%arg0: vector<8xf16>) {
// -----
-llvm.func @invalid_xevm_truncf_1(%arg0: vector<8xf16>) {
- // expected-error at +1 {{op src and dst vector types should have the same number of elements}}
- %0 = xevm.truncf %arg0 { src_etype = f16, dst_etype = bf8 } : (vector<8xf16>) -> 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/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir
new file mode 100644
index 0000000000000..6aefef1742184
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir
@@ -0,0 +1,147 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=lane zebin-chip=cri" \
+// 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
+
+// XFAIL: *
+module @gemm attributes {gpu.container_module} {
+ gpu.module @kernel {
+ gpu.func @block_scaled_dpas_e2m1(%a: !llvm.ptr<1>, %b: !llvm.ptr<1>, %c: !llvm.ptr<1>) kernel {
+ // TODO: some values are related can be derived from others like the following.
+ // %M = arith.constant 8 : i32
+ // %N = arith.constant 16 : i32
+ // %K = arith.constant 8 : i32
+ // %load_a_elem_bitwidth = arith.constant 32 : i32
+ // %a_elem_bitwidth = arith.constant 16 : i32
+ // %mx_elem_bitwidth = arith.constant 8 : i32
+ // %load_a_pack_ratio = arith.divsi %load_a_elem_bitwidth, %a_elem_bitwidth : i32
+ // %mx_pack_ratio = arith.divsi %load_a_elem_bitwidth, %mx_elem_bitwidth : i32
+ // %load_a_K = arith.muli %K, %load_a_pack_ratio : i32
+ // %load_b_K = arith.muli %K, %mx_pack_ratio : i32
+
+ %x = arith.constant 0 : i32
+ %y = arith.constant 0 : i32
+ // Use constant A for simplifying the example. In real case, A will be loaded from memory.
+ %loaded_a_casted = arith.constant dense<1.0> : vector<16xf16>
+ %a_trunc_partial = xevm.truncf %loaded_a_casted { src_etype = f16, dst_etype = e2m1 } : (vector<16xf16>) -> vector<8xi8>
+ %a_trunc = vector.shuffle %a_trunc_partial, %a_trunc_partial [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] : vector<8xi8>, vector<8xi8>
+ %a_trunc_casted = vector.bitcast %a_trunc : vector<16xi8> to vector<8xi16>
+
+ %base_width_b = arith.constant 16 : i32
+ %base_height_b = arith.constant 32 : i32
+ %base_pitch_b = arith.constant 16 : i32
+ // B is already in e2m1, and it will be used as is for MMA.
+ // So the blockload2d op is configured to load normally with 8bit element bitwidth
+ // with pack_register request.
+ %loaded_b = xevm.blockload2d %b, %base_width_b, %base_height_b, %base_pitch_b, %x, %y
+ <{elem_size_in_bits=8 : i32, tile_width=16 : i32, tile_height=32 : i32, v_blocks=1 : i32,
+ transpose=false, pack_register=true}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
+
+ // Note: scale is not computed. Constant values are used for simplifying the example
+ %scale_a = arith.constant 1.0 : f8E8M0FNU
+ %scale_b = arith.constant 1.0 : f8E8M0FNU
+ %scale_a_casted = arith.bitcast %scale_a : f8E8M0FNU to i8
+ %scale_b_casted = arith.bitcast %scale_b : f8E8M0FNU to i8
+ // Note: c is not loaded. constant vector is used for simplifying the example
+ %loaded_c_casted = arith.constant dense<0.0> : vector<8xf32>
+
+ %c_result = xevm.mma_mx %a_trunc_casted, %loaded_b, %scale_a_casted, %scale_b_casted, %loaded_c_casted
+ {shape=<m=8, n=16, k=32>, types=<d=f32, a=e2m1, b=e2m1, c=f32>}
+ : (vector<8xi16>, vector<8xi32>, i8, i8, vector<8xf32>) -> vector<8xf32>
+ %c_result_casted = vector.bitcast %c_result : vector<8xf32> to vector<8xi32>
+
+ %base_width_c = arith.constant 64 : i32
+ %base_height_c = arith.constant 8 : i32
+ %base_pitch_c = arith.constant 64 : i32
+ xevm.blockstore2d %c, %base_width_c, %base_height_c, %base_pitch_c, %x, %y, %c_result_casted
+ <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32}>
+ : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)
+ gpu.return
+ }
+ }
+
+ func.func @test(%a : memref<8x64xf16>, %b : memref<64x16xf4E2M1FN>, %c : memref<8x16xf32>) -> memref<8x16xf32> attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+
+ %memref_a = gpu.alloc() : memref<8x64xf16>
+ gpu.memcpy %memref_a, %a : memref<8x64xf16>, memref<8x64xf16>
+ %a_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_a : memref<8x64xf16> -> index
+ %a_ptr_as_i64 = arith.index_cast %a_ptr_as_idx : index to i64
+ %a_ptr = llvm.inttoptr %a_ptr_as_i64 : i64 to !llvm.ptr
+ %a_ptr_casted = llvm.addrspacecast %a_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ %memref_b = gpu.alloc() : memref<64x16xf4E2M1FN>
+ gpu.memcpy %memref_b, %b : memref<64x16xf4E2M1FN>, memref<64x16xf4E2M1FN>
+ %b_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_b : memref<64x16xf4E2M1FN> -> index
+ %b_ptr_as_i64 = arith.index_cast %b_ptr_as_idx : index to i64
+ %b_ptr = llvm.inttoptr %b_ptr_as_i64 : i64 to !llvm.ptr
+ %b_ptr_casted = llvm.addrspacecast %b_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ %memref_c = gpu.alloc() : memref<8x16xf32>
+ gpu.memcpy %memref_c, %c : memref<8x16xf32>, memref<8x16xf32>
+ %c_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_c : memref<8x16xf32> -> index
+ %c_ptr_as_i64 = arith.index_cast %c_ptr_as_idx : index to i64
+ %c_ptr = llvm.inttoptr %c_ptr_as_i64 : i64 to !llvm.ptr
+ %c_ptr_casted = llvm.addrspacecast %c_ptr : !llvm.ptr to !llvm.ptr<1>
+
+ gpu.launch_func @kernel::@block_scaled_dpas_e2m1 blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)
+ args(%a_ptr_casted : !llvm.ptr<1>, %b_ptr_casted : !llvm.ptr<1>, %c_ptr_casted : !llvm.ptr<1>)
+ gpu.dealloc %memref_a : memref<8x64xf16>
+ gpu.dealloc %memref_b : memref<64x16xf4E2M1FN>
+ %res = memref.alloc() : memref<8x16xf32>
+ gpu.memcpy %res, %memref_c : memref<8x16xf32>, memref<8x16xf32>
+ gpu.dealloc %memref_c : memref<8x16xf32>
+ return %res : memref<8x16xf32>
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c64 = arith.constant 64 : index
+ %c1f16 = arith.constant 1.0 : f16
+ %c1e2m1 = arith.constant 1.0 : f4E2M1FN
+ %c0f32 = arith.constant 0.0 : f32
+
+ %A = memref.alloc() : memref<8x64xf16>
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.for %j = %c0 to %c64 step %c1 {
+ memref.store %c1f16, %A[%i, %j] : memref<8x64xf16>
+ }
+ }
+
+ %B = memref.alloc() : memref<64x16xf4E2M1FN>
+ scf.for %i = %c0 to %c64 step %c1 {
+ scf.for %j = %c0 to %c16 step %c1 {
+ memref.store %c1e2m1, %B[%i, %j] : memref<64x16xf4E2M1FN>
+ }
+ }
+
+ %C = memref.alloc() : memref<8x16xf32>
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.for %j = %c0 to %c16 step %c1 {
+ memref.store %c0f32, %C[%i, %j] : memref<8x16xf32>
+ }
+ }
+
+ %C_res = call @test(%A, %B, %C) : (memref<8x64xf16>, memref<64x16xf4E2M1FN>, memref<8x16xf32>) -> memref<8x16xf32>
+ %C_cast = memref.cast %C_res : memref<8x16xf32> to memref<*xf32>
+ call @printMemrefF32(%C_cast) : (memref<*xf32>) -> ()
+
+ // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
+ // CHECK-COUNT-8: [64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64]
+ memref.dealloc %A : memref<8x64xf16>
+ memref.dealloc %B : memref<64x16xf4E2M1FN>
+ memref.dealloc %C : memref<8x16xf32>
+ memref.dealloc %C_res : memref<8x16xf32>
+ return
+ }
+ func.func private @printMemrefF32(%ptr : memref<*xf32>) attributes { llvm.emit_c_interface }
+
+}
>From ca7cfdc668c7b04bb753a8139d8e296b1bd547ab Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 27 Apr 2026 15:57:42 -0700
Subject: [PATCH 07/10] Handle fp4 dst type first as bf16 does not need to be
converted to f16 first in such case.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 107 ++++++++++--------
1 file changed, 57 insertions(+), 50 deletions(-)
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index 3e3d67256c5f0..e6e502f4028ef 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1167,56 +1167,9 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
/*targetMem1=*/LLVM::ModRefInfo::NoModRef);
auto funcAttrs = convergentNoUnwindWillReturnAttrs;
funcAttrs.memEffectsAttr = memAttr;
- // BF16 type needs some preprocessing before conversion,
- // First extended to F32 and then truncated to F16.
- if (srcEtype == TruncfSrcElemTypes::BF16) {
- // Step 1: Extend to F32
- // Use float16 __builtin_IB_bftof_16(short16)
- src = LLVM::BitcastOp::create(
- rewriter, op.getLoc(),
- VectorType::get(vecSrcTy.getShape(), rewriter.getI16Type()), src);
- std::string fnName = "__builtin_IB_bftof_16";
- SmallVector<Type> argTypes{src.getType()};
- SmallVector<Value> args{src};
- Type resTy = VectorType::get(vecSrcTy.getShape(), rewriter.getF32Type());
- src = createDeviceFunctionCall(rewriter, fnName, resTy, argTypes, args,
- {}, funcAttrs, op.getOperation())
- ->getResult(0);
- // Step 2: Truncf to F16
- // Use half16 convert_half16(float16)
- std::string truncFnName = "convert_half16";
- SmallVector<Type> truncArgTypes{src.getType()};
- SmallVector<Value> truncArgs{src};
- truncFnName = mangle(truncFnName, truncArgTypes);
- resTy = VectorType::get(vecSrcTy.getShape(), rewriter.getF16Type());
- src =
- createDeviceFunctionCall(rewriter, truncFnName, resTy, truncArgTypes,
- truncArgs, {}, funcAttrs, op.getOperation())
- ->getResult(0);
- }
- if (dstEtype == TruncfDstElemTypes::BF8) {
- // Use char16 __builtin_IB_hftobf8_16(half16)
- std::string fnName = "__builtin_IB_hftobf8_16";
- SmallVector<Type> argTypes{src.getType()};
- SmallVector<Value> args{src};
- Value result =
- createDeviceFunctionCall(rewriter, fnName, vecDstTy, argTypes, args,
- {}, funcAttrs, op.getOperation())
- ->getResult(0);
- rewriter.replaceOp(op, result);
- } else if (dstEtype == TruncfDstElemTypes::F8) {
- // Use char16 __builtin_IB_hftohf8_16(half16)
- std::string fnName = "__builtin_IB_hftohf8_16";
- SmallVector<Type> argTypes{src.getType()};
- SmallVector<Value> args{src};
- Value result =
- createDeviceFunctionCall(rewriter, fnName, vecDstTy, argTypes, args,
- {}, funcAttrs, op.getOperation())
- ->getResult(0);
-
- rewriter.replaceOp(op, result);
- } else if (dstEtype == TruncfDstElemTypes::E2M1) {
+ // Handle the case where dst type is fp4 first.
+ if (dstEtype == TruncfDstElemTypes::E2M1) {
// Convert 8 elements at a time.
// To convert 8 elements, vector<8xf16>:
// Use:
@@ -1227,9 +1180,10 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
rewriter, op.getLoc(), VectorType::get(8, rewriter.getI32Type()),
src);
+ std::string fnName = "__builtin_IB_dnscl_";
+ fnName += (srcEtype == TruncfSrcElemTypes::F16) ? "hf16" : "bf16";
auto genDnscl = [&](Value input, Value idx0, Value idx1, Value dstTy,
Value mode) -> Value {
- std::string fnName = "__builtin_IB_dnscl_hf16";
Value arg1 =
LLVM::ExtractElementOp::create(rewriter, op.getLoc(), input, idx0)
->getResult(0);
@@ -1281,6 +1235,59 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
Value result =
LLVM::BitcastOp::create(rewriter, op.getLoc(), vecDstTy, combined);
rewriter.replaceOp(op, result);
+ return success();
+ }
+
+ // Handle the case where dst type is fp8.
+ // BF16 type needs some preprocessing before conversion,
+ // First extended to F32 and then truncated to F16.
+ if (srcEtype == TruncfSrcElemTypes::BF16) {
+ // Step 1: Extend to F32
+ // Use float16 __builtin_IB_bftof_16(short16)
+ src = LLVM::BitcastOp::create(
+ rewriter, op.getLoc(),
+ VectorType::get(vecSrcTy.getShape(), rewriter.getI16Type()), src);
+ std::string fnName = "__builtin_IB_bftof_16";
+ SmallVector<Type> argTypes{src.getType()};
+ SmallVector<Value> args{src};
+ Type resTy = VectorType::get(vecSrcTy.getShape(), rewriter.getF32Type());
+ src = createDeviceFunctionCall(rewriter, fnName, resTy, argTypes, args,
+ {}, funcAttrs, op.getOperation())
+ ->getResult(0);
+ // Step 2: Truncf to F16
+ // Use half16 convert_half16(float16)
+ std::string truncFnName = "convert_half16";
+ SmallVector<Type> truncArgTypes{src.getType()};
+ SmallVector<Value> truncArgs{src};
+ truncFnName = mangle(truncFnName, truncArgTypes);
+ resTy = VectorType::get(vecSrcTy.getShape(), rewriter.getF16Type());
+ src =
+ createDeviceFunctionCall(rewriter, truncFnName, resTy, truncArgTypes,
+ truncArgs, {}, funcAttrs, op.getOperation())
+ ->getResult(0);
+ }
+ if (dstEtype == TruncfDstElemTypes::BF8) {
+ // Use char16 __builtin_IB_hftobf8_16(half16)
+ std::string fnName = "__builtin_IB_hftobf8_16";
+ SmallVector<Type> argTypes{src.getType()};
+ SmallVector<Value> args{src};
+ Value result =
+ createDeviceFunctionCall(rewriter, fnName, vecDstTy, argTypes, args,
+ {}, funcAttrs, op.getOperation())
+ ->getResult(0);
+
+ rewriter.replaceOp(op, result);
+ } else if (dstEtype == TruncfDstElemTypes::F8) {
+ // Use char16 __builtin_IB_hftohf8_16(half16)
+ std::string fnName = "__builtin_IB_hftohf8_16";
+ SmallVector<Type> argTypes{src.getType()};
+ SmallVector<Value> args{src};
+ Value result =
+ createDeviceFunctionCall(rewriter, fnName, vecDstTy, argTypes, args,
+ {}, funcAttrs, op.getOperation())
+ ->getResult(0);
+
+ rewriter.replaceOp(op, result);
} else {
return rewriter.notifyMatchFailure(
op, "Unsupported src, dst element type pair.");
>From a1ed92a0d3918d9740b72f91fea8504ab46d97c9 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 27 Apr 2026 16:32:19 -0700
Subject: [PATCH 08/10] [MLIR][XeVM] Add more xevm.truncf conversion tests.
---
.../XeVMToLLVM/xevm_mx-to-llvm.mlir | 162 +++++++++++++++++-
1 file changed, 154 insertions(+), 8 deletions(-)
diff --git a/mlir/test/Conversion/XeVMToLLVM/xevm_mx-to-llvm.mlir b/mlir/test/Conversion/XeVMToLLVM/xevm_mx-to-llvm.mlir
index d0a7bc744d311..097d406be6f35 100644
--- a/mlir/test/Conversion/XeVMToLLVM/xevm_mx-to-llvm.mlir
+++ b/mlir/test/Conversion/XeVMToLLVM/xevm_mx-to-llvm.mlir
@@ -1,22 +1,168 @@
// RUN: mlir-opt --convert-xevm-to-llvm --split-input-file %s | FileCheck %s
+// CHECK: llvm.func spir_funccc @__builtin_IB_hftobf8_16(vector<16xf16>) -> vector<16xi8>
+// CHECK-SAME: attributes {convergent, memory_effects = #llvm.memory_effects<other = none,
+// CHECK-SAME: argMem = none, inaccessibleMem = none, errnoMem = none,
+// CHECK-SAME: targetMem0 = none, targetMem1 = none>, no_unwind, will_return}
// CHECK-LABEL: llvm.func @truncf_f16_to_bf8
// CHECK-SAME: %[[ARG0:.*]]: vector<16xf16>
llvm.func @truncf_f16_to_bf8(%src: vector<16xf16>) -> vector<16xi8> {
- // CHECK: %[[VAR0:.*]] = llvm.shufflevector %[[ARG0]], %[[ARG0]] [0, 1, 2, 3, 4, 5, 6, 7] : vector<16xf16>
- // CHECK: %[[VAR1:.*]] = llvm.shufflevector %[[ARG0]], %[[ARG0]] [8, 9, 10, 11, 12, 13, 14, 15] : vector<16xf16>
- // CHECK: %[[VAR2:.*]] = llvm.bitcast %[[VAR0]] : vector<8xf16> to vector<16xi8>
- // CHECK: %[[VAR3:.*]] = llvm.bitcast %[[VAR1]] : vector<8xf16> to vector<16xi8>
- // CHECK: %[[VAR4:.*]] = llvm.shufflevector %[[VAR2]], %[[VAR2]] [1, 3, 5, 7, 9, 11, 13, 15] : vector<16xi8>
- // CHECK: %[[VAR5:.*]] = llvm.shufflevector %[[VAR3]], %[[VAR3]] [1, 3, 5, 7, 9, 11, 13, 15] : vector<16xi8>
- // CHECK: %[[VAR6:.*]] = llvm.shufflevector %4, %5 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] : vector<8xi8>
+ // CHECK: %[[VAR0:.*]] = llvm.call spir_funccc @__builtin_IB_hftobf8_16(%[[ARG0]])
+ // CHECK-SAME: {convergent, function_type = !llvm.func<vector<16xi8> (vector<16xf16>)>,
+ // CHECK-SAME: linkage = #llvm.linkage<external>, memory_effects = #llvm.memory_effects<other = none,
+ // CHECK-SAME: argMem = none, inaccessibleMem = none, errnoMem = none,
+ // CHECK-SAME: targetMem0 = none, targetMem1 = none>,
+ // CHECK-SAME: no_unwind, sym_name = "__builtin_IB_hftobf8_16",
+ // CHECK-SAME: visibility_ = 0 : i64, will_return} :
+ // CHECK-SAME: (vector<16xf16>) -> vector<16xi8>
%dst = xevm.truncf %src { src_etype = f16, dst_etype = bf8 } : (vector<16xf16>) -> vector<16xi8>
llvm.return %dst : vector<16xi8>
}
// -----
-// CHECK-LABEL: llvm.func spir_funccc @__builtin_IB_sub_group16_bdpas_f_f_bf8_bf8_8_8
+// CHECK: llvm.func spir_funccc @__builtin_IB_hftohf8_16(vector<16xf16>) -> vector<16xi8>
+// CHECK-SAME: attributes {convergent, memory_effects = #llvm.memory_effects<other = none,
+// CHECK-SAME: argMem = none, inaccessibleMem = none, errnoMem = none,
+// CHECK-SAME: targetMem0 = none, targetMem1 = none>, no_unwind, will_return}
+// CHECK-LABEL: llvm.func @truncf_f16_to_hf8
+// CHECK-SAME: %[[ARG0:.*]]: vector<16xf16>
+llvm.func @truncf_f16_to_hf8(%src: vector<16xf16>) -> vector<16xi8> {
+ // CHECK: %[[VAR0:.*]] = llvm.call spir_funccc @__builtin_IB_hftohf8_16(%[[ARG0]])
+ // CHECK-SAME: {convergent, function_type = !llvm.func<vector<16xi8> (vector<16xf16>)>,
+ // CHECK-SAME: linkage = #llvm.linkage<external>, memory_effects = #llvm.memory_effects<other = none,
+ // CHECK-SAME: argMem = none, inaccessibleMem = none, errnoMem = none,
+ // CHECK-SAME: targetMem0 = none, targetMem1 = none>,
+ // CHECK-SAME: no_unwind, sym_name = "__builtin_IB_hftohf8_16",
+ // CHECK-SAME: visibility_ = 0 : i64, will_return} :
+ // CHECK-SAME: (vector<16xf16>) -> vector<16xi8>
+ %dst = xevm.truncf %src { src_etype = f16, dst_etype = f8 } : (vector<16xf16>) -> vector<16xi8>
+ llvm.return %dst : vector<16xi8>
+}
+
+// -----
+
+// CHECK: llvm.func spir_funccc @__builtin_IB_hftobf8_16(vector<16xf16>) -> vector<16xi8>
+// CHECK: llvm.func spir_funccc @_Z14convert_half16Dv16_f(vector<16xf32>) -> vector<16xf16>
+// CHECK: llvm.func spir_funccc @__builtin_IB_bftof_16(vector<16xi16>) -> vector<16xf32>
+// CHECK-LABEL: llvm.func @truncf_bf16_to_bf8
+// CHECK-SAME: %[[ARG0:.*]]: vector<16xbf16>
+llvm.func @truncf_bf16_to_bf8(%src: vector<16xbf16>) -> vector<16xi8> {
+ // CHECK: %[[VAR0:.*]] = llvm.bitcast %[[ARG0]] : vector<16xbf16> to vector<16xi16>
+ // CHECK: %[[VAR1:.*]] = llvm.call spir_funccc @__builtin_IB_bftof_16(%[[VAR0]])
+ // CHECK-SAME: : (vector<16xi16>) -> vector<16xf32>
+ // CHECK: %[[VAR2:.*]] = llvm.call spir_funccc @_Z14convert_half16Dv16_f(%[[VAR1]])
+ // CHECK-SAME: : (vector<16xf32>) -> vector<16xf16>
+ // CHECK: %[[VAR3:.*]] = llvm.call spir_funccc @__builtin_IB_hftobf8_16(%[[VAR2]])
+ // CHECK-SAME: : (vector<16xf16>) -> vector<16xi8>
+ %dst = xevm.truncf %src { src_etype = bf16, dst_etype = bf8 } : (vector<16xbf16>) -> vector<16xi8>
+ llvm.return %dst : vector<16xi8>
+}
+
+// -----
+
+// CHECK: llvm.func spir_funccc @__builtin_IB_hftohf8_16(vector<16xf16>) -> vector<16xi8>
+// CHECK: llvm.func spir_funccc @_Z14convert_half16Dv16_f(vector<16xf32>) -> vector<16xf16>
+// CHECK: llvm.func spir_funccc @__builtin_IB_bftof_16(vector<16xi16>) -> vector<16xf32>
+// CHECK-LABEL: llvm.func @truncf_bf16_to_hf8
+// CHECK-SAME: %[[ARG0:.*]]: vector<16xbf16>
+llvm.func @truncf_bf16_to_hf8(%src: vector<16xbf16>) -> vector<16xi8> {
+ // CHECK: %[[VAR0:.*]] = llvm.bitcast %[[ARG0]] : vector<16xbf16> to vector<16xi16>
+ // CHECK: %[[VAR1:.*]] = llvm.call spir_funccc @__builtin_IB_bftof_16(%[[VAR0]])
+ // CHECK-SAME: : (vector<16xi16>) -> vector<16xf32>
+ // CHECK: %[[VAR2:.*]] = llvm.call spir_funccc @_Z14convert_half16Dv16_f(%[[VAR1]])
+ // CHECK-SAME: : (vector<16xf32>) -> vector<16xf16>
+ // CHECK: %[[VAR3:.*]] = llvm.call spir_funccc @__builtin_IB_hftohf8_16(%[[VAR2]])
+ // CHECK-SAME: : (vector<16xf16>) -> vector<16xi8>
+ %dst = xevm.truncf %src { src_etype = bf16, dst_etype = f8 } : (vector<16xbf16>) -> vector<16xi8>
+ llvm.return %dst : vector<16xi8>
+}
+
+// -----
+
+// CHECK: llvm.func spir_funccc @__builtin_IB_dnscl_hf16(i32, i32, i32, i32) -> i32
+// CHECK-LABEL: llvm.func @truncf_f16_to_e2m1
+// CHECK-SAME: %[[ARG0:.*]]: vector<16xf16>
+llvm.func @truncf_f16_to_e2m1(%src: vector<16xf16>) -> vector<8xi8> {
+ // CHECK: %[[UNDEF:.*]] = llvm.mlir.undef : vector<2xi32>
+ // CHECK: %[[C7:.*]] = llvm.mlir.constant(7 : i32) : i32
+ // CHECK: %[[C6:.*]] = llvm.mlir.constant(6 : i32) : i32
+ // CHECK: %[[C5:.*]] = llvm.mlir.constant(5 : i32) : i32
+ // CHECK: %[[C4:.*]] = llvm.mlir.constant(4 : i32) : i32
+ // CHECK: %[[C3:.*]] = llvm.mlir.constant(3 : i32) : i32
+ // CHECK: %[[C2:.*]] = llvm.mlir.constant(2 : i32) : i32
+ // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK: %[[BC:.*]] = llvm.bitcast %[[ARG0]] : vector<16xf16> to vector<8xi32>
+ // CHECK: %[[E0:.*]] = llvm.extractelement %[[BC]][%[[C0]] : i32] : vector<8xi32>
+ // CHECK: %[[E2:.*]] = llvm.extractelement %[[BC]][%[[C2]] : i32] : vector<8xi32>
+ // CHECK: %[[CALL0:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_hf16(%[[E0]], %[[E2]], %[[C1]], %[[C0]])
+ // CHECK-SAME: : (i32, i32, i32, i32) -> i32
+ // CHECK: %[[E1:.*]] = llvm.extractelement %[[BC]][%[[C1]] : i32] : vector<8xi32>
+ // CHECK: %[[E3:.*]] = llvm.extractelement %[[BC]][%[[C3]] : i32] : vector<8xi32>
+ // CHECK: %[[CALL1:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_hf16(%[[E1]], %[[E3]], %[[C1]], %[[C3]])
+ // CHECK-SAME: : (i32, i32, i32, i32) -> i32
+ // CHECK: %[[OR0:.*]] = llvm.or %[[CALL0]], %[[CALL1]] : i32
+ // CHECK: %[[E4:.*]] = llvm.extractelement %[[BC]][%[[C4]] : i32] : vector<8xi32>
+ // CHECK: %[[E6:.*]] = llvm.extractelement %[[BC]][%[[C6]] : i32] : vector<8xi32>
+ // CHECK: %[[CALL2:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_hf16(%[[E4]], %[[E6]], %[[C1]], %[[C0]])
+ // CHECK-SAME: : (i32, i32, i32, i32) -> i32
+ // CHECK: %[[E5:.*]] = llvm.extractelement %[[BC]][%[[C5]] : i32] : vector<8xi32>
+ // CHECK: %[[E7:.*]] = llvm.extractelement %[[BC]][%[[C7]] : i32] : vector<8xi32>
+ // CHECK: %[[CALL3:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_hf16(%[[E5]], %[[E7]], %[[C1]], %[[C3]])
+ // CHECK-SAME: : (i32, i32, i32, i32) -> i32
+ // CHECK: %[[OR1:.*]] = llvm.or %[[CALL2]], %[[CALL3]] : i32
+ // CHECK: %[[INS0:.*]] = llvm.insertelement %[[OR0]], %[[UNDEF]][%[[C0]] : i32] : vector<2xi32>
+ // CHECK: %[[INS1:.*]] = llvm.insertelement %[[OR1]], %[[INS0]][%[[C1]] : i32] : vector<2xi32>
+ // CHECK: %[[RES:.*]] = llvm.bitcast %[[INS1]] : vector<2xi32> to vector<8xi8>
+ %dst = xevm.truncf %src { src_etype = f16, dst_etype = e2m1 } : (vector<16xf16>) -> vector<8xi8>
+ llvm.return %dst : vector<8xi8>
+}
+
+// -----
+
+// CHECK: llvm.func spir_funccc @__builtin_IB_dnscl_bf16(i32, i32, i32, i32) -> i32
+// CHECK-LABEL: llvm.func @truncf_bf16_to_e2m1
+// CHECK-SAME: %[[ARG0:.*]]: vector<16xbf16>
+llvm.func @truncf_bf16_to_e2m1(%src: vector<16xbf16>) -> vector<8xi8> {
+ // CHECK: %[[UNDEF:.*]] = llvm.mlir.undef : vector<2xi32>
+ // CHECK: %[[C7:.*]] = llvm.mlir.constant(7 : i32) : i32
+ // CHECK: %[[C6:.*]] = llvm.mlir.constant(6 : i32) : i32
+ // CHECK: %[[C5:.*]] = llvm.mlir.constant(5 : i32) : i32
+ // CHECK: %[[C4:.*]] = llvm.mlir.constant(4 : i32) : i32
+ // CHECK: %[[C3:.*]] = llvm.mlir.constant(3 : i32) : i32
+ // CHECK: %[[C2:.*]] = llvm.mlir.constant(2 : i32) : i32
+ // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK: %[[BC:.*]] = llvm.bitcast %[[ARG0]] : vector<16xbf16> to vector<8xi32>
+ // CHECK: %[[E0:.*]] = llvm.extractelement %[[BC]][%[[C0]] : i32] : vector<8xi32>
+ // CHECK: %[[E2:.*]] = llvm.extractelement %[[BC]][%[[C2]] : i32] : vector<8xi32>
+ // CHECK: %[[CALL0:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_bf16(%[[E0]], %[[E2]], %[[C1]], %[[C0]])
+ // CHECK-SAME: : (i32, i32, i32, i32) -> i32
+ // CHECK: %[[E1:.*]] = llvm.extractelement %[[BC]][%[[C1]] : i32] : vector<8xi32>
+ // CHECK: %[[E3:.*]] = llvm.extractelement %[[BC]][%[[C3]] : i32] : vector<8xi32>
+ // CHECK: %[[CALL1:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_bf16(%[[E1]], %[[E3]], %[[C1]], %[[C3]])
+ // CHECK-SAME: : (i32, i32, i32, i32) -> i32
+ // CHECK: %[[OR0:.*]] = llvm.or %[[CALL0]], %[[CALL1]] : i32
+ // CHECK: %[[E4:.*]] = llvm.extractelement %[[BC]][%[[C4]] : i32] : vector<8xi32>
+ // CHECK: %[[E6:.*]] = llvm.extractelement %[[BC]][%[[C6]] : i32] : vector<8xi32>
+ // CHECK: %[[CALL2:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_bf16(%[[E4]], %[[E6]], %[[C1]], %[[C0]])
+ // CHECK-SAME: : (i32, i32, i32, i32) -> i32
+ // CHECK: %[[E5:.*]] = llvm.extractelement %[[BC]][%[[C5]] : i32] : vector<8xi32>
+ // CHECK: %[[E7:.*]] = llvm.extractelement %[[BC]][%[[C7]] : i32] : vector<8xi32>
+ // CHECK: %[[CALL3:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_bf16(%[[E5]], %[[E7]], %[[C1]], %[[C3]])
+ // CHECK-SAME: : (i32, i32, i32, i32) -> i32
+ // CHECK: %[[OR1:.*]] = llvm.or %[[CALL2]], %[[CALL3]] : i32
+ // CHECK: %[[INS0:.*]] = llvm.insertelement %[[OR0]], %[[UNDEF]][%[[C0]] : i32] : vector<2xi32>
+ // CHECK: %[[INS1:.*]] = llvm.insertelement %[[OR1]], %[[INS0]][%[[C1]] : i32] : vector<2xi32>
+ // CHECK: %[[RES:.*]] = llvm.bitcast %[[INS1]] : vector<2xi32> to vector<8xi8>
+ %dst = xevm.truncf %src { src_etype = bf16, dst_etype = e2m1 } : (vector<16xbf16>) -> vector<8xi8>
+ llvm.return %dst : vector<8xi8>
+}
+
+// -----
+
+// CHECK: llvm.func spir_funccc @__builtin_IB_sub_group16_bdpas_f_f_bf8_bf8_8_8
// CHECK-SAME: (vector<8xf32>, vector<8xi16>, vector<8xi32>, i8, i8) -> vector<8xf32>
// CHECK-SAME: attributes {convergent, memory_effects = #llvm.memory_effects<other = none,
// CHECK-SAME: argMem = none, inaccessibleMem = none, errnoMem = none,
>From c43bfcfc42e221ed6622c76344f8d56dd21d70a2 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 28 Apr 2026 15:33:45 -0700
Subject: [PATCH 09/10] Use correct dnscl mode. Update e2m1 conversion and
integration tests.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 4 +-
.../XeVMToLLVM/xevm_mx-to-llvm.mlir | 8 ++--
.../XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir | 39 ++++++++++---------
3 files changed, 27 insertions(+), 24 deletions(-)
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index e6e502f4028ef..bafb1b08fa1ed 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1209,7 +1209,7 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
Value three = LLVM::ConstantOp::create(rewriter, op.getLoc(),
rewriter.getI32Type(), 3);
Value even = genDnscl(cast, zero, two, one, zero);
- Value odd = genDnscl(cast, one, three, one, three);
+ Value odd = genDnscl(cast, one, three, one, two);
Value firstHalf = LLVM::OrOp::create(rewriter, op.getLoc(), even, odd);
Value four = LLVM::ConstantOp::create(rewriter, op.getLoc(),
rewriter.getI32Type(), 4);
@@ -1220,7 +1220,7 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
Value seven = LLVM::ConstantOp::create(rewriter, op.getLoc(),
rewriter.getI32Type(), 7);
even = genDnscl(cast, four, six, one, zero);
- odd = genDnscl(cast, five, seven, one, three);
+ odd = genDnscl(cast, five, seven, one, two);
Value secondHalf = LLVM::OrOp::create(rewriter, op.getLoc(), even, odd);
// Create vector<2xi32> from two i32 values and then bitcast to
// vector<8xi8> to match the dst type.
diff --git a/mlir/test/Conversion/XeVMToLLVM/xevm_mx-to-llvm.mlir b/mlir/test/Conversion/XeVMToLLVM/xevm_mx-to-llvm.mlir
index 097d406be6f35..9706887182003 100644
--- a/mlir/test/Conversion/XeVMToLLVM/xevm_mx-to-llvm.mlir
+++ b/mlir/test/Conversion/XeVMToLLVM/xevm_mx-to-llvm.mlir
@@ -100,7 +100,7 @@ llvm.func @truncf_f16_to_e2m1(%src: vector<16xf16>) -> vector<8xi8> {
// CHECK-SAME: : (i32, i32, i32, i32) -> i32
// CHECK: %[[E1:.*]] = llvm.extractelement %[[BC]][%[[C1]] : i32] : vector<8xi32>
// CHECK: %[[E3:.*]] = llvm.extractelement %[[BC]][%[[C3]] : i32] : vector<8xi32>
- // CHECK: %[[CALL1:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_hf16(%[[E1]], %[[E3]], %[[C1]], %[[C3]])
+ // CHECK: %[[CALL1:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_hf16(%[[E1]], %[[E3]], %[[C1]], %[[C2]])
// CHECK-SAME: : (i32, i32, i32, i32) -> i32
// CHECK: %[[OR0:.*]] = llvm.or %[[CALL0]], %[[CALL1]] : i32
// CHECK: %[[E4:.*]] = llvm.extractelement %[[BC]][%[[C4]] : i32] : vector<8xi32>
@@ -109,7 +109,7 @@ llvm.func @truncf_f16_to_e2m1(%src: vector<16xf16>) -> vector<8xi8> {
// CHECK-SAME: : (i32, i32, i32, i32) -> i32
// CHECK: %[[E5:.*]] = llvm.extractelement %[[BC]][%[[C5]] : i32] : vector<8xi32>
// CHECK: %[[E7:.*]] = llvm.extractelement %[[BC]][%[[C7]] : i32] : vector<8xi32>
- // CHECK: %[[CALL3:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_hf16(%[[E5]], %[[E7]], %[[C1]], %[[C3]])
+ // CHECK: %[[CALL3:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_hf16(%[[E5]], %[[E7]], %[[C1]], %[[C2]])
// CHECK-SAME: : (i32, i32, i32, i32) -> i32
// CHECK: %[[OR1:.*]] = llvm.or %[[CALL2]], %[[CALL3]] : i32
// CHECK: %[[INS0:.*]] = llvm.insertelement %[[OR0]], %[[UNDEF]][%[[C0]] : i32] : vector<2xi32>
@@ -141,7 +141,7 @@ llvm.func @truncf_bf16_to_e2m1(%src: vector<16xbf16>) -> vector<8xi8> {
// CHECK-SAME: : (i32, i32, i32, i32) -> i32
// CHECK: %[[E1:.*]] = llvm.extractelement %[[BC]][%[[C1]] : i32] : vector<8xi32>
// CHECK: %[[E3:.*]] = llvm.extractelement %[[BC]][%[[C3]] : i32] : vector<8xi32>
- // CHECK: %[[CALL1:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_bf16(%[[E1]], %[[E3]], %[[C1]], %[[C3]])
+ // CHECK: %[[CALL1:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_bf16(%[[E1]], %[[E3]], %[[C1]], %[[C2]])
// CHECK-SAME: : (i32, i32, i32, i32) -> i32
// CHECK: %[[OR0:.*]] = llvm.or %[[CALL0]], %[[CALL1]] : i32
// CHECK: %[[E4:.*]] = llvm.extractelement %[[BC]][%[[C4]] : i32] : vector<8xi32>
@@ -150,7 +150,7 @@ llvm.func @truncf_bf16_to_e2m1(%src: vector<16xbf16>) -> vector<8xi8> {
// CHECK-SAME: : (i32, i32, i32, i32) -> i32
// CHECK: %[[E5:.*]] = llvm.extractelement %[[BC]][%[[C5]] : i32] : vector<8xi32>
// CHECK: %[[E7:.*]] = llvm.extractelement %[[BC]][%[[C7]] : i32] : vector<8xi32>
- // CHECK: %[[CALL3:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_bf16(%[[E5]], %[[E7]], %[[C1]], %[[C3]])
+ // CHECK: %[[CALL3:.*]] = llvm.call spir_funccc @__builtin_IB_dnscl_bf16(%[[E5]], %[[E7]], %[[C1]], %[[C2]])
// CHECK-SAME: : (i32, i32, i32, i32) -> i32
// CHECK: %[[OR1:.*]] = llvm.or %[[CALL2]], %[[CALL3]] : i32
// CHECK: %[[INS0:.*]] = llvm.insertelement %[[OR0]], %[[UNDEF]][%[[C0]] : i32] : vector<2xi32>
diff --git a/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir
index 6aefef1742184..a4e56fbaee324 100644
--- a/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir
+++ b/mlir/test/Integration/Dialect/XeVM/GPU/xevm_block_scaled_dpas_e2m1.mlir
@@ -28,6 +28,8 @@ module @gemm attributes {gpu.container_module} {
%loaded_a_casted = arith.constant dense<1.0> : vector<16xf16>
%a_trunc_partial = xevm.truncf %loaded_a_casted { src_etype = f16, dst_etype = e2m1 } : (vector<16xf16>) -> vector<8xi8>
%a_trunc = vector.shuffle %a_trunc_partial, %a_trunc_partial [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] : vector<8xi8>, vector<8xi8>
+ // %a_trunc = arith.constant dense<0x22> : vector<16xi8>
+
%a_trunc_casted = vector.bitcast %a_trunc : vector<16xi8> to vector<8xi16>
%base_width_b = arith.constant 16 : i32
@@ -41,16 +43,16 @@ module @gemm attributes {gpu.container_module} {
transpose=false, pack_register=true}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
// Note: scale is not computed. Constant values are used for simplifying the example
- %scale_a = arith.constant 1.0 : f8E8M0FNU
- %scale_b = arith.constant 1.0 : f8E8M0FNU
- %scale_a_casted = arith.bitcast %scale_a : f8E8M0FNU to i8
- %scale_b_casted = arith.bitcast %scale_b : f8E8M0FNU to i8
+ %scale_a = arith.constant dense<1.0> : vector<2xf8E8M0FNU>
+ %scale_b = arith.constant dense<1.0> : vector<2xf8E8M0FNU>
+ %scale_a_casted = vector.bitcast %scale_a : vector<2xf8E8M0FNU> to vector<2xi8>
+ %scale_b_casted = vector.bitcast %scale_b : vector<2xf8E8M0FNU> to vector<2xi8>
// Note: c is not loaded. constant vector is used for simplifying the example
%loaded_c_casted = arith.constant dense<0.0> : vector<8xf32>
%c_result = xevm.mma_mx %a_trunc_casted, %loaded_b, %scale_a_casted, %scale_b_casted, %loaded_c_casted
- {shape=<m=8, n=16, k=32>, types=<d=f32, a=e2m1, b=e2m1, c=f32>}
- : (vector<8xi16>, vector<8xi32>, i8, i8, vector<8xf32>) -> vector<8xf32>
+ {shape=<m=8, n=16, k=64>, types=<d=f32, a=e2m1, b=e2m1, c=f32>}
+ : (vector<8xi16>, vector<8xi32>, vector<2xi8>, vector<2xi8>, vector<8xf32>) -> vector<8xf32>
%c_result_casted = vector.bitcast %c_result : vector<8xf32> to vector<8xi32>
%base_width_c = arith.constant 64 : i32
@@ -63,7 +65,7 @@ module @gemm attributes {gpu.container_module} {
}
}
- func.func @test(%a : memref<8x64xf16>, %b : memref<64x16xf4E2M1FN>, %c : memref<8x16xf32>) -> memref<8x16xf32> attributes {llvm.emit_c_interface} {
+ func.func @test(%a : memref<8x64xf16>, %b : memref<32x16xi8>, %c : memref<8x16xf32>) -> memref<8x16xf32> attributes {llvm.emit_c_interface} {
%c1 = arith.constant 1 : index
%c16 = arith.constant 16 : index
@@ -74,9 +76,9 @@ module @gemm attributes {gpu.container_module} {
%a_ptr = llvm.inttoptr %a_ptr_as_i64 : i64 to !llvm.ptr
%a_ptr_casted = llvm.addrspacecast %a_ptr : !llvm.ptr to !llvm.ptr<1>
- %memref_b = gpu.alloc() : memref<64x16xf4E2M1FN>
- gpu.memcpy %memref_b, %b : memref<64x16xf4E2M1FN>, memref<64x16xf4E2M1FN>
- %b_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_b : memref<64x16xf4E2M1FN> -> index
+ %memref_b = gpu.alloc() : memref<32x16xi8>
+ gpu.memcpy %memref_b, %b : memref<32x16xi8>, memref<32x16xi8>
+ %b_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_b : memref<32x16xi8> -> index
%b_ptr_as_i64 = arith.index_cast %b_ptr_as_idx : index to i64
%b_ptr = llvm.inttoptr %b_ptr_as_i64 : i64 to !llvm.ptr
%b_ptr_casted = llvm.addrspacecast %b_ptr : !llvm.ptr to !llvm.ptr<1>
@@ -91,7 +93,7 @@ module @gemm attributes {gpu.container_module} {
gpu.launch_func @kernel::@block_scaled_dpas_e2m1 blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)
args(%a_ptr_casted : !llvm.ptr<1>, %b_ptr_casted : !llvm.ptr<1>, %c_ptr_casted : !llvm.ptr<1>)
gpu.dealloc %memref_a : memref<8x64xf16>
- gpu.dealloc %memref_b : memref<64x16xf4E2M1FN>
+ gpu.dealloc %memref_b : memref<32x16xi8>
%res = memref.alloc() : memref<8x16xf32>
gpu.memcpy %res, %memref_c : memref<8x16xf32>, memref<8x16xf32>
gpu.dealloc %memref_c : memref<8x16xf32>
@@ -104,9 +106,10 @@ module @gemm attributes {gpu.container_module} {
%c1 = arith.constant 1 : index
%c8 = arith.constant 8 : index
%c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
%c64 = arith.constant 64 : index
%c1f16 = arith.constant 1.0 : f16
- %c1e2m1 = arith.constant 1.0 : f4E2M1FN
+ %c1e2m1packed = arith.constant 0x22 : i8
%c0f32 = arith.constant 0.0 : f32
%A = memref.alloc() : memref<8x64xf16>
@@ -116,10 +119,10 @@ module @gemm attributes {gpu.container_module} {
}
}
- %B = memref.alloc() : memref<64x16xf4E2M1FN>
- scf.for %i = %c0 to %c64 step %c1 {
+ %B = memref.alloc() : memref<32x16xi8>
+ scf.for %i = %c0 to %c32 step %c1 {
scf.for %j = %c0 to %c16 step %c1 {
- memref.store %c1e2m1, %B[%i, %j] : memref<64x16xf4E2M1FN>
+ memref.store %c1e2m1packed, %B[%i, %j] : memref<32x16xi8>
}
}
@@ -130,14 +133,14 @@ module @gemm attributes {gpu.container_module} {
}
}
- %C_res = call @test(%A, %B, %C) : (memref<8x64xf16>, memref<64x16xf4E2M1FN>, memref<8x16xf32>) -> memref<8x16xf32>
+ %C_res = call @test(%A, %B, %C) : (memref<8x64xf16>, memref<32x16xi8>, memref<8x16xf32>) -> memref<8x16xf32>
%C_cast = memref.cast %C_res : memref<8x16xf32> to memref<*xf32>
call @printMemrefF32(%C_cast) : (memref<*xf32>) -> ()
// CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}
- // CHECK-COUNT-8: [64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64]
+ // CHECK-COUNT-8: [64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64]
memref.dealloc %A : memref<8x64xf16>
- memref.dealloc %B : memref<64x16xf4E2M1FN>
+ memref.dealloc %B : memref<32x16xi8>
memref.dealloc %C : memref<8x16xf32>
memref.dealloc %C_res : memref<8x16xf32>
return
>From 32080f3ac05d640c12589ef0ca47a5ce312e9906 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Fri, 1 May 2026 08:41:52 -0700
Subject: [PATCH 10/10] Add small comment about the matching MLIR types for bf8
and f8.
---
mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index bafb1b08fa1ed..61f8e8c059c63 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1266,7 +1266,7 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
truncArgs, {}, funcAttrs, op.getOperation())
->getResult(0);
}
- if (dstEtype == TruncfDstElemTypes::BF8) {
+ if (dstEtype == TruncfDstElemTypes::BF8) { // Float8E5M2Type
// Use char16 __builtin_IB_hftobf8_16(half16)
std::string fnName = "__builtin_IB_hftobf8_16";
SmallVector<Type> argTypes{src.getType()};
@@ -1277,7 +1277,7 @@ class TruncfToOCLPattern : public OpConversionPattern<TruncfOp> {
->getResult(0);
rewriter.replaceOp(op, result);
- } else if (dstEtype == TruncfDstElemTypes::F8) {
+ } else if (dstEtype == TruncfDstElemTypes::F8) { // Float8E4M3FNType
// Use char16 __builtin_IB_hftohf8_16(half16)
std::string fnName = "__builtin_IB_hftohf8_16";
SmallVector<Type> argTypes{src.getType()};
More information about the Mlir-commits
mailing list