[Mlir-commits] [mlir] [mlir][xevm] Fold xevm.truncf of an xevm.extf round trip (PR #225875)
Sang Ik Lee
llvmlistbot at llvm.org
Thu Sep 24 07:43:43 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/225875
>From ee8c758e2907b4cc08d08f13ea103b6e6eb8a459 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 23 Sep 2026 18:06:15 +0000
Subject: [PATCH] [mlir][xevm] Fold xevm.truncf of an xevm.extf round trip
A xevm.extf followed by a xevm.truncf through the same pair of formats returns the value it started from, so the pair can be dropped. Fold it on the truncf, the way arith::TruncFOp::fold recognizes truncf(extf(a)), and gate it on APFloatBase::isLosslesslyConvertibleTo of the two formats the element type attributes stand for.
That gate excludes bf8: it is IEEE-like, so extending a signaling NaN quiets it and the original value cannot be recovered, which is the same reason arith does not fold an f8E5M2 round trip. f8 (E4M3, NaN only) and e2m1 (finite only) do fold.
Mark both ops Pure so the extf left without users is removed and the pair folds to nothing.
---
mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td | 5 +-
mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp | 76 +++++++++++++++++
.../Dialect/LLVMIR/xevm-canonicalize.mlir | 83 +++++++++++++++++++
3 files changed, 162 insertions(+), 2 deletions(-)
create mode 100644 mlir/test/Dialect/LLVMIR/xevm-canonicalize.mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
index 2cd7bbb07244c..fdf6263094bb2 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
@@ -627,7 +627,7 @@ def XeVM_TruncfDstElemTypeAttr : XeVM_Attr<"TruncfDstElemType", "dst_etype"> {
}
def XeVM_TruncfOp
- : XeVM_Op<"truncf">,
+ : XeVM_Op<"truncf", [Pure]>,
Results<(outs AnyTypeOf<[FixedVectorOfRankAndType<[1], [I8, I<4>]>, I8,
I<4>]>:$dst)>,
Arguments<(ins AnyTypeOf<[FixedVectorOfRankAndType<[1], [F16, BF16]>, F16,
@@ -644,6 +644,7 @@ def XeVM_TruncfOp
$src ` ` `{` $src_etype `,` $dst_etype `}` attr-dict `:` functional-type(operands, results)
}];
+ let hasFolder = 1;
let hasVerifier = 1;
}
@@ -672,7 +673,7 @@ def XeVM_ExtfDstElemTypeAttr : XeVM_Attr<"ExtfDstElemType", "ext_dst_etype"> {
}
def XeVM_ExtfOp
- : XeVM_Op<"extf">,
+ : XeVM_Op<"extf", [Pure]>,
Results<(outs AnyTypeOf<[FixedVectorOfRankAndType<[1], [F16, BF16]>, F16,
BF16]>:$dst)>,
Arguments<(ins AnyTypeOf<[FixedVectorOfRankAndType<[1], [I8, I<4>]>, I8,
diff --git a/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
index 809e520b59a87..5680ab7ab006e 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/XeVMDialect.cpp
@@ -10,6 +10,7 @@
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/TypeUtilities.h"
+#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/FileSystem.h"
@@ -422,6 +423,81 @@ LogicalResult ExtfOp::verify() {
getNarrowFloatBitWidth(getSrcEtype().getEtype()));
}
+/// Float semantics the element type attributes of `xevm.truncf` and `xevm.extf`
+/// stand for. The narrow formats are the OCP FP8 and FP4 ones: `bf8` is
+/// E5M2, `f8` is E4M3 and `e2m1` is FP4.
+static const llvm::fltSemantics *getFloatSemantics(TruncfSrcElemTypes etype) {
+ switch (etype) {
+ case TruncfSrcElemTypes::F16:
+ return &llvm::APFloat::IEEEhalf();
+ case TruncfSrcElemTypes::BF16:
+ return &llvm::APFloat::BFloat();
+ }
+ return nullptr;
+}
+
+static const llvm::fltSemantics *getFloatSemantics(ExtfDstElemTypes etype) {
+ switch (etype) {
+ case ExtfDstElemTypes::F16:
+ return &llvm::APFloat::IEEEhalf();
+ case ExtfDstElemTypes::BF16:
+ return &llvm::APFloat::BFloat();
+ }
+ return nullptr;
+}
+
+static const llvm::fltSemantics *getFloatSemantics(TruncfDstElemTypes etype) {
+ switch (etype) {
+ case TruncfDstElemTypes::BF8:
+ return &llvm::APFloat::Float8E5M2();
+ case TruncfDstElemTypes::F8:
+ return &llvm::APFloat::Float8E4M3FN();
+ case TruncfDstElemTypes::E2M1:
+ return &llvm::APFloat::Float4E2M1FN();
+ }
+ return nullptr;
+}
+
+static const llvm::fltSemantics *getFloatSemantics(ExtfSrcElemTypes etype) {
+ switch (etype) {
+ case ExtfSrcElemTypes::BF8:
+ return &llvm::APFloat::Float8E5M2();
+ case ExtfSrcElemTypes::F8:
+ return &llvm::APFloat::Float8E4M3FN();
+ case ExtfSrcElemTypes::E2M1:
+ return &llvm::APFloat::Float4E2M1FN();
+ }
+ return nullptr;
+}
+
+/// truncf(extf(a)) -> a, when the two ops convert through the same pair of
+/// formats and every narrow value survives the round trip. `bf8` does not
+/// qualify: it is IEEE-like, and extending a signaling NaN quiets it, so the
+/// original value cannot be recovered. This mirrors `arith.truncf`, which gates
+/// the same fold on `APFloatBase::isLosslesslyConvertibleTo`.
+OpFoldResult TruncfOp::fold(FoldAdaptor) {
+ auto extfOp = getSrc().getDefiningOp<ExtfOp>();
+ if (!extfOp)
+ return {};
+
+ const llvm::fltSemantics *narrowSem =
+ getFloatSemantics(getDstEtype().getEtype());
+ const llvm::fltSemantics *wideSem =
+ getFloatSemantics(getSrcEtype().getEtype());
+ if (narrowSem != getFloatSemantics(extfOp.getSrcEtype().getEtype()) ||
+ wideSem != getFloatSemantics(extfOp.getDstEtype().getEtype()))
+ return {};
+
+ Value narrowSrc = extfOp.getSrc();
+ if (narrowSrc.getType() != getDst().getType())
+ return {};
+
+ if (!llvm::APFloatBase::isLosslesslyConvertibleTo(*narrowSem, *wideSem))
+ return {};
+
+ return narrowSrc;
+}
+
LogicalResult BitcastShuffleOp::verify() {
Type srcTy = getSrc().getType();
Type resTy = getRes().getType();
diff --git a/mlir/test/Dialect/LLVMIR/xevm-canonicalize.mlir b/mlir/test/Dialect/LLVMIR/xevm-canonicalize.mlir
new file mode 100644
index 0000000000000..1cead8514051f
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/xevm-canonicalize.mlir
@@ -0,0 +1,83 @@
+// RUN: mlir-opt -canonicalize %s | FileCheck %s
+
+// An f8 (E4M3) value survives the round trip through f16, so the pair folds
+// away and the packed source is returned directly.
+// CHECK-LABEL: @fold_f8_round_trip
+// CHECK-SAME: (%[[ARG0:.*]]: vector<8xi8>)
+// CHECK-NOT: xevm.extf
+// CHECK-NOT: xevm.truncf
+// CHECK: return %[[ARG0]] : vector<8xi8>
+func.func @fold_f8_round_trip(%arg0: vector<8xi8>) -> vector<8xi8> {
+ %ext = xevm.extf %arg0 { src_etype=f8, dst_etype=f16 } : (vector<8xi8>) -> vector<8xf16>
+ %trunc = xevm.truncf %ext { src_etype=f16, dst_etype=f8 } : (vector<8xf16>) -> vector<8xi8>
+ return %trunc : vector<8xi8>
+}
+
+// e2m1 has no infinities and no NaNs, so it round trips through bf16 as well.
+// CHECK-LABEL: @fold_e2m1_round_trip
+// CHECK-SAME: (%[[ARG0:.*]]: vector<8xi4>)
+// CHECK-NOT: xevm.extf
+// CHECK-NOT: xevm.truncf
+// CHECK: return %[[ARG0]] : vector<8xi4>
+func.func @fold_e2m1_round_trip(%arg0: vector<8xi4>) -> vector<8xi4> {
+ %ext = xevm.extf %arg0 { src_etype=e2m1, dst_etype=bf16 } : (vector<8xi4>) -> vector<8xbf16>
+ %trunc = xevm.truncf %ext { src_etype=bf16, dst_etype=e2m1 } : (vector<8xbf16>) -> vector<8xi4>
+ return %trunc : vector<8xi4>
+}
+
+// CHECK-LABEL: @fold_f8_round_trip_scalar
+// CHECK-SAME: (%[[ARG0:.*]]: i8)
+// CHECK-NOT: xevm.extf
+// CHECK-NOT: xevm.truncf
+// CHECK: return %[[ARG0]] : i8
+func.func @fold_f8_round_trip_scalar(%arg0: i8) -> i8 {
+ %ext = xevm.extf %arg0 { src_etype=f8, dst_etype=f16 } : (i8) -> f16
+ %trunc = xevm.truncf %ext { src_etype=f16, dst_etype=f8 } : (f16) -> i8
+ return %trunc : i8
+}
+
+// bf8 is IEEE-like: extending a signaling NaN quiets it, so the narrow value is
+// not recoverable and the round trip must be kept.
+// CHECK-LABEL: @no_fold_bf8_round_trip
+// CHECK: %[[EXT:.*]] = xevm.extf %{{.*}} {src_etype = bf8, dst_etype = f16}
+// CHECK: %[[TRUNC:.*]] = xevm.truncf %[[EXT]] {src_etype = f16, dst_etype = bf8}
+// CHECK: return %[[TRUNC]]
+func.func @no_fold_bf8_round_trip(%arg0: vector<8xi8>) -> vector<8xi8> {
+ %ext = xevm.extf %arg0 { src_etype=bf8, dst_etype=f16 } : (vector<8xi8>) -> vector<8xf16>
+ %trunc = xevm.truncf %ext { src_etype=f16, dst_etype=bf8 } : (vector<8xf16>) -> vector<8xi8>
+ return %trunc : vector<8xi8>
+}
+
+// The narrow formats of the two ops differ, so the pair is not a round trip.
+// CHECK-LABEL: @no_fold_narrow_format_mismatch
+// CHECK: %[[EXT:.*]] = xevm.extf %{{.*}} {src_etype = f8, dst_etype = f16}
+// CHECK: %[[TRUNC:.*]] = xevm.truncf %[[EXT]] {src_etype = f16, dst_etype = bf8}
+// CHECK: return %[[TRUNC]]
+func.func @no_fold_narrow_format_mismatch(%arg0: vector<8xi8>) -> vector<8xi8> {
+ %ext = xevm.extf %arg0 { src_etype=f8, dst_etype=f16 } : (vector<8xi8>) -> vector<8xf16>
+ %trunc = xevm.truncf %ext { src_etype=f16, dst_etype=bf8 } : (vector<8xf16>) -> vector<8xi8>
+ return %trunc : vector<8xi8>
+}
+
+// The wide formats of the two ops differ: the value is extended as f16 and
+// truncated as if it were bf16.
+// CHECK-LABEL: @no_fold_wide_format_mismatch
+// CHECK: %[[EXT:.*]] = xevm.extf %{{.*}} {src_etype = f8, dst_etype = f16}
+// CHECK: %[[TRUNC:.*]] = xevm.truncf %[[EXT]] {src_etype = bf16, dst_etype = f8}
+// CHECK: return %[[TRUNC]]
+func.func @no_fold_wide_format_mismatch(%arg0: vector<8xi8>) -> vector<8xi8> {
+ %ext = xevm.extf %arg0 { src_etype=f8, dst_etype=f16 } : (vector<8xi8>) -> vector<8xf16>
+ %trunc = xevm.truncf %ext { src_etype=bf16, dst_etype=f8 } : (vector<8xf16>) -> vector<8xi8>
+ return %trunc : vector<8xi8>
+}
+
+// The other order is lossy: the truncation to f8 cannot be undone.
+// CHECK-LABEL: @no_fold_lossy_round_trip
+// CHECK: %[[TRUNC:.*]] = xevm.truncf %{{.*}} {src_etype = f16, dst_etype = f8}
+// CHECK: %[[EXT:.*]] = xevm.extf %[[TRUNC]] {src_etype = f8, dst_etype = f16}
+// CHECK: return %[[EXT]]
+func.func @no_fold_lossy_round_trip(%arg0: vector<8xf16>) -> vector<8xf16> {
+ %trunc = xevm.truncf %arg0 { src_etype=f16, dst_etype=f8 } : (vector<8xf16>) -> vector<8xi8>
+ %ext = xevm.extf %trunc { src_etype=f8, dst_etype=f16 } : (vector<8xi8>) -> vector<8xf16>
+ return %ext : vector<8xf16>
+}
More information about the Mlir-commits
mailing list