coreylowman · coreylowman · Feb 6, 2023 · Feb 4, 2023
diff --git a/src/tensor_ops/max_to/max_to.cu b/src/tensor_ops/max_to/max_to.cu
@@ -10,15 +10,11 @@ __device__ __forceinline__ float atomicMaxf(float * addr, float value) {
     }
 }
 
-// atomicMax is not implemented for floats,
-// solution copied https://stackoverflow.com/questions/17399119/how-do-i-use-atomicmax-on-floating-point-values-in-cuda
-__device__ __forceinline__ float atomicMaxf(double * addr, double value) {
+__device__ __forceinline__ double atomicMaxf(double * addr, double value) {
     if (signbit(value)) {
-        int high = __double2hiint(value);
-        int low = __double2loint(value);
-        return __uint_as_double(atomicMin((unsigned int *)addr, __double_as_uint(value)));        
+        return __longlong_as_double(atomicMin((unsigned long long int *)addr, __double_as_longlong(value)));
     } else {
-        return __int_as_double(atomicMax((int *)addr, __double_as_int(value)));
+        return __longlong_as_double(atomicMax((long long int *)addr, __double_as_longlong(value)));
     }
 }
 

diff --git a/src/tensor_ops/min_to/min_to.cu b/src/tensor_ops/min_to/min_to.cu
@@ -10,13 +10,11 @@ __device__ __forceinline__ float atomicMinf(float * addr, float value) {
     }
 }
 
-// atomicMax is not implemented for double,
-// solution copied https://stackoverflow.com/questions/17399119/how-do-i-use-atomicmax-on-floating-point-values-in-cuda
 __device__ __forceinline__ double atomicMinf(double * addr, double value) {
     if (signbit(value)) {
-        return __uint_as_double(atomicMax((unsigned int *)addr, __double_as_uint(value)));
+        return __longlong_as_double(atomicMax((unsigned long long int *)addr, __double_as_longlong(value)));
     } else {
-        return __int_as_double(atomicMin((int *)addr, __double_as_int(value)));
+        return __longlong_as_double(atomicMin((long long int *)addr, __double_as_longlong(value)));
     }
 }
 

diff --git a/src/tensor_ops/utilities/cuda_utils.cuh b/src/tensor_ops/utilities/cuda_utils.cuh
@@ -4,6 +4,8 @@
 __device__ double atomicAdd(double* a, double b) { return b; }
 #endif
 
+#include "cuda_fp16.h"
+
 __device__ unsigned int get_strided_index(
     unsigned int idx,
     const size_t num_dims,