1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
|
#ifndef CAFFE2_UTILS_GPU_ATOMICS_H_
#define CAFFE2_UTILS_GPU_ATOMICS_H_
#include <cuda_runtime.h>
namespace caffe2 {
namespace {
template <typename T>
inline __device__ void gpu_atomic_add(T* address, const T val) {
atomicAdd(address, val);
}
template <>
inline __device__ void gpu_atomic_add(float* address, const float val) {
#if defined(USE_ROCM) && defined(__gfx908__)
atomicAddNoRet(address, val);
#else
atomicAdd(address, val);
#endif
}
} // namespace
} // namespace caffe2
#endif // CAFFE2_UTILS_GPU_ATOMICS_H_
|
