Reapply "[compiler-rt][nsan] Add support for nan detection" (#105909)

This reverts commit 1f89cd4.

Author: alexander-shaposhnikov

compiler-rt/lib/nsan/nsan.cpp

@@ -409,21 +409,21 @@ __nsan_dump_shadow_mem(const u8 *addr, size_t size_bytes, size_t bytes_per_line,
   }
 }
 
-alignas(16) SANITIZER_INTERFACE_ATTRIBUTE
+alignas(64) SANITIZER_INTERFACE_ATTRIBUTE
     thread_local uptr __nsan_shadow_ret_tag = 0;
 
-alignas(16) SANITIZER_INTERFACE_ATTRIBUTE
+alignas(64) SANITIZER_INTERFACE_ATTRIBUTE
     thread_local char __nsan_shadow_ret_ptr[kMaxVectorWidth *
                                             sizeof(__float128)];
 
-alignas(16) SANITIZER_INTERFACE_ATTRIBUTE
+alignas(64) SANITIZER_INTERFACE_ATTRIBUTE
     thread_local uptr __nsan_shadow_args_tag = 0;
 
 // Maximum number of args. This should be enough for anyone (tm). An alternate
 // scheme is to have the generated code create an alloca and make
 // __nsan_shadow_args_ptr point ot the alloca.
 constexpr const int kMaxNumArgs = 128;
-alignas(16) SANITIZER_INTERFACE_ATTRIBUTE
+alignas(64) SANITIZER_INTERFACE_ATTRIBUTE
     thread_local char __nsan_shadow_args_ptr[kMaxVectorWidth * kMaxNumArgs *
                                              sizeof(__float128)];
 
@@ -445,6 +445,32 @@ int32_t checkFT(const FT value, ShadowFT Shadow, CheckTypeT CheckType,
   const InternalFT check_value = value;
   const InternalFT check_shadow = Shadow;
 
+  // We only check for NaNs in the value, not the shadow.
+  if (flags().check_nan && isnan(check_value)) {
+    GET_CALLER_PC_BP;
+    BufferedStackTrace stack;
+    stack.Unwind(pc, bp, nullptr, false);
+    if (GetSuppressionForStack(&stack, CheckKind::Consistency)) {
+      // FIXME: optionally print.
+      return flags().resume_after_suppression ? kResumeFromValue
+                                              : kContinueWithShadow;
+    }
+    Decorator D;
+    Printf("%s", D.Warning());
+    Printf("WARNING: NumericalStabilitySanitizer: NaN detected\n");
+    Printf("%s", D.Default());
+    stack.Print();
+    if (flags().halt_on_error) {
+      if (common_flags()->abort_on_error)
+        Printf("ABORTING\n");
+      else
+        Printf("Exiting\n");
+      Die();
+    }
+    // Performing other tests for NaN values is meaningless when dealing with numbers.
+    return kResumeFromValue;
+  }
+
   // See this article for an interesting discussion of how to compare floats:
   // https://randomascii.wordpress.com/2012/02/25/comparing-floating-point-numbers-2012-edition/
   static constexpr const FT Eps = FTInfo<FT>::kEpsilon;

compiler-rt/lib/nsan/nsan_flags.inc

@@ -48,3 +48,5 @@ NSAN_FLAG(bool, enable_loadtracking_stats, false,
           "due to invalid or unknown types.")
 NSAN_FLAG(bool, poison_in_free, true, "")
 NSAN_FLAG(bool, print_stats_on_exit, false, "If true, print stats on exit.")
+NSAN_FLAG(bool, check_nan, false,
+          "If true, check the floating-point number is nan")

compiler-rt/test/nsan/nan.cpp

@@ -0,0 +1,25 @@
+// RUN: %clangxx_nsan -O0 -g %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=0 %run %t 2>&1 | FileCheck %s
+
+// RUN: %clangxx_nsan -O3 -g %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=0 %run %t 2>&1 | FileCheck %s
+
+// RUN: %clangxx_nsan -O0 -g %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=1 not %run %t
+
+#include <cmath>
+#include <cstdio>
+
+// This function returns a NaN value for triggering the NaN detection.
+__attribute__((noinline)) float ReturnNaN(float p, float q) {
+  float ret = p / q;
+  return ret;
+  // CHECK: WARNING: NumericalStabilitySanitizer: NaN detected
+}
+
+int main() {
+  float val = ReturnNaN(0., 0.);
+  printf("%f\n", val);
+  // CHECK: WARNING: NumericalStabilitySanitizer: NaN detected
+  return 0;
+}

compiler-rt/test/nsan/softmax.cpp

@@ -0,0 +1,54 @@
+// RUN: %clangxx_nsan -O0 -g -DSOFTMAX=softmax %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=0,log2_max_relative_error=19 %run %t 2>&1 | FileCheck %s
+
+// RUN: %clangxx_nsan -O3 -g -DSOFTMAX=softmax %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=0,log2_max_relative_error=19 %run %t 2>&1 | FileCheck %s
+
+// RUN: %clangxx_nsan -O0 -g -DSOFTMAX=stable_softmax %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=1,log2_max_relative_error=19 %run %t 
+
+// RUN: %clangxx_nsan -O3 -g -DSOFTMAX=stable_softmax %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=1,log2_max_relative_error=19 %run %t
+
+#include<iostream>
+#include<vector>
+#include<algorithm>
+#include<cmath>
+
+// unstable softmax
+template <typename T>
+__attribute__((noinline)) void softmax(std::vector<T> &values) {
+    T sum_exp = 0.0;
+    for (auto &i: values) {
+      i = std::exp(i);
+      sum_exp += i;
+    }
+    for (auto &i: values) {
+      i /= sum_exp;
+    }
+}
+
+// use max value to avoid overflow
+// \sigma_i exp(x_i) / \sum_j exp(x_j) = \sigma_i exp(x_i - max(x)) / \sum_j exp(x_j - max(x))
+template <typename T>
+__attribute__((noinline)) void stable_softmax(std::vector<T> &values) {
+  T sum_exp = 0.0;
+  T max_values = *std::max_element(values.begin(), values.end());
+  for (auto &i: values) {
+    i = std::exp(i - max_values);
+    sum_exp += i;
+  }
+  for (auto &i:values) {
+    i /= sum_exp;
+  }
+}
+
+int main() {
+  std::vector<double> data = {1000, 1001, 1002};
+  SOFTMAX(data);
+  for (auto i: data) {
+    printf("%f", i);
+    // CHECK: WARNING: NumericalStabilitySanitizer: NaN detected
+  }
+  return 0;
+}

compiler-rt/test/nsan/vec_sqrt.cpp

@@ -0,0 +1,34 @@
+// RUN: %clangxx_nsan -O0 -g -mavx %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=0 %run %t 2>&1 | FileCheck %s
+// RUN: %clangxx_nsan -O3 -g -mavx %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=0 %run %t 2>&1 | FileCheck %s
+
+#include <cmath>
+#include <immintrin.h>
+#include <iostream>
+
+void simd_sqrt(const float *input, float *output, size_t size) {
+  size_t i = 0;
+  for (; i + 7 < size; i += 8) {
+    __m256 vec = _mm256_loadu_ps(&input[i]);
+    __m256 result = _mm256_sqrt_ps(vec);
+    _mm256_storeu_ps(&output[i], result);
+  }
+  for (; i < size; ++i) {
+    output[i] = std::sqrt(input[i]);
+    // CHECK: WARNING: NumericalStabilitySanitizer: NaN detected
+  }
+}
+
+int main() {
+  float input[] = {1.0,  2.0,   -3.0,  4.0,   5.0,   6.0,  7.0,
+                   8.0,  9.0,   -10.0, 11.0,  12.0,  13.0, 14.0,
+                   15.0, -16.0, 17.0,  -18.0, -19.0, -20.0};
+  float output[20];
+  simd_sqrt(input, output, 20);
+  for (int i = 0; i < 20; ++i) {
+    std::cout << output[i] << std::endl;
+    // CHECK: WARNING: NumericalStabilitySanitizer: NaN detected
+  }
+  return 0;
+}

compiler-rt/test/nsan/vec_sqrt_ext.cpp

@@ -0,0 +1,25 @@
+// RUN: %clangxx_nsan -O0 -g -mavx %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=0 %run %t 2>&1 | FileCheck %s
+// RUN: %clangxx_nsan -O3 -g -mavx %s -o %t
+// RUN: NSAN_OPTIONS=check_nan=true,halt_on_error=0 %run %t 2>&1 | FileCheck %s
+#include <iostream>
+#include <cmath>
+
+typedef float v8sf __attribute__ ((vector_size(32)));
+
+v8sf simd_sqrt(v8sf a) {
+  return __builtin_elementwise_sqrt(a);
+  // CHECK: WARNING: NumericalStabilitySanitizer: NaN detected
+}
+
+int main() {
+  v8sf a = {-1.0, -2.0, -3.0, 4.0, 5.0, 6.0, 7.0, 8.0};
+  a = simd_sqrt(a);
+
+  // This prevents DCE.
+  for (size_t i = 0; i < 8; ++i) {
+    std::cout << a[i] << std::endl;
+    // CHECK: WARNING: NumericalStabilitySanitizer: NaN detected
+  }
+  return 0;
+}