#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_