diff --git a/CMakeLists.txt b/CMakeLists.txt index 67e87d907141..22b619675447 100644 --- a/CMakeLists.txt +++ b/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() diff --git a/cmake/config.cmake b/cmake/config.cmake index 8a7a0f1fdd29..c1a7d39edc3d 100644 --- a/cmake/config.cmake +++ b/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: diff --git a/cmake/modules/LibInfo.cmake b/cmake/modules/LibInfo.cmake index 17a56ac439e2..611f014e3222 100644 --- a/cmake/modules/LibInfo.cmake +++ b/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}" diff --git a/cmake/utils/FindNCCL.cmake b/cmake/utils/FindNCCL.cmake new file mode 100644 index 000000000000..0cabaf74f879 --- /dev/null +++ b/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) diff --git a/src/support/libinfo.cc b/src/support/libinfo.cc index 23091fffd23a..0c39adef535a 100644 --- a/src/support/libinfo.cc +++ b/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 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},