[Mlir-commits] [llvm] [mlir] [APFloat] Don't turn a Float8E8M0FNU NaN into an Inf on conversion (PR #214919)

[曾鈜寬 Tseng Hung Kuan]

https://github.com/Tim096 updated https://github.com/llvm/llvm-project/pull/214919

>From 3b1949997d91944d2258acb3a60364e2a0f5daf6 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <p76091014 at gs.ncku.edu.tw>
Date: Sat, 8 Aug 2026 10:47:00 +0800
Subject: [PATCH] [APFloat] Don't turn a Float8E8M0FNU NaN into an Inf on
 conversion

Float8E8M0FNU has a precision of 1, so it stores no significand bits and
its NaN carries no payload. Converting it to a format that has infinities
shifted that empty significand into the target, leaving the NaN exponent
with an all-zero significand -- which is the encoding of an infinity.

The APFloat object still reported fcNaN, but bitcastToAPInt() returned Inf
bits, so anything that stores the bit pattern (ConstantFP, MLIR's
FloatAttr) silently ended up with an infinity. In MLIR this showed up when
folding arith.extf on an f8E8M0FNU constant, where the OCP MXFP spec uses
NaN as the scale of an invalid block:

  %c = arith.constant 0xFF : f8E8M0FNU
  %0 = arith.extf %c : f8E8M0FNU to f32
  // folded to 0x7F800000 (+Inf) instead of a NaN

Create a new NaN in the target semantics when the source has no
significand. Verified exhaustively over all 256 Float8E8M0FNU values
against f16, bf16, f32, f64 and f128.
---
 llvm/lib/Support/APFloat.cpp              |  5 ++++
 llvm/unittests/ADT/APFloatTest.cpp        | 30 +++++++++++++++++++++++
 mlir/test/Dialect/Arith/canonicalize.mlir | 19 ++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/llvm/lib/Support/APFloat.cpp b/llvm/lib/Support/APFloat.cpp
index be255bd19964b..63007d88ed4d4 100644
--- a/llvm/lib/Support/APFloat.cpp
+++ b/llvm/lib/Support/APFloat.cpp
@@ -2557,6 +2557,11 @@ APFloat::opStatus IEEEFloat::convert(const fltSemantics &toSemantics,
         semantics->nanEncoding != fltNanEncoding::NegativeZero)
       makeNaN(false, false);
 
+    // If the source has no significand, there are no payload bits to carry
+    // over, and an all-zero significand would encode an Inf. Create a new NaN.
+    if (!APFloat::hasSignificand(fromSemantics))
+      makeNaN(false, sign);
+
     *losesInfo = lostFraction != lfExactlyZero || X86SpecialNan;
 
     // For x87 extended precision, we want to make a NaN, not a special NaN if
diff --git a/llvm/unittests/ADT/APFloatTest.cpp b/llvm/unittests/ADT/APFloatTest.cpp
index 55d39e7f0a414..4ef2814e1a7b3 100644
--- a/llvm/unittests/ADT/APFloatTest.cpp
+++ b/llvm/unittests/ADT/APFloatTest.cpp
@@ -2304,6 +2304,36 @@ TEST(APFloatTest, Float8E8M0FNUValues) {
   EXPECT_EQ(0x1.0p-127, test.convertToDouble());
 }
 
+// Test that a Float8E8M0FNU NaN, which has no payload bits, does not
+// become an Inf when converted to a format that has infinities.
+TEST(APFloatTest, Float8E8M0FNUNaNConvert) {
+  APFloat nan = APFloat(APFloat::Float8E8M0FNU(), "nan");
+  EXPECT_TRUE(nan.isNaN());
+  EXPECT_EQ(APInt(8, 0xff), nan.bitcastToAPInt());
+
+  const std::pair<const fltSemantics *, uint64_t> ToSemantics[] = {
+      {&APFloat::IEEEhalf(), 0x7e00},
+      {&APFloat::BFloat(), 0x7fc0},
+      {&APFloat::IEEEsingle(), 0x7fc00000},
+      {&APFloat::IEEEdouble(), 0x7ff8000000000000ULL},
+  };
+
+  for (const auto &[Sem, ExpectedBits] : ToSemantics) {
+    APFloat test = nan;
+    bool losesInfo = true;
+    APFloat::opStatus status =
+        test.convert(*Sem, APFloat::rmNearestTiesToEven, &losesInfo);
+    EXPECT_EQ(status, APFloat::opOK);
+    EXPECT_FALSE(losesInfo);
+    EXPECT_TRUE(test.isNaN());
+    EXPECT_FALSE(test.isInfinity());
+    APInt bits = test.bitcastToAPInt();
+    EXPECT_EQ(APInt(APFloat::getSizeInBits(*Sem), ExpectedBits), bits);
+    // The stored bit pattern, not just the category, has to say NaN.
+    EXPECT_TRUE(APFloat(*Sem, bits).isNaN());
+  }
+}
+
 TEST(APFloatTest, getLargest) {
   EXPECT_EQ(3.402823466e+38f, APFloat::getLargest(APFloat::IEEEsingle()).convertToFloat());
   EXPECT_EQ(1.7976931348623158e+308, APFloat::getLargest(APFloat::IEEEdouble()).convertToDouble());
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 0c06aa6e861a4..89632305fa99d 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -1120,6 +1120,25 @@ func.func @extFPVectorConstant() -> vector<2xf128> {
   return %0 : vector<2xf128>
 }
 
+// A f8E8M0FNU NaN has no payload bits; folding must not turn it into an Inf.
+// CHECK-LABEL: @extFPConstantE8M0NaN
+//       CHECK:   %[[cres:.+]] = arith.constant 0x7FC00000 : f32
+//       CHECK:   return %[[cres]]
+func.func @extFPConstantE8M0NaN() -> f32 {
+  %cst = arith.constant 0xFF : f8E8M0FNU
+  %0 = arith.extf %cst : f8E8M0FNU to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: @extFPVectorConstantE8M0NaN
+//       CHECK:   %[[cres:.+]] = arith.constant dense<[1.000000e+00, 0x7FC00000]> : vector<2xf32>
+//       CHECK:   return %[[cres]]
+func.func @extFPVectorConstantE8M0NaN() -> vector<2xf32> {
+  %cst = arith.constant dense<[1.000000e+00, 0xFF]> : vector<2xf8E8M0FNU>
+  %0 = arith.extf %cst : vector<2xf8E8M0FNU> to vector<2xf32>
+  return %0 : vector<2xf32>
+}
+
 // CHECK-LABEL: @truncExtf
 //       CHECK-NOT:  truncf
 //       CHECK:   return  %arg0