gloo/cuda: use torch dtype bf16

Author: d4l3k

CMakeLists.txt

@@ -42,6 +42,7 @@ if(${USE_TCP_OPENSSL_LINK} AND ${USE_TCP_OPENSSL_LOAD})
 endif()
 option(USE_CUDA "Build with CUDA support" OFF)
 option(GLOO_USE_CUDA_TOOLKIT "Build CUDA with FindCUDATookit.cmake and enable_language(CUDA)" OFF)
+option(GLOO_USE_TORCH_DTYPES "Build CUDA kernels with pytorch dtypes" 0)
 
 if(MSVC)
   message(STATUS "MSVC detected")

gloo/CMakeLists.txt

@@ -171,9 +171,19 @@ target_link_libraries(gloo PRIVATE ${gloo_DEPENDENCY_LIBS})
 target_include_directories(gloo INTERFACE $<INSTALL_INTERFACE:include>)
 if(USE_CUDA)
   target_include_directories(gloo_cuda INTERFACE $<INSTALL_INTERFACE:include>)
+
+  message(STATUS "GLOO_USE_TORCH_DTYPES : ${GLOO_USE_TORCH_DTYPES} ${GLOO_TORCH_DIR}")
+  if(GLOO_USE_TORCH_DTYPES)
+    target_include_directories(gloo_cuda PRIVATE ${GLOO_TORCH_DIR})
+  endif()
 endif()
 if(USE_ROCM)
   target_include_directories(gloo_hip INTERFACE $<INSTALL_INTERFACE:include>)
+
+  message(STATUS "GLOO_USE_TORCH_DTYPES : ${GLOO_USE_TORCH_DTYPES} ${GLOO_TORCH_DIR}")
+  if(GLOO_USE_TORCH_DTYPES)
+    target_include_directories(gloo_hip PRIVATE ${GLOO_TORCH_DIR})
+  endif()
 endif()
 
 # Install if necessary.

gloo/config.h.in

@@ -37,3 +37,5 @@ static_assert(
 #cmakedefine01 GLOO_HAVE_TRANSPORT_TCP_TLS
 #cmakedefine01 GLOO_HAVE_TRANSPORT_IBVERBS
 #cmakedefine01 GLOO_HAVE_TRANSPORT_UV
+
+#cmakedefine01 GLOO_USE_TORCH_DTYPES

gloo/cuda.cu

@@ -391,4 +391,13 @@ DELEGATE_SIMPLE_CUDA_BINARY_COMPARE(double, cudaMax, >);
 DELEGATE_HALF_PRECISION_CUDA_BINARY_COMPARE(cudaMin, <);
 DELEGATE_HALF_PRECISION_CUDA_BINARY_COMPARE(cudaMax, >);
 
+#if GLOO_USE_TORCH_DTYPES
+using BFloat16 = c10::BFloat16;
+INSTANTIATE_COPY_ASYNC(BFloat16);
+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, >);
+#endif
+
 } // namespace gloo

gloo/cuda_allreduce_bcube.cc

@@ -514,4 +514,8 @@ INSTANTIATE_TEMPLATE(float);
 INSTANTIATE_TEMPLATE(double);
 INSTANTIATE_TEMPLATE(float16);
 
+#if GLOO_USE_TORCH_DTYPES
+INSTANTIATE_TEMPLATE(c10::BFloat16);
+#endif
+
 } // namespace gloo

gloo/cuda_allreduce_halving_doubling.cc

@@ -657,4 +657,8 @@ INSTANTIATE_TEMPLATE(float);
 INSTANTIATE_TEMPLATE(double);
 INSTANTIATE_TEMPLATE(float16);
 
+#if GLOO_USE_TORCH_DTYPES
+INSTANTIATE_TEMPLATE(c10::BFloat16);
+#endif
+
 } // namespace gloo

gloo/cuda_allreduce_local.cc

@@ -76,4 +76,8 @@ INSTANTIATE_TEMPLATE(float);
 INSTANTIATE_TEMPLATE(double);
 INSTANTIATE_TEMPLATE(float16);
 
+#if GLOO_USE_TORCH_DTYPES
+INSTANTIATE_TEMPLATE(c10::BFloat16);
+#endif
+
 } // namespace gloo

gloo/cuda_allreduce_ring.cc

@@ -188,4 +188,8 @@ INSTANTIATE_TEMPLATE(float);
 INSTANTIATE_TEMPLATE(double);
 INSTANTIATE_TEMPLATE(float16);
 
+#if GLOO_USE_TORCH_DTYPES
+INSTANTIATE_TEMPLATE(c10::BFloat16);
+#endif
+
 } // namespace gloo

gloo/cuda_allreduce_ring_chunked.cc

@@ -365,4 +365,8 @@ INSTANTIATE_TEMPLATE(float);
 INSTANTIATE_TEMPLATE(double);
 INSTANTIATE_TEMPLATE(float16);
 
+#if GLOO_USE_TORCH_DTYPES
+INSTANTIATE_TEMPLATE(c10::BFloat16);
+#endif
+
 } // namespace gloo

gloo/cuda_broadcast_one_to_all.cc

@@ -197,4 +197,8 @@ INSTANTIATE_TEMPLATE(float);
 INSTANTIATE_TEMPLATE(double);
 INSTANTIATE_TEMPLATE(float16);
 
+#if GLOO_USE_TORCH_DTYPES
+INSTANTIATE_TEMPLATE(c10::BFloat16);
+#endif
+
 } // namespace gloo

gloo/cuda_private.h

@@ -20,6 +20,10 @@
 #include "gloo/cuda.h"
 #include "gloo/transport/device.h"
 
+#if GLOO_USE_TORCH_DTYPES
+#include <c10/util/BFloat16.h>
+#endif
+
 namespace gloo {
 
 #define CUDA_CHECK(condition)       \