This PR adds support to __half and nb_bfloat16 to myAtomicReduce (#1585)

Authors:
  - Nicolas Blin (https://github.com/Kh4ster)

Approvers:
  - Louis Sugy (https://github.com/Nyrio)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #1585

cpp/include/raft/util/cuda_utils.cuh

@@ -20,6 +20,11 @@
 #include <stdint.h>
 #include <type_traits>
 
+#if defined(_RAFT_HAS_CUDA)
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#endif
+
 #include <raft/core/cudart_utils.hpp>
 #include <raft/core/math.hpp>
 #include <raft/core/operators.hpp>
@@ -79,6 +84,35 @@ DI void myAtomicReduce(float* address, float val, ReduceLambda op)
   } while (assumed != old);
 }
 
+// Needed for atomicCas on ushort
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)
+template <typename ReduceLambda>
+DI void myAtomicReduce(__half* address, __half val, ReduceLambda op)
+{
+  unsigned short int* address_as_uint = (unsigned short int*)address;
+  unsigned short int old              = *address_as_uint, assumed;
+  do {
+    assumed = old;
+    old = atomicCAS(address_as_uint, assumed, __half_as_ushort(op(val, __ushort_as_half(assumed))));
+  } while (assumed != old);
+}
+#endif
+
+// Needed for nv_bfloat16 support
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+template <typename ReduceLambda>
+DI void myAtomicReduce(nv_bfloat16* address, nv_bfloat16 val, ReduceLambda op)
+{
+  unsigned short int* address_as_uint = (unsigned short int*)address;
+  unsigned short int old              = *address_as_uint, assumed;
+  do {
+    assumed = old;
+    old     = atomicCAS(
+      address_as_uint, assumed, __bfloat16_as_ushort(op(val, __ushort_as_bfloat16(assumed))));
+  } while (assumed != old);
+}
+#endif
+
 template <typename ReduceLambda>
 DI void myAtomicReduce(int* address, int val, ReduceLambda op)
 {