From 2c4786b06188bdbcd898808da8712430f22ad29f Mon Sep 17 00:00:00 2001 From: Lesheng Jin Date: Sat, 26 Aug 2023 05:03:00 -0400 Subject: [PATCH] [CMake] Add RCCL to TVM and TVM Runtime --- CMakeLists.txt | 14 ++++++++++ cmake/config.cmake | 6 +++++ cmake/modules/LibInfo.cmake | 1 + cmake/utils/FindRCCL.cmake | 52 +++++++++++++++++++++++++++++++++++++ src/support/libinfo.cc | 5 ++++ 5 files changed, 78 insertions(+) create mode 100644 cmake/utils/FindRCCL.cmake diff --git a/CMakeLists.txt b/CMakeLists.txt index 22b619675447..7b69145806ad 100644 --- a/CMakeLists.txt +++ b/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() diff --git a/cmake/config.cmake b/cmake/config.cmake index 08cded07df3e..f0b4d9b5c508 100644 --- a/cmake/config.cmake +++ b/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) diff --git a/cmake/modules/LibInfo.cmake b/cmake/modules/LibInfo.cmake index 611f014e3222..ad153cce0455 100644 --- a/cmake/modules/LibInfo.cmake +++ b/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}" diff --git a/cmake/utils/FindRCCL.cmake b/cmake/utils/FindRCCL.cmake new file mode 100644 index 000000000000..93d8c8744630 --- /dev/null +++ b/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) diff --git a/src/support/libinfo.cc b/src/support/libinfo.cc index 0c39adef535a..53f9292d162f 100644 --- a/src/support/libinfo.cc +++ b/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 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},