[CMake] Add NCCL to TVM and TVM Runtime

This PR introduces NCCL in the cmake system.
NCCL is NVIDIA's library for distributed communication.

Author: junrushao

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)
@@ -832,3 +834,9 @@ 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)
+  find_nccl(${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/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)