xref: /llvm-project/compiler-rt/test/nsan/softmax.cpp (revision 52ae891036e3ab1f668eb103c46ca57257901c6b)
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
19 template <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))
33 template <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 
46 int 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