[CMake] Add RCCL to TVM and TVM Runtime

CMakeLists.txt

@@ -11,6 +11,7 @@ include(cmake/utils/FindOpenCL.cmake)
 include(cmake/utils/FindVulkan.cmake)
 include(cmake/utils/FindLLVM.cmake)
 include(cmake/utils/FindROCM.cmake)
+include(cmake/utils/FindRCCL.cmake)
 include(cmake/utils/FindEthosN.cmake)
 
 if(EXISTS ${CMAKE_BINARY_DIR}/config.cmake)
@@ -43,6 +44,7 @@ tvm_option(USE_KHRONOS_SPIRV "Whether to use spirv-tools.and SPIRV-Headers from
 tvm_option(USE_SPIRV_KHR_INTEGER_DOT_PRODUCT "whether enable SPIRV_KHR_DOT_PRODUCT" OFF)
 tvm_option(USE_METAL "Build with Metal" OFF)
 tvm_option(USE_ROCM "Build with ROCM" OFF)
+tvm_option(USE_RCCL "Build with RCCL" OFF)
 tvm_option(ROCM_PATH "The path to rocm" /opt/rocm)
 tvm_option(USE_HEXAGON "Build with Hexagon support" OFF)
 tvm_option(USE_HEXAGON_SDK "Path to the Hexagon SDK root (required for Hexagon support)" /path/to/sdk)
@@ -431,6 +433,13 @@ if(USE_CUDA AND USE_NCCL)
   list(APPEND RUNTIME_SRCS ${RUNTIME_NCCL_SRC})
 endif()
 
+if(USE_ROCM AND USE_RCCL)
+  message(STATUS "Build with RCCL...")
+  find_rccl(${USE_RCCL})
+  tvm_file_glob(GLOB RUNTIME_RCCL_SRC src/runtime/disco/rccl/*.cc)
+  list(APPEND RUNTIME_SRCS ${RUNTIME_RCCL_SRC})
+endif()
+
 if(USE_AOT_EXECUTOR)
   message(STATUS "Build with AOT Executor support...")
   file(GLOB RUNTIME_AOT_EXECUTOR_SRCS src/runtime/aot_executor/*.cc)
@@ -847,3 +856,8 @@ if(USE_CUDA AND USE_NCCL)
   target_link_libraries(tvm PRIVATE nccl)
   target_link_libraries(tvm_runtime PRIVATE nccl)
 endif()
+
+if(USE_ROCM AND USE_RCCL)
+  target_link_libraries(tvm PRIVATE rccl)
+  target_link_libraries(tvm_runtime PRIVATE rccl)
+endif()

cmake/config.cmake

@@ -62,6 +62,12 @@ set(USE_NCCL OFF)
 # - /path/to/rocm: use specific path to rocm
 set(USE_ROCM OFF)
 
+# Whether to enable RCCL support:
+# - ON: enable RCCL with cmake's auto search
+# - OFF: disable RCCL
+# - /path/to/rccl: use specific path to rccl
+set(USE_RCCL OFF)
+
 # Whether enable SDAccel runtime
 set(USE_SDACCEL OFF)
 

cmake/modules/LibInfo.cmake

@@ -104,6 +104,7 @@ function(add_lib_info src_file)
     TVM_INFO_USE_RELAY_DEBUG="${USE_RELAY_DEBUG}"
     TVM_INFO_USE_ROCBLAS="${USE_ROCBLAS}"
     TVM_INFO_USE_ROCM="${USE_ROCM}"
+    TVM_INFO_USE_RCCL="${USE_RCCL}"
     TVM_INFO_USE_RPC="${USE_RPC}"
     TVM_INFO_USE_RTTI="${USE_RTTI}"
     TVM_INFO_USE_RUST_EXT="${USE_RUST_EXT}"

cmake/utils/FindRCCL.cmake

@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# Variables used by this module, they can change the default behaviour and need
+# to be set before calling find_package:
+#
+# This module defines
+#  Rccl_FOUND, whether rccl has been found
+#  RCCL_INCLUDE_DIR, directory containing header
+#  RCCL_LIBRARY, directory containing rccl library
+# This module assumes that the user has already called find_package(rocm)
+
+macro(find_rccl use_rccl)
+  if(${use_rccl} MATCHES ${IS_FALSE_PATTERN})
+    return()
+  endif()
+  if(${use_rccl} MATCHES ${IS_TRUE_PATTERN})
+    find_path(RCCL_INCLUDE_DIR NAMES rccl.h)
+    find_library(RCCL_LIBRARY NAMES rccl)
+  else()
+    find_path(RCCL_INCLUDE_DIR NAMES rccl.h HINTS ${use_rccl} ${use_rccl}/include)
+    find_library(RCCL_LIBRARY NAMES rccl HINTS ${use_rccl} ${use_rccl}/lib)
+  endif()
+  include(FindPackageHandleStandardArgs)
+  find_package_handle_standard_args(Rccl DEFAULT_MSG RCCL_INCLUDE_DIR RCCL_LIBRARY)
+  if (Rccl_FOUND)
+    message(STATUS "Found RCCL_LIBRARY: ${RCCL_LIBRARY}")
+    message(STATUS "Found RCCL_INCLUDE_DIR: ${RCCL_INCLUDE_DIR}")
+    add_library(rccl SHARED IMPORTED)
+    set_target_properties(rccl
+      PROPERTIES
+      INTERFACE_INCLUDE_DIRECTORIES "${RCCL_INCLUDE_DIR}"
+      IMPORTED_LOCATION             "${RCCL_LIBRARY}")
+  else()
+    message(STATUS "RCCL not found")
+  endif()
+  mark_as_advanced(RCCL_INCLUDE_DIR RCCL_LIBRARY)
+endmacro(find_rccl)

src/support/libinfo.cc

@@ -71,6 +71,10 @@
 #define TVM_INFO_ROCM_PATH "NOT-FOUND"
 #endif
 
+#ifndef TVM_INFO_USE_RCCL
+#define TVM_INFO_USE_RCCL "NOT-FOUND"
+#endif
+
 #ifndef TVM_INFO_USE_HEXAGON
 #define TVM_INFO_USE_HEXAGON "NOT-FOUND"
 #endif
@@ -325,6 +329,7 @@ TVM_DLL Map<String, String> GetLibInfo() {
       {"USE_RELAY_DEBUG", TVM_INFO_USE_RELAY_DEBUG},
       {"USE_ROCBLAS", TVM_INFO_USE_ROCBLAS},
       {"USE_ROCM", TVM_INFO_USE_ROCM},
+      {"USE_RCCL", TVM_INFO_USE_RCCL},
       {"USE_RPC", TVM_INFO_USE_RPC},
       {"USE_RTTI", TVM_INFO_USE_RTTI},
       {"USE_RUST_EXT", TVM_INFO_USE_RUST_EXT},