[Mlir-commits] [mlir] [mlir][arith] Fix f8E5M2 extf/truncf expansion for bf16 (PR #230976)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 11 07:45:01 PDT 2026
https://github.com/SomnusGG created https://github.com/llvm/llvm-project/pull/230976
The f8E5M2 extf/truncf expansions use f16 as their intermediate type. For a
bf16 result (extf) or a bf16 source (truncf), they emitted
`arith.extf f16 to bf16` or `arith.truncf bf16 to f16`. These ops are invalid
because both types are 16 bits wide:
'arith.extf' op operand type 'f16' and result type 'bf16' are cast
incompatible
Reproducer (`mlir-opt -arith-expand="include-f8e5m2=true"`):
%0 = arith.extf %arg0 : f8E5M2 to bf16
%1 = arith.truncf %arg1 : bf16 to f8E5M2
`arith.scaling_extf` and `arith.scaling_truncf` hit the same problem when they
convert between bf16 and f8E5M2.
Go through f32 when the other type is 16 bits wide and is not f16:
* extf: f16 -> f32 -> bf16. Every f8E5M2 value is exactly representable in
bf16, so this does not round.
* truncf: bf16 -> f32 -> f16 -> f8E5M2. Converting bf16 to f16 rounds only
two kinds of value. Values below 2^-17 (half the smallest f8E5M2
subnormal) become 0 either way. Values above 65504 are past the f8E5M2
overflow threshold and become inf either way. So this does not
introduce double rounding.
This adds lit tests and execution tests in
test-arith-expand-truncf-extf.mlir. I also compared the expansion with an
exact round-to-nearest-even reference: all 256 f8E5M2 inputs for extf, and
about 2000 bf16 inputs for truncf. Every result matched.
Known issues found during this work but not addressed here:
* truncf from f32/f64 to f8E4M3FN/f8E5M2 rounds twice (once to f16, then to
the f8 type). Inputs very close to a tie between two f8 values can round
the wrong way.
* The f8E8M0FNU extf expansion returns +0 for the 0x00 encoding. It should
return 2^-127, which is an f32 subnormal.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
>From d8fdca53575c6e6e1045d69f4c86696f40d5a869 Mon Sep 17 00:00:00 2001
From: SomnusGG <a69254758 at gmail.com>
Date: Sun, 11 Oct 2026 22:14:49 +0800
Subject: [PATCH] [mlir][arith] Fix f8E5M2 extf/truncf expansion for bf16
The f8E5M2 extf/truncf expansions use f16 as their intermediate type. For a
bf16 result (extf) or a bf16 source (truncf), they emitted
`arith.extf f16 to bf16` or `arith.truncf bf16 to f16`. These ops are invalid
because both types are 16 bits wide:
'arith.extf' op operand type 'f16' and result type 'bf16' are cast
incompatible
Reproducer (`mlir-opt -arith-expand="include-f8e5m2=true"`):
%0 = arith.extf %arg0 : f8E5M2 to bf16
%1 = arith.truncf %arg1 : bf16 to f8E5M2
`arith.scaling_extf` and `arith.scaling_truncf` hit the same problem when they
convert between bf16 and f8E5M2.
Go through f32 when the other type is 16 bits wide and is not f16:
* extf: f16 -> f32 -> bf16. Every f8E5M2 value is exactly representable in
bf16, so this does not round.
* truncf: bf16 -> f32 -> f16 -> f8E5M2. Converting bf16 to f16 rounds only
two kinds of value. Values below 2^-17 (half the smallest f8E5M2
subnormal) become 0 either way. Values above 65504 are past the f8E5M2
overflow threshold and become inf either way. So this does not
introduce double rounding.
This adds lit tests and execution tests in
test-arith-expand-truncf-extf.mlir. I also compared the expansion with an
exact round-to-nearest-even reference: all 256 f8E5M2 inputs for extf, and
about 2000 bf16 inputs for truncf. Every result matched.
Known issues found during this work but not addressed here:
* truncf from f32/f64 to f8E4M3FN/f8E5M2 rounds twice (once to f16, then to
the f8 type). Inputs very close to a tie between two f8 values can round
the wrong way.
* The f8E8M0FNU extf expansion returns +0 for the 0x00 encoding. It should
return 2^-127, which is an f32 subnormal.
Co-Authored-By: Claude Opus 5.5 <noreply at anthropic.com>
---
.../Dialect/Arith/Transforms/ExpandOps.cpp | 25 ++++++--
mlir/test/Dialect/Arith/expand-ops.mlir | 57 ++++++++++++++++++
.../CPU/test-arith-expand-truncf-extf.mlir | 58 ++++++++++++++++++-
3 files changed, 135 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index 2e02d44bdbe2c..432f8b20f8bad 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -695,11 +695,17 @@ struct F8E5M2ExtFOpConverter : public OpRewritePattern<arith::ExtFOp> {
Value result = f16;
if (!resultETy.isF16()) {
- if (resultETy.getIntOrFloatBitWidth() < 16)
- result = arith::TruncFOp::create(b, resultTy, f16, nullptr,
+ if (resultETy.getIntOrFloatBitWidth() <= 16) {
+ // There is no direct extf/truncf between F16 and another 16-bit type
+ // such as BF16, so go through F32. Every F8E5M2 value is exactly
+ // representable in F32 and in BF16, so this does not round.
+ Type f32Ty = cloneToShapedType(operandTy, b.getF32Type());
+ Value f32 = arith::ExtFOp::create(b, f32Ty, f16, op.getFastmathAttr());
+ result = arith::TruncFOp::create(b, resultTy, f32, nullptr,
op.getFastmathAttr());
- else
+ } else {
result = arith::ExtFOp::create(b, resultTy, f16, op.getFastmathAttr());
+ }
}
rewriter.replaceOp(op, result);
return success();
@@ -735,9 +741,20 @@ struct F8E5M2TruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
Type f16Ty = cloneToShapedType(operandTy, b.getF16Type());
Value h16 = operand;
- if (!operandETy.isF16())
+ if (!operandETy.isF16()) {
+ // There is no direct truncf from another 16-bit type such as BF16 to
+ // F16, so go through F32. BF16 to F16 only rounds values that are below
+ // half the smallest F8E5M2 subnormal (they become 0 either way) or above
+ // the F8E5M2 overflow threshold (they become infinity either way), so
+ // this does not introduce double rounding.
+ if (operandETy.getIntOrFloatBitWidth() <= 16) {
+ Type f32Ty = cloneToShapedType(operandTy, b.getF32Type());
+ operand =
+ arith::ExtFOp::create(b, f32Ty, operand, op.getFastmathAttr());
+ }
h16 = arith::TruncFOp::create(b, f16Ty, operand, nullptr,
op.getFastmathAttr());
+ }
Value isNan = arith::CmpFOp::create(b, arith::CmpFPredicate::UNE, h16, h16);
Value h16Bits = arith::BitcastOp::create(b, i16Ty, h16);
diff --git a/mlir/test/Dialect/Arith/expand-ops.mlir b/mlir/test/Dialect/Arith/expand-ops.mlir
index 0342d6a7036cf..eb13dc51fb585 100644
--- a/mlir/test/Dialect/Arith/expand-ops.mlir
+++ b/mlir/test/Dialect/Arith/expand-ops.mlir
@@ -751,6 +751,63 @@ func.func @truncf_vector_f32_to_f8E5M2(%arg0 : vector<4xf32>) -> vector<4xf8E5M2
// -----
+// BF16 has the same width as the F16 pivot type, so there is no direct extf
+// between them: the expansion goes through F32. It used to emit an invalid
+// `arith.extf f16 to bf16`.
+func.func @extf_f8E5M2_to_bf16(%arg0 : f8E5M2) -> bf16 {
+ %0 = arith.extf %arg0 : f8E5M2 to bf16
+ return %0 : bf16
+}
+
+// CHECK-LABEL: @extf_f8E5M2_to_bf16
+// CHECK: %[[F16:.+]] = arith.bitcast %{{.+}} : i16 to f16
+// CHECK: arith.extf %[[F16]] : f16 to f32
+// CHECK: %[[RESULT:.+]] = arith.bitcast %{{.+}} : i16 to bf16
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @extf_vector_f8E5M2_to_bf16(%arg0 : vector<4xf8E5M2>) -> vector<4xbf16> {
+ %0 = arith.extf %arg0 : vector<4xf8E5M2> to vector<4xbf16>
+ return %0 : vector<4xbf16>
+}
+
+// CHECK-LABEL: @extf_vector_f8E5M2_to_bf16
+// CHECK: arith.extf %{{.+}} : vector<4xf16> to vector<4xf32>
+// CHECK: arith.bitcast %{{.+}} : vector<4xi16> to vector<4xbf16>
+// CHECK: return
+
+// -----
+
+// BF16 has the same width as the F16 pivot type, so there is no direct truncf
+// between them: the expansion goes through F32. It used to emit an invalid
+// `arith.truncf bf16 to f16`.
+func.func @truncf_bf16_to_f8E5M2(%arg0 : bf16) -> f8E5M2 {
+ %0 = arith.truncf %arg0 : bf16 to f8E5M2
+ return %0 : f8E5M2
+}
+
+// CHECK-LABEL: @truncf_bf16_to_f8E5M2
+// CHECK: %[[F32:.+]] = arith.bitcast %{{.+}} : i32 to f32
+// CHECK: %[[H16:.+]] = arith.truncf %[[F32]] : f32 to f16
+// CHECK: arith.cmpf une, %[[H16]], %[[H16]] : f16
+// CHECK: %[[RESULT:.+]] = arith.bitcast %{{.+}} : i8 to f8E5M2
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @truncf_vector_bf16_to_f8E5M2(%arg0 : vector<4xbf16>) -> vector<4xf8E5M2> {
+ %0 = arith.truncf %arg0 : vector<4xbf16> to vector<4xf8E5M2>
+ return %0 : vector<4xf8E5M2>
+}
+
+// CHECK-LABEL: @truncf_vector_bf16_to_f8E5M2
+// CHECK: arith.truncf %{{.+}} : vector<4xf32> to vector<4xf16>
+// CHECK: arith.bitcast %{{.+}} : vector<4xi8> to vector<4xf8E5M2>
+// CHECK: return
+
+// -----
+
func.func @extf_f8E4M3FN_to_f32(%arg0 : f8E4M3FN) -> f32 {
%0 = arith.extf %arg0 : f8E4M3FN to f32
return %0 : f32
diff --git a/mlir/test/Integration/Dialect/Arith/CPU/test-arith-expand-truncf-extf.mlir b/mlir/test/Integration/Dialect/Arith/CPU/test-arith-expand-truncf-extf.mlir
index fb4d789fb3cd4..e255b0be4e510 100644
--- a/mlir/test/Integration/Dialect/Arith/CPU/test-arith-expand-truncf-extf.mlir
+++ b/mlir/test/Integration/Dialect/Arith/CPU/test-arith-expand-truncf-extf.mlir
@@ -1,5 +1,5 @@
// Check various edge cases for truncf/extf ops involving f32 and the f4E2M1FN,
-// f8E4M3FN and f8E5M2 types.
+// f8E4M3FN and f8E5M2 types, and between bf16 and f8E5M2.
// RUN: mlir-opt %s --convert-func-to-llvm \
// RUN: --arith-expand="include-f4e2m1=true include-f8e4m3fn=true include-f8e5m2=true" \
@@ -58,6 +58,23 @@ func.func @check_truncf_f8E5M2(%in : f32) -> () {
return
}
+// BF16 has the same width as the F16 pivot type used by the F8E5M2 expansion.
+func.func @check_extf_f8E5M2_to_bf16(%in : f8E5M2) -> () {
+ %res = arith.extf %in : f8E5M2 to bf16
+ %bitcast = arith.bitcast %res : bf16 to i16
+ %ext = arith.extui %bitcast : i16 to i64
+ vector.print %ext : i64
+ return
+}
+
+func.func @check_truncf_bf16_to_f8E5M2(%in : bf16) -> () {
+ %trunc = arith.truncf %in : bf16 to f8E5M2
+ %bitcast = arith.bitcast %trunc : f8E5M2 to i8
+ %res = arith.extui %bitcast : i8 to i64
+ vector.print %res : i64
+ return
+}
+
func.func @entry() {
%zero = arith.constant 0.0 : f32
%half = arith.constant 0.5 : f32
@@ -318,5 +335,44 @@ func.func @entry() {
// CHECK: -inf
%lowerE5m2 = arith.truncf %lowerThanMin : f32 to f8E5M2
func.call @check_extf_f8E5M2(%lowerE5m2) : (f8E5M2) -> ()
+
+ // BF16 <-> F8E5M2 checks.
+ %bf16One = arith.constant 1.0 : bf16
+ %bf16E5m2Max = arith.constant 57344.0 : bf16
+ %bf16Ovf = arith.constant 70000.0 : bf16
+ %bf16Inf = arith.constant 0x7F80 : bf16
+ %bf16Nan = arith.constant 0x7FC0 : bf16
+ // 2^-16, the smallest F8E5M2 subnormal.
+ %bf16MinSub = arith.constant 0x3780 : bf16
+ // -2^-17, a tie between -0 and the smallest negative subnormal.
+ %bf16NegTie = arith.constant 0xB700 : bf16
+
+ // CHECK: 60
+ func.call @check_truncf_bf16_to_f8E5M2(%bf16One) : (bf16) -> ()
+ // CHECK: 123
+ func.call @check_truncf_bf16_to_f8E5M2(%bf16E5m2Max) : (bf16) -> ()
+ // CHECK: 124
+ func.call @check_truncf_bf16_to_f8E5M2(%bf16Ovf) : (bf16) -> ()
+ // CHECK: 124
+ func.call @check_truncf_bf16_to_f8E5M2(%bf16Inf) : (bf16) -> ()
+ // CHECK: 126
+ func.call @check_truncf_bf16_to_f8E5M2(%bf16Nan) : (bf16) -> ()
+ // CHECK: 1
+ func.call @check_truncf_bf16_to_f8E5M2(%bf16MinSub) : (bf16) -> ()
+ // CHECK: 128
+ func.call @check_truncf_bf16_to_f8E5M2(%bf16NegTie) : (bf16) -> ()
+
+ %e5m2OneAndAQuarter = arith.constant 0x3D : f8E5M2
+ %e5m2MaxBits = arith.constant 0x7B : f8E5M2
+ %e5m2NegInf = arith.constant 0xFC : f8E5M2
+ %e5m2MinSub = arith.constant 0x01 : f8E5M2
+ // CHECK: 16288
+ func.call @check_extf_f8E5M2_to_bf16(%e5m2OneAndAQuarter) : (f8E5M2) -> ()
+ // CHECK: 18272
+ func.call @check_extf_f8E5M2_to_bf16(%e5m2MaxBits) : (f8E5M2) -> ()
+ // CHECK: 65408
+ func.call @check_extf_f8E5M2_to_bf16(%e5m2NegInf) : (f8E5M2) -> ()
+ // CHECK: 14208
+ func.call @check_extf_f8E5M2_to_bf16(%e5m2MinSub) : (f8E5M2) -> ()
return
}
More information about the Mlir-commits
mailing list