1// RUN: %clangxx_nsan -O0 -g -DSOFTMAX=softmax %s -o %t
2// RUN: env NSAN_OPTIONS=check_nan=true,halt_on_error=0,log2_max_relative_error=19 %run %t 2>&1 | FileCheck %s
3
4// RUN: %clangxx_nsan -O3 -g -DSOFTMAX=softmax %s -o %t
5// RUN: env NSAN_OPTIONS=check_nan=true,halt_on_error=0,log2_max_relative_error=19 %run %t 2>&1 | FileCheck %s
6
7// RUN: %clangxx_nsan -O0 -g -DSOFTMAX=stable_softmax %s -o %t
8// RUN: env NSAN_OPTIONS=check_nan=true,halt_on_error=1,log2_max_relative_error=19 %run %t
9
10// RUN: %clangxx_nsan -O3 -g -DSOFTMAX=stable_softmax %s -o %t
11// RUN: env NSAN_OPTIONS=check_nan=true,halt_on_error=1,log2_max_relative_error=19 %run %t
12
13#include<iostream>
14#include<vector>
15#include<algorithm>
16#include<cmath>
17
18// unstable softmax
19template <typename T>
20__attribute__((noinline)) void softmax(std::vector<T> &values) {
21 T sum_exp = 0.0;
22 for (auto &i: values) {
23 i = std::exp(i);
24 sum_exp += i;
25 }
26 for (auto &i: values) {
27 i /= sum_exp;
28 }
29}
30
31// use max value to avoid overflow
32// \sigma_i exp(x_i) / \sum_j exp(x_j) = \sigma_i exp(x_i - max(x)) / \sum_j exp(x_j - max(x))
33template <typename T>
34__attribute__((noinline)) void stable_softmax(std::vector<T> &values) {
35 T sum_exp = 0.0;
36 T max_values = *std::max_element(values.begin(), values.end());
37 for (auto &i: values) {
38 i = std::exp(i - max_values);
39 sum_exp += i;
40 }
41 for (auto &i:values) {
42 i /= sum_exp;
43 }
44}
45
46int main() {
47 std::vector<double> data = {1000, 1001, 1002};
48 SOFTMAX(data);
49 for (auto i: data) {
50 printf("%f", i);
51 // CHECK: WARNING: NumericalStabilitySanitizer: NaN detected
52 }
53 return 0;
54}
55

source code of compiler-rt/test/nsan/softmax.cpp