[gloo] Enable using c10::Half for gloo cuda

Differential Revision: D75909352

Pull Request resolved: #449

Author: fduwjj

gloo/cuda.cu

@@ -283,7 +283,7 @@ static inline int cudaGetBlocks(const int N) {
 #define DELEGATE_SIMPLE_CUDA_BINARY_OPERATOR(T, Funcname, op)       \
   __global__ void _Kernel_##T##_##Funcname(                         \
       T* dst, const T* src, const int n) {                          \
-    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);    \
+    for (auto i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);   \
          i += blockDim.x * gridDim.x) {                             \
       dst[i] = dst[i] op src[i];                                    \
     }                                                               \
@@ -301,7 +301,7 @@ static inline int cudaGetBlocks(const int N) {
 #define DELEGATE_HALF_PRECISION_CUDA_BINARY_OPERATOR(Funcname, op)             \
   __global__ void _Kernel_half_##Funcname(                                     \
       half* dst, const half* src, const int n) {                               \
-    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);               \
+    for (auto i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);              \
          i += blockDim.x * gridDim.x) {                                        \
       float r = __half2float(dst[i]) op __half2float(src[i]);                  \
       dst[i] = __float2half(r);                                                \
@@ -337,7 +337,7 @@ DELEGATE_HALF_PRECISION_CUDA_BINARY_OPERATOR(cudaProduct, *);
 #define DELEGATE_SIMPLE_CUDA_BINARY_COMPARE(T, Funcname, op)        \
   __global__ void _Kernel_##T##_##Funcname(                         \
       T* dst, const T* src, const int n) {                          \
-    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);    \
+    for (auto i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);   \
          i += blockDim.x * gridDim.x) {                             \
       if (src[i] op dst[i]) {                                       \
         dst[i] = src[i];                                            \
@@ -357,7 +357,7 @@ DELEGATE_HALF_PRECISION_CUDA_BINARY_OPERATOR(cudaProduct, *);
 #define DELEGATE_HALF_PRECISION_CUDA_BINARY_COMPARE(Funcname, op)              \
   __global__ void _Kernel_half_##Funcname(                                     \
       half* dst, const half* src, const int n) {                               \
-    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);               \
+    for (auto i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);              \
          i += blockDim.x * gridDim.x) {                                        \
       if (__half2float(src[i]) op __half2float(dst[i])) {                      \
         dst[i] = src[i];                                                       \
@@ -398,6 +398,12 @@ DELEGATE_SIMPLE_CUDA_BINARY_OPERATOR(BFloat16, cudaSum, +);
 DELEGATE_SIMPLE_CUDA_BINARY_OPERATOR(BFloat16, cudaProduct, *);
 DELEGATE_SIMPLE_CUDA_BINARY_COMPARE(BFloat16, cudaMin, <);
 DELEGATE_SIMPLE_CUDA_BINARY_COMPARE(BFloat16, cudaMax, >);
+using Half = c10::Half;
+INSTANTIATE_COPY_ASYNC(Half);
+DELEGATE_SIMPLE_CUDA_BINARY_OPERATOR(Half, cudaSum, +);
+DELEGATE_SIMPLE_CUDA_BINARY_OPERATOR(Half, cudaProduct, *);
+DELEGATE_SIMPLE_CUDA_BINARY_COMPARE(Half, cudaMin, <);
+DELEGATE_SIMPLE_CUDA_BINARY_COMPARE(Half, cudaMax, >);
 #endif
 
 } // namespace gloo

gloo/cuda_allreduce_bcube.cc

@@ -516,6 +516,7 @@ INSTANTIATE_TEMPLATE(float16);
 
 #if GLOO_USE_TORCH_DTYPES
 INSTANTIATE_TEMPLATE(c10::BFloat16);
+INSTANTIATE_TEMPLATE(c10::Half);
 #endif
 
 } // namespace gloo

gloo/cuda_allreduce_halving_doubling.cc

@@ -659,6 +659,7 @@ INSTANTIATE_TEMPLATE(float16);
 
 #if GLOO_USE_TORCH_DTYPES
 INSTANTIATE_TEMPLATE(c10::BFloat16);
+INSTANTIATE_TEMPLATE(c10::Half);
 #endif
 
 } // namespace gloo

gloo/cuda_allreduce_local.cc

@@ -78,6 +78,7 @@ INSTANTIATE_TEMPLATE(float16);
 
 #if GLOO_USE_TORCH_DTYPES
 INSTANTIATE_TEMPLATE(c10::BFloat16);
+INSTANTIATE_TEMPLATE(c10::Half);
 #endif
 
 } // namespace gloo

gloo/cuda_allreduce_ring.cc

@@ -190,6 +190,7 @@ INSTANTIATE_TEMPLATE(float16);
 
 #if GLOO_USE_TORCH_DTYPES
 INSTANTIATE_TEMPLATE(c10::BFloat16);
+INSTANTIATE_TEMPLATE(c10::Half);
 #endif
 
 } // namespace gloo

gloo/cuda_allreduce_ring_chunked.cc

@@ -367,6 +367,7 @@ INSTANTIATE_TEMPLATE(float16);
 
 #if GLOO_USE_TORCH_DTYPES
 INSTANTIATE_TEMPLATE(c10::BFloat16);
+INSTANTIATE_TEMPLATE(c10::Half);
 #endif
 
 } // namespace gloo

gloo/cuda_broadcast_one_to_all.cc

@@ -199,6 +199,7 @@ INSTANTIATE_TEMPLATE(float16);
 
 #if GLOO_USE_TORCH_DTYPES
 INSTANTIATE_TEMPLATE(c10::BFloat16);
+INSTANTIATE_TEMPLATE(c10::Half);
 #endif
 
 } // namespace gloo

gloo/cuda_private.h

@@ -22,6 +22,7 @@
 
 #if GLOO_USE_TORCH_DTYPES
 #include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
 #endif
 
 namespace gloo {