[libc-commits] [libc] [libc][math] Add test and fix atan2f crashing when flush-denorm-to-zero (FTZ) and denorm-as-zero (DAZ) modes are set. (PR #112828)

[via libc-commits]

https://github.com/lntue created https://github.com/llvm/llvm-project/pull/112828

None

>From eed7ec5bbead2f66c38f3d303c43177672b03341 Mon Sep 17 00:00:00 2001
From: Tue Ly <lntue.h at gmail.com>
Date: Fri, 18 Oct 2024 03:40:56 +0000
Subject: [PATCH] [libc][math] Add test and fix atan2f crashing when
 flush-denorm-to-zero (FTZ) and denorm-as-zero (DAZ) modes are set.

---
 libc/src/math/generic/atan2f.cpp         | 14 +++++----
 libc/test/UnitTest/FPMatcher.h           | 26 +++++++++++++++++
 libc/test/src/math/smoke/atan2f_test.cpp | 37 ++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/libc/src/math/generic/atan2f.cpp b/libc/src/math/generic/atan2f.cpp
index e4b297c00f012c..a2e5499809a340 100644
--- a/libc/src/math/generic/atan2f.cpp
+++ b/libc/src/math/generic/atan2f.cpp
@@ -246,12 +246,18 @@ LLVM_LIBC_FUNCTION(float, atan2f, (float y, float x)) {
   uint32_t y_abs = y_bits.uintval();
   uint32_t max_abs = x_abs > y_abs ? x_abs : y_abs;
   uint32_t min_abs = x_abs <= y_abs ? x_abs : y_abs;
+  float num_f = FPBits(min_abs).get_val();
+  float den_f = FPBits(max_abs).get_val();
+  double num_d = static_cast<double>(num_f);
+  double den_d = static_cast<double>(den_f);
 
-  if (LIBC_UNLIKELY(max_abs >= 0x7f80'0000U || min_abs == 0U)) {
+  if (LIBC_UNLIKELY(max_abs >= 0x7f80'0000U || num_d == 0.0)) {
     if (x_bits.is_nan() || y_bits.is_nan())
       return FPBits::quiet_nan().get_val();
-    size_t x_except = x_abs == 0 ? 0 : (x_abs == 0x7f80'0000 ? 2 : 1);
-    size_t y_except = y_abs == 0 ? 0 : (y_abs == 0x7f80'0000 ? 2 : 1);
+    double x_d = static_cast<double>(x);
+    double y_d = static_cast<double>(y);
+    size_t x_except = (x_d == 0.0) ? 0 : (x_abs == 0x7f80'0000 ? 2 : 1);
+    size_t y_except = (y_d == 0.0) ? 0 : (y_abs == 0x7f80'0000 ? 2 : 1);
 
     // Exceptional cases:
     //   EXCEPT[y_except][x_except][x_is_neg]
@@ -275,8 +281,6 @@ LLVM_LIBC_FUNCTION(float, atan2f, (float y, float x)) {
   bool recip = x_abs < y_abs;
   double final_sign = IS_NEG[(x_sign != y_sign) != recip];
   fputil::DoubleDouble const_term = CONST_ADJ[x_sign][y_sign][recip];
-  double num_d = static_cast<double>(FPBits(min_abs).get_val());
-  double den_d = static_cast<double>(FPBits(max_abs).get_val());
   double q_d = num_d / den_d;
 
   double k_d = fputil::nearest_integer(q_d * 0x1.0p4f);
diff --git a/libc/test/UnitTest/FPMatcher.h b/libc/test/UnitTest/FPMatcher.h
index 5220b1245bf3a5..b0f9495da4c068 100644
--- a/libc/test/UnitTest/FPMatcher.h
+++ b/libc/test/UnitTest/FPMatcher.h
@@ -15,6 +15,7 @@
 #include "src/__support/FPUtil/FPBits.h"
 #include "src/__support/FPUtil/fpbits_str.h"
 #include "src/__support/macros/config.h"
+#include "src/__support/macros/properties/architectures.h"
 #include "test/UnitTest/RoundingModeUtils.h"
 #include "test/UnitTest/StringUtils.h"
 #include "test/UnitTest/Test.h"
@@ -175,6 +176,31 @@ template <typename T> struct FPTest : public Test {
   };
 };
 
+// Add facility to test Flush-Denormal-To-Zero (FTZ) and Denormal-As-Zero (DAZ)
+// modes.
+// These tests to ensure that our implementations will not crash under these
+// modes.
+#if defined(LIBC_TARGET_ARCH_IS_X86_64) && __has_builtin(__builtin_ia32_stmxcsr)
+
+#define LIBC_TEST_FTZ_DAZ
+
+static constexpr unsigned FTZ = 0x8000; // Flush denormal to zero
+static constexpr unsigned DAZ = 0x0040; // Denormal as zero
+
+struct ModifyMXCSR {
+  ModifyMXCSR(unsigned flags) {
+    old_mxcsr = __builtin_ia32_stmxcsr();
+    __builtin_ia32_ldmxcsr(old_mxcsr | flags);
+  }
+
+  ~ModifyMXCSR() { __builtin_ia32_ldmxcsr(old_mxcsr); }
+
+private:
+  unsigned old_mxcsr;
+};
+
+#endif
+
 } // namespace testing
 } // namespace LIBC_NAMESPACE_DECL
 
diff --git a/libc/test/src/math/smoke/atan2f_test.cpp b/libc/test/src/math/smoke/atan2f_test.cpp
index 32a28cfdfeaa62..94ec18d8f6b143 100644
--- a/libc/test/src/math/smoke/atan2f_test.cpp
+++ b/libc/test/src/math/smoke/atan2f_test.cpp
@@ -58,3 +58,40 @@ TEST_F(LlvmLibcAtan2fTest, SpecialNumbers) {
   // EXPECT_FP_EXCEPTION(0);
   EXPECT_MATH_ERRNO(0);
 }
+
+#ifdef LIBC_TEST_FTZ_DAZ
+
+using namespace LIBC_NAMESPACE::testing;
+
+TEST_F(LlvmLibcAtan2fTest, FTZMode) {
+  ModifyMXCSR mxcsr(FTZ);
+
+  EXPECT_FP_EQ(0x1.921fb6p-1f,
+               LIBC_NAMESPACE::atan2f(min_denormal, min_denormal));
+  EXPECT_FP_EQ(0x1.000002p-23f,
+               LIBC_NAMESPACE::atan2f(min_denormal, max_denormal));
+  EXPECT_FP_EQ(0x1.921fb4p0f,
+               LIBC_NAMESPACE::atan2f(max_denormal, min_denormal));
+  EXPECT_FP_EQ(0x1.921fb6p-1f,
+               LIBC_NAMESPACE::atan2f(max_denormal, max_denormal));
+}
+
+TEST_F(LlvmLibcAtan2fTest, DAZMode) {
+  ModifyMXCSR mxcsr(DAZ);
+
+  EXPECT_FP_EQ(0.0f, LIBC_NAMESPACE::atan2f(min_denormal, min_denormal));
+  EXPECT_FP_EQ(0.0f, LIBC_NAMESPACE::atan2f(min_denormal, max_denormal));
+  EXPECT_FP_EQ(0.0f, LIBC_NAMESPACE::atan2f(max_denormal, min_denormal));
+  EXPECT_FP_EQ(0.0f, LIBC_NAMESPACE::atan2f(max_denormal, max_denormal));
+}
+
+TEST_F(LlvmLibcAtan2fTest, FTZDAZMode) {
+  ModifyMXCSR mxcsr(FTZ | DAZ);
+
+  EXPECT_FP_EQ(0.0f, LIBC_NAMESPACE::atan2f(min_denormal, min_denormal));
+  EXPECT_FP_EQ(0.0f, LIBC_NAMESPACE::atan2f(min_denormal, max_denormal));
+  EXPECT_FP_EQ(0.0f, LIBC_NAMESPACE::atan2f(max_denormal, min_denormal));
+  EXPECT_FP_EQ(0.0f, LIBC_NAMESPACE::atan2f(max_denormal, max_denormal));
+}
+
+#endif