Repository: pytorch
Issue: Incorrect exponential calculation on Jetson devices with float32 dtype
Labels: module: cuda,triaged,module: arm,module: jetson

## 🐛 Bug

The `exp` function in torch and libtorch (https://pytorch.org/docs/stable/generated/torch.exp.html#torch.exp) give incorrect output when run on Jetson devices (tested on Xavier NX and Nano) with `dtype=torch.float32`.

## To Reproduce
test.cpp:
```
#include <iostream>
#include <torch/script.h> // One-stop header.

int main(int argc, char *argv[])
{
        c10::InferenceMode guard(true); // Note: gives the same result without this.
        auto t = torch::ones({3, 3}, torch::dtype(torch::kFloat32));
        std::cout << "t:\n" << t << "\n";
        std::cout << "t.exp():\n" << t.exp() << "\n";
        return 0;
}
```

CMakeLists.txt:
```
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(test)

set(CMAKE_CXX_STANDARD 14)

find_package(Torch REQUIRED
        HINTS "/usr/local/lib/python3.8/dist-packages/torch")

add_executable(test test.cpp)
target_link_libraries(test PUBLIC "${TORCH_LIBRARIES}")
```

Compile and run:
```
br@NX:~/tmp/test$ Torch_DIR=/usr/local/lib/python3.6/dist-packages/torch cmake -GNinja .
-- Looking for pthread.h
-- Looking for pthread.h - found
-- Performing Test CMAKE_HAVE_LIBC_PTHREAD
-- Performing Test CMAKE_HAVE_LIBC_PTHREAD - Failed
-- Looking for pthread_create in pthreads
-- Looking for pthread_create in pthreads - not found
-- Looking for pthread_create in pthread
-- Looking for pthread_create in pthread - found
-- Found Threads: TRUE  
-- Found CUDA: /usr/local/cuda (found version "10.2") 
-- Caffe2: CUDA detected: 10.2
-- Caffe2: CUDA nvcc is: /usr/local/cuda/bin/nvcc
-- Caffe2: CUDA toolkit directory: /usr/local/cuda
-- Caffe2: Header version is: 10.2
-- Found CUDNN: /usr/lib/aarch64-linux-gnu/libcudnn.so  
-- Found cuDNN: v8.0.0  (include: /usr/include, library: /usr/lib/aarch64-linux-gnu/libcudnn.so)
-- /usr/local/cuda/lib64/libnvrtc.so shorthash is c13e41e1
-- Autodetected CUDA architecture(s):  7.2
-- Added CUDA NVCC flags for: -gencode;arch=compute_72,code=sm_72
CMake Warning at /usr/local/lib/python3.6/dist-packages/torch/share/cmake/Torch/TorchConfig.cmake:22 (message):
  static library kineto_LIBRARY-NOTFOUND not found.
Call Stack (most recent call first):
  /usr/local/lib/python3.6/dist-packages/torch/share/cmake/Torch/TorchConfig.cmake:127 (append_torchlib_if_found)
  CMakeLists.txt:6 (find_package)


-- Found Torch: /usr/local/lib/python3.6/dist-packages/torch/lib/libtorch.so  
-- Configuring done
-- Generating done
-- Build files have been written to: /home/br/tmp/test
br@NX:~/tmp/test$ ninja && ./test
[2/2] Linking CXX executable test
t:
 1  1  1
 1  1  1
 1  1  1
[ CPUFloatType{3,3} ]
t.exp():
 2.7183  1.0000  1.0000
 1.0000  9.4901  9.4901
 9.4901  9.4901  2.7183
[ CPUFloatType{3,3} ]
```

In python:
```
br@NX:~/tmp/test$ ipython
Python 3.6.9 (default, Jan 26 2021, 15:33:00) 
Type 'copyright', 'credits' or 'license' for more information
IPython 7.16.1 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import torch

In [2]: t=torch.ones((3,3), dtype=torch.float64); t.exp()
Out[2]: 
tensor([[2.7183, 2.7183, 2.7183],
        [2.7183, 2.7183, 2.7183],
        [2.7183, 2.7183, 2.7183]], dtype=torch.float64)

In [3]: t=torch.ones((3,3), dtype=torch.float32); t.exp()
Out[3]: 
tensor([[2.7183, 1.0000, 1.0000],
        [1.0000, 0.0000, 1.0000],
        [1.0000, 1.0000, 2.7183]])
```

## Expected behavior
Similar values as when calculated with float64.

## Environment

```
$ python3 collect_env.py 
Collecting environment information...
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (aarch64)
GCC version: (Ubuntu/Linaro 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.21.0-rc1
Libc version: glibc-2.25

Python version: 3.6.9 (default, Jan 26 2021, 15:33:00)  [GCC 8.4.0] (64-bit runtime)
Python platform: Linux-4.9.201-tegra-aarch64-with-Ubuntu-18.04-bionic
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/aarch64-linux-gnu/libcudnn.so.8.0.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv_infer.so.8.0.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv_train.so.8.0.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn_infer.so.8.0.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn_train.so.8.0.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops_infer.so.8.0.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops_train.so.8.0.0
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.9.0
[pip3] torchvision==0.10.0
[conda] Could not collect
```

## Additional Context
I've noticed this behaviour on older versions of libtorch due to getting strange results in a custom NMS function, but have only now figured out the source of those strange results, so I'm not sure how old this error is.

cc @malfet @ngimel