Build AllReduce only mode (#443)

* Initial commit of all_reduce_only support

* Working AllReduce only build

* Removing printfs and restoring release build

* Restore P2P index

* Updates to build_allreduce_only mode.

* cleaning up macro ifdefs

[ROCm/rccl commit: 7e55b211c5]

projects/rccl/CMakeLists.txt

@@ -47,6 +47,7 @@ set(AMDGPU_TARGETS "${DEFAULT_AMDGPU_TARGETS}" CACHE STRING "List of specific ma
 option(BUILD_TESTS "Build test programs" OFF)
 option(INSTALL_DEPENDENCIES "Force install dependencies" OFF)
 option(BUILD_ADDRESS_SANITIZER "Build with address sanitizer enabled" OFF)
+option(BUILD_ALLREDUCE_ONLY "Build AllReduce + sum + float kernel only" OFF)
 
 # parse version from Makefile NCCL_MAJOR, NCCL_MINOR, NCCL_PATCH must exist
 # NCCL_SUFFIX is optional NCCL_VERSION formatting is ((X) * 1000 + (Y) * 100 +
@@ -203,6 +204,10 @@ if(COLLTRACE)
   add_definitions(-DENABLE_COLLTRACE)
 endif()
 
+if (BUILD_ALLREDUCE_ONLY)
+  add_definitions(-DBUILD_ALLREDUCE_ONLY)
+endif()
+
 foreach(target ${AMDGPU_TARGETS})
   target_link_libraries(rccl PRIVATE --amdgpu-target=${target})
 endforeach()

projects/rccl/install.sh

@@ -19,6 +19,7 @@ function display_help()
     echo "    [--hcc] Build library using deprecated hcc compiler (default:hip-clang)."
     echo "    [--prefix] Specify custom directory to install RCCL to (default: /opt/rocm)."
     echo "    [--address-sanitizer] Build with address sanitizer enabled"
+    echo "    [--build_allreduce_only] Build only AllReduce + sum + float kernel"
 }
 
 # #################################################
@@ -36,6 +37,7 @@ build_hip_clang=true
 clean_build=true
 install_dependencies=false
 build_static=false
+build_allreduce_only=false
 
 # #################################################
 # Parameter parsing
@@ -44,7 +46,7 @@ build_static=false
 # check if we have a modern version of getopt that can handle whitespace and long parameters
 getopt -T
 if [[ $? -eq 4 ]]; then
-    GETOPT_PARSE=$(getopt --name "${0}" --longoptions help,install,dependencies,package_build,tests_build,run_tests_quick,static,run_tests_all,hcc,hip-clang,no_clean,prefix:,address-sanitizer --options hidptrs -- "$@")
+    GETOPT_PARSE=$(getopt --name "${0}" --longoptions help,install,dependencies,package_build,tests_build,run_tests_quick,static,run_tests_all,hcc,hip-clang,no_clean,prefix:,address-sanitizer,build_allreduce_only --options hidptrs -- "$@")
 else
     echo "Need a new version of getopt"
     exit 1
@@ -97,6 +99,9 @@ while true; do
     --address-sanitizer)
         build_address_sanitizer=true
         shift ;;
+    --build_allreduce_only)
+        build_allreduce_only=true
+        shift ;;
     --prefix)
         install_prefix=${2}
         shift 2 ;;
@@ -194,6 +199,9 @@ if ($install_dependencies); then
     cmake_common_options="${cmake_common_options} -DINSTALL_DEPENDENCIES=ON"
 fi
 
+if ($build_allreduce_only); then
+    cmake_common_options="${cmake_common_options} -DBUILD_ALLREDUCE_ONLY=ON"
+fi
 check_exit_code "$?"
 
 if ($build_tests) || (($run_tests) && [[ ! -f ./test/UnitTests ]]); then

projects/rccl/src/collectives/device/common.h

@@ -112,6 +112,9 @@ static const __device__ constexpr ncclKernelFunc_t ncclFuncs[]{
 // variable. There is no host pointer to a device-side function, which
 // confuses clang. This will be fixed in the next clang release.
 #if defined(__HIP_DEVICE_COMPILE__)
+#if defined(BUILD_ALLREDUCE_ONLY)
+  NCCL_FUNC_NAME(AllReduce, RING, SIMPLE, Sum, float)
+#else
   NCCL_FUNCS2B(Broadcast),
   NCCL_FUNCS2A(Reduce),
   NCCL_FUNCS2B(AllGather),
@@ -119,6 +122,7 @@ static const __device__ constexpr ncclKernelFunc_t ncclFuncs[]{
   NCCL_FUNCS2C(AllReduce),
   NCCL_FUNC_NAME(SendRecv, RING, SIMPLE, Sum, int8_t),
 #endif
+#endif
 };
 
 template<unsigned short f, unsigned short l>
@@ -143,6 +147,10 @@ static_assert(FUNC_INDEX_P2P == 2250, "Wrong P2P function index");
 inline
 __device__
 void NCCL_CALL_FUNCTIONS(struct ncclWorkElem* const c) noexcept {
+#if defined(BUILD_ALLREDUCE_ONLY)
+  assert(c->funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE));
+  ncclFunction_AllReduce_RING_SIMPLE_Sum_float(c);
+#else
   if (c->funcIndex < 450) {
     if (c->funcIndex % 9 == 0) ncclFunction_Broadcast_TREE_LL_Sum_int8_t(c);
     else if (c->funcIndex % 9 == 1) ncclFunction_Broadcast_TREE_LL_Sum_int8_t(c);
@@ -168,6 +176,7 @@ void NCCL_CALL_FUNCTIONS(struct ncclWorkElem* const c) noexcept {
   }
   else if (c->funcIndex < 2250) Caller<1350, 2250>::call(c);
   else ncclFunction_SendRecv_RING_SIMPLE_Sum_int8_t(c);
+#endif
 }
 
 template <ncclFunc_t FUNCTION, int ALGO, int PROTO, class REDOP, typename T, int UNROLL>

projects/rccl/test/CMakeLists.txt

@@ -12,38 +12,51 @@ if(BUILD_TESTS)
 
   include_directories(${GTEST_INCLUDE_DIRS})
 
-  # Collect source files for tests
-  set(TEST_SOURCES_SINGLE_PROCESS
-    test_AllGather.cpp
-    test_AllReduce.cpp
-    test_AllReduceGroup.cpp
-    test_Broadcast.cpp
-    test_Reduce.cpp
-    test_ReduceScatter.cpp
-    test_GroupCalls.cpp
-    test_CombinedCalls.cpp
-    test_AllReduceAbort.cpp
-    test_BroadcastAbort.cpp
-    test_Scatter.cpp
-    test_Gather.cpp
-    test_AllToAll.cpp
-    test_AllToAllv.cpp
-  )
-
-  set(TEST_SOURCES_MULTI_PROCESS
-    test_AllGatherMultiProcess.cpp
-    test_AllReduceMultiProcess.cpp
-    test_AllReduceGroupMultiProcess.cpp
-    test_AllToAllMultiProcess.cpp
-    test_BroadcastMultiProcess.cpp
-    test_CombinedCallsMultiProcess.cpp
-    test_GatherMultiProcess.cpp
-    test_GroupCallsMultiProcess.cpp
-    test_ReduceMultiProcess.cpp
-    test_ReduceScatterMultiProcess.cpp
-    test_ScatterMultiProcess.cpp
-  )
-
+  if(BUILD_ALLREDUCE_ONLY)
+    set(TEST_SOURCES_SINGLE_PROCESS
+      test_AllReduce.cpp
+      test_AllReduceAbort.cpp
+      test_AllReduceGroup.cpp
+    )
+  else()
+    # Collect source files for tests
+    set(TEST_SOURCES_SINGLE_PROCESS
+      test_AllGather.cpp
+      test_AllReduce.cpp
+      test_AllReduceGroup.cpp
+      test_Broadcast.cpp
+      test_Reduce.cpp
+      test_ReduceScatter.cpp
+      test_GroupCalls.cpp
+      test_CombinedCalls.cpp
+      test_AllReduceAbort.cpp
+      test_BroadcastAbort.cpp
+      test_Scatter.cpp
+      test_Gather.cpp
+      test_AllToAll.cpp
+      test_AllToAllv.cpp
+    )
+  endif()
+  if(BUILD_ALLREDUCE_ONLY)
+    set(TEST_SOURCES_MULTI_PROCESS
+      test_AllReduceMultiProcess.cpp
+      test_AllReduceGroupMultiProcess.cpp
+    )
+  else()
+    set(TEST_SOURCES_MULTI_PROCESS
+      test_AllGatherMultiProcess.cpp
+      test_AllReduceMultiProcess.cpp
+      test_AllReduceGroupMultiProcess.cpp
+      test_AllToAllMultiProcess.cpp
+      test_BroadcastMultiProcess.cpp
+      test_CombinedCallsMultiProcess.cpp
+      test_GatherMultiProcess.cpp
+      test_GroupCallsMultiProcess.cpp
+      test_ReduceMultiProcess.cpp
+      test_ReduceScatterMultiProcess.cpp
+      test_ScatterMultiProcess.cpp
+    )
+  endif()
   add_executable(UnitTests ${TEST_SOURCES_SINGLE_PROCESS})
   target_include_directories(UnitTests PRIVATE ${ROCM_PATH} ${GTEST_INCLUDE_DIRS})
   target_link_libraries(UnitTests PRIVATE ${GTEST_BOTH_LIBRARIES})

projects/rccl/test/test_AllReduce.cpp

@@ -35,7 +35,23 @@ namespace CorrectnessTests
 
     dataset.Release();
   }
-
+#if defined(BUILD_ALLREDUCE_ONLY)
+  INSTANTIATE_TEST_SUITE_P(AllReduceCorrectnessSweep,
+                           AllReduceCorrectnessTest,
+                           testing::Combine(
+                             // Reduction operator
+                             testing::Values(ncclSum),
+                             // Data types
+                             testing::Values(ncclFloat32),
+                             // Number of elements
+                             testing::Values(1024, 1048576),
+                             // Number of devices
+                             testing::Values(2,3,4,5,6,7,8),
+                             // In-place or not
+                             testing::Values(false, true),
+                             testing::Values("RCCL_ENABLE_CLIQUE=0", "RCCL_ENABLE_CLIQUE=1")),
+                           CorrectnessTest::PrintToStringParamName());
+#else
   INSTANTIATE_TEST_SUITE_P(AllReduceCorrectnessSweep,
                            AllReduceCorrectnessTest,
                            testing::Combine(
@@ -60,4 +76,5 @@ namespace CorrectnessTests
                              testing::Values(false, true),
                              testing::Values("RCCL_ENABLE_CLIQUE=0", "RCCL_ENABLE_CLIQUE=1")),
                            CorrectnessTest::PrintToStringParamName());
+#endif
 } // namespace

projects/rccl/test/test_AllReduceGroup.cpp

@@ -47,7 +47,23 @@ namespace CorrectnessTests
         dataset2.Release();
         dataset3.Release();
   }
-
+#if defined(BUILD_ALLREDUCE_ONLY)
+  INSTANTIATE_TEST_SUITE_P(AllReduceGroupCorrectnessSweep,
+                           AllReduceGroupCorrectnessTest,
+                           testing::Combine(
+                             // Reduction operator
+                             testing::Values(ncclSum),
+                             // Data types
+                             testing::Values(ncclFloat32),
+                             // Number of elements
+                             testing::Values(1024, 1048576),
+                             // Number of devices
+                             testing::Values(2,3,4,5,6,7,8),
+                             // In-place or not
+                             testing::Values(false, true),
+                             testing::Values("RCCL_ENABLE_CLIQUE=0", "RCCL_ENABLE_CLIQUE=1")),
+                           CorrectnessTest::PrintToStringParamName());
+#else
   INSTANTIATE_TEST_SUITE_P(AllReduceGroupCorrectnessSweep,
                            AllReduceGroupCorrectnessTest,
                            testing::Combine(
@@ -63,4 +79,5 @@ namespace CorrectnessTests
                              testing::Values(false, true),
                              testing::Values("RCCL_ENABLE_CLIQUE=0", "RCCL_ENABLE_CLIQUE=1")),
                            CorrectnessTest::PrintToStringParamName());
+#endif
 } // namespace