[CMAKE] Misc improvment of Util (#16900)

This PR updates the utils so tvm_option can take in list argument. Also introduces a flag for MSCCLPP.
apache · Apr 18, 2024 · fe52709 · fe52709
1 parent 59376ee
commit fe52709
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 9 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -28,6 +28,7 @@ endif()
 # 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_MSCCL "Build with MSCCL" 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)
@@ -940,8 +941,8 @@ endif()
 
 if(USE_CUDA AND USE_NCCL)
   find_library(LIBRT rt)
-  target_link_libraries(tvm PRIVATE nccl msccl ${LIBRT})
-  target_link_libraries(tvm_runtime PRIVATE nccl msccl ${LIBRT})
+  target_link_libraries(tvm PRIVATE nccl ${LIBRT})
+  target_link_libraries(tvm_runtime PRIVATE nccl ${LIBRT})
 endif()
 
 if(USE_ROCM AND USE_RCCL)

diff --git a/cmake/config.cmake b/cmake/config.cmake
@@ -54,6 +54,11 @@ set(USE_CUDA OFF)
 # - /path/to/nccl: use specific path to nccl
 set(USE_NCCL OFF)
 
+# Whether to enable MSCCL support:
+# - ON: enable MSCCL
+# - OFF: disable MSCCL
+set(USE_MSCCL OFF)
+
 # Whether to enable NVTX support (must have USE_CUDA enabled):
 # - ON: enable NCCL with cmake's auto search
 # - OFF: disable NCCL

diff --git a/cmake/modules/LibInfo.cmake b/cmake/modules/LibInfo.cmake
@@ -72,6 +72,7 @@ function(add_lib_info src_file)
     TVM_INFO_USE_CUDA="${USE_CUDA}"
     TVM_INFO_USE_NVTX="${USE_NVTX}"
     TVM_INFO_USE_NCCL="${USE_NCCL}"
+    TVM_INFO_USE_MSCCL="${USE_MSCCL}"
     TVM_INFO_USE_CUDNN="${USE_CUDNN}"
     TVM_INFO_USE_CUSTOM_LOGGING="${USE_CUSTOM_LOGGING}"
     TVM_INFO_USE_CUTLASS="${USE_CUTLASS}"

diff --git a/cmake/modules/contrib/MSCCLPP.cmake b/cmake/modules/contrib/MSCCLPP.cmake
@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-if(USE_CUDA AND USE_NCCL)
+if(USE_CUDA AND USE_NCCL AND USE_MSCCL)
   include(FetchContent)
   FetchContent_Declare(
     mscclpp
@@ -46,5 +46,5 @@ if(USE_CUDA AND USE_NCCL)
     FILE_SET HEADERS DESTINATION ${INSTALL_PREFIX}/include)
   install(TARGETS mscclpp EXPORT ${PROJECT_NAME}Targets DESTINATION lib${LIB_SUFFIX})
   install(TARGETS msccl EXPORT ${PROJECT_NAME}Targets DESTINATION lib${LIB_SUFFIX})
-
+  list(APPEND TVM_RUNTIME_LINKER_LIBS msccl)
 endif()
diff --git a/cmake/utils/Utils.cmake b/cmake/utils/Utils.cmake
@@ -46,11 +46,8 @@ macro(tvm_option variable description value)
 
   if(${__condition})
     if("${__value}" MATCHES ";")
-      if(${__value})
-        __tvm_option(${variable} "${description}" ON)
-      else()
-        __tvm_option(${variable} "${description}" OFF)
-      endif()
+      # list values directly pass through
+      __tvm_option(${variable} "${description}" "${__value}")
     elseif(DEFINED ${__value})
       if(${__value})
         __tvm_option(${variable} "${description}" ON)

diff --git a/src/support/libinfo.cc b/src/support/libinfo.cc
@@ -47,6 +47,10 @@
 #define TVM_INFO_USE_NCCL "NOT-FOUND"
 #endif
 
+#ifndef TVM_INFO_USE_MSCCLPP
+#define TVM_INFO_USE_MSCCLPP "NOT-FOUND"
+#endif
+
 #ifndef TVM_INFO_CUDA_VERSION
 #define TVM_INFO_CUDA_VERSION "NOT-FOUND"
 #endif
@@ -308,6 +312,7 @@ TVM_DLL Map<String, String> GetLibInfo() {
       {"USE_CUDA", TVM_INFO_USE_CUDA},
       {"USE_NVTX", TVM_INFO_USE_NVTX},
       {"USE_NCCL", TVM_INFO_USE_NCCL},
+      {"USE_MSCCL", TVM_INFO_USE_MSCCL},
       {"USE_CUDNN", TVM_INFO_USE_CUDNN},
       {"USE_CUSTOM_LOGGING", TVM_INFO_USE_CUSTOM_LOGGING},
       {"USE_CUTLASS", TVM_INFO_USE_CUTLASS},