[CMake] Add NCCL to TVM and TVM Runtime

CMakeLists.txt

@@ -6,6 +6,7 @@ include(cmake/utils/Utils.cmake)
 include(cmake/utils/Summary.cmake)
 include(cmake/utils/Linker.cmake)
 include(cmake/utils/FindCUDA.cmake)
+include(cmake/utils/FindNCCL.cmake)
 include(cmake/utils/FindOpenCL.cmake)
 include(cmake/utils/FindVulkan.cmake)
 include(cmake/utils/FindLLVM.cmake)
@@ -25,6 +26,7 @@ endif()
 # and add set(OPTION VALUE) to override these build options.
 # Alernatively, use cmake -DOPTION=VALUE through command-line.
 tvm_option(USE_CUDA "Build with CUDA" OFF)
+tvm_option(USE_NCCL "Build with NCCL" OFF)
 tvm_option(USE_OPENCL "Build with OpenCL" OFF)
 tvm_option(USE_OPENCL_ENABLE_HOST_PTR "Enable OpenCL memory object access to host" OFF)
 tvm_option(USE_OPENCL_GTEST "Path to OpenCL specific gtest version for runtime cpp tests." /path/to/opencl/gtest)
@@ -338,6 +340,7 @@ list(APPEND COMPILER_SRCS "src/target/datatype/myfloat/myfloat.cc")
 tvm_file_glob(GLOB RUNTIME_SRCS
   src/runtime/*.cc
   src/runtime/vm/*.cc
+  src/runtime/disco/*.cc
   src/runtime/minrpc/*.cc
 )
 
@@ -421,6 +424,13 @@ if(USE_PROFILER)
   list(APPEND RUNTIME_SRCS ${RUNTIME_VM_PROFILER_SRCS})
 endif(USE_PROFILER)
 
+if(USE_CUDA AND USE_NCCL)
+  message(STATUS "Build with NCCL...")
+  find_nccl(${USE_NCCL})
+  tvm_file_glob(GLOB RUNTIME_NCCL_SRC src/runtime/disco/nccl/*.cc)
+  list(APPEND RUNTIME_SRCS ${RUNTIME_NCCL_SRC})
+endif()
+
 if(USE_AOT_EXECUTOR)
   message(STATUS "Build with AOT Executor support...")
   file(GLOB RUNTIME_AOT_EXECUTOR_SRCS src/runtime/aot_executor/*.cc)
@@ -832,3 +842,8 @@ if(USE_CUDA AND USE_CUTLASS)
   target_link_libraries(tvm PRIVATE fpA_intB_gemm)
   target_link_libraries(tvm_runtime PRIVATE fpA_intB_gemm)
 endif()
+
+if(USE_CUDA AND USE_NCCL)
+  target_link_libraries(tvm PRIVATE nccl)
+  target_link_libraries(tvm_runtime PRIVATE nccl)
+endif()

cmake/config.cmake

@@ -48,6 +48,12 @@
 # - /path/to/cuda: use specific path to cuda toolkit
 set(USE_CUDA OFF)
 
+# Whether to enable NCCL support:
+# - ON: enable NCCL with cmake's auto search
+# - OFF: disable NCCL
+# - /path/to/nccl: use specific path to nccl
+set(USE_NCCL OFF)
+
 # Whether enable ROCM runtime
 #
 # Possible values:

cmake/modules/LibInfo.cmake

@@ -64,6 +64,7 @@ function(add_lib_info src_file)
     TVM_INFO_USE_CPP_RTVM="${USE_CPP_RTVM}"
     TVM_INFO_USE_CUBLAS="${USE_CUBLAS}"
     TVM_INFO_USE_CUDA="${USE_CUDA}"
+    TVM_INFO_USE_NCCL="${USE_NCCL}"
     TVM_INFO_USE_CUDNN="${USE_CUDNN}"
     TVM_INFO_USE_CUSTOM_LOGGING="${USE_CUSTOM_LOGGING}"
     TVM_INFO_USE_CUTLASS="${USE_CUTLASS}"

cmake/utils/FindNCCL.cmake

@@ -0,0 +1,56 @@
+# 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:
+#
+#  NCCL_ROOT - When set, this path is inspected instead of standard library
+#              locations as the root of the NCCL installation.
+#              The environment variable NCCL_ROOT overrides this variable.
+#
+# This module defines
+#  Nccl_FOUND, whether nccl has been found
+#  NCCL_INCLUDE_DIR, directory containing header
+#  NCCL_LIBRARY, directory containing nccl library
+# This module assumes that the user has already called find_package(CUDA)
+
+macro(find_nccl use_nccl)
+  if(${use_nccl} MATCHES ${IS_FALSE_PATTERN})
+    return()
+  endif()
+  if(${use_nccl} MATCHES ${IS_TRUE_PATTERN})
+    find_path(NCCL_INCLUDE_DIR NAMES nccl.h)
+    find_library(NCCL_LIBRARY NAMES nccl)
+  else()
+    find_path(NCCL_INCLUDE_DIR NAMES nccl.h HINTS ${use_nccl} ${use_nccl}/include)
+    find_library(NCCL_LIBRARY NAMES nccl HINTS ${use_nccl} ${use_nccl}/lib)
+  endif()
+  include(FindPackageHandleStandardArgs)
+  find_package_handle_standard_args(Nccl DEFAULT_MSG NCCL_INCLUDE_DIR NCCL_LIBRARY)
+  if (Nccl_FOUND)
+    message(STATUS "Found NCCL_LIBRARY: ${NCCL_LIBRARY}")
+    message(STATUS "Found NCCL_INCLUDE_DIR: ${NCCL_INCLUDE_DIR}")
+    add_library(nccl SHARED IMPORTED)
+    set_target_properties(nccl
+      PROPERTIES
+      INTERFACE_INCLUDE_DIRECTORIES "${NCCL_INCLUDE_DIR}"
+      IMPORTED_LOCATION             "${NCCL_LIBRARY}")
+  else()
+    message(STATUS "NCCL not found")
+  endif()
+  mark_as_advanced(NCCL_INCLUDE_DIR NCCL_LIBRARY)
+endmacro(find_nccl)

src/support/libinfo.cc

@@ -35,6 +35,10 @@
 #define TVM_INFO_USE_CUDA "NOT-FOUND"
 #endif
 
+#ifndef TVM_INFO_USE_NCCL
+#define TVM_INFO_USE_NCCL "NOT-FOUND"
+#endif
+
 #ifndef TVM_INFO_CUDA_VERSION
 #define TVM_INFO_CUDA_VERSION "NOT-FOUND"
 #endif
@@ -281,6 +285,7 @@ TVM_DLL Map<String, String> GetLibInfo() {
       {"USE_CPP_RTVM", TVM_INFO_USE_CPP_RTVM},
       {"USE_CUBLAS", TVM_INFO_USE_CUBLAS},
       {"USE_CUDA", TVM_INFO_USE_CUDA},
+      {"USE_NCCL", TVM_INFO_USE_NCCL},
       {"USE_CUDNN", TVM_INFO_USE_CUDNN},
       {"USE_CUSTOM_LOGGING", TVM_INFO_USE_CUSTOM_LOGGING},
       {"USE_CUTLASS", TVM_INFO_USE_CUTLASS},