[BUILD] Enable path option for ROCM, CUDA, Vulkan, simplify optional build

CMakeLists.txt

@@ -3,7 +3,11 @@ project(tvm C CXX)
 
 # Utility functions
 include(cmake/util/Util.cmake)
+include(cmake/util/FindCUDA.cmake)
+include(cmake/util/FindVulkan.cmake)
 include(cmake/util/FindLLVM.cmake)
+include(cmake/util/FindROCM.cmake)
+
 
 if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/build/private/local_config.cmake)
   include(${CMAKE_CURRENT_SOURCE_DIR}/build/private/local_config.cmake)

cmake/config.cmake

@@ -22,12 +22,22 @@
 #---------------------------------------------
 # Backend runtimes.
 #---------------------------------------------
-# whether enable CUDA during compile
+
+# Whether enable CUDA during compile,
+#
+# Possible values:
+# - ON: enable CUDA with cmake's auto search
+# - OFF: disbale CUDA
+# - /path/to/cuda: use specific path to cuda toolkit
 set(USE_CUDA OFF)
 
-# ROCM
+# Whether enable ROCM runtime
+#
+# Possible values:
+# - ON: enable ROCM with cmake's auto search
+# - OFF: disbale ROCM
+# - /path/to/rocm: use specific path to rocm
 set(USE_ROCM OFF)
-set(ROCM_PATH "/opt/rocm")
 
 # Whether enable OpenCL runtime
 set(USE_OPENCL OFF)
@@ -36,6 +46,11 @@ set(USE_OPENCL OFF)
 set(USE_METAL OFF)
 
 # Whether enable Vulkan runtime
+#
+# Possible values:
+# - ON: enable Vulkan with cmake's auto search
+# - OFF: disbale vulkan
+# - /path/to/vulkan-sdk: use specific path to vulkan-sdk
 set(USE_VULKAN OFF)
 
 # Whether enable OpenGL runtime
@@ -54,9 +69,9 @@ set(USE_GRAPH_RUNTIME_DEBUG OFF)
 # Requires LLVM version >= 4.0
 #
 # Possible values:
-# - ON: enable llvm with cmake's find llvm
+# - ON: enable llvm with cmake's find search
 # - OFF: disbale llvm
-# - /path/to/llvm-config enable specific LLVM when multiple llvm-dev is available.
+# - /path/to/llvm-config: enable specific LLVM when multiple llvm-dev is available.
 set(USE_LLVM OFF)
 
 #---------------------------------------------

cmake/modules/CUDA.cmake

@@ -1,5 +1,5 @@
 # CUDA Module
-find_package(CUDA QUIET)
+find_cuda(${USE_CUDA})
 
 if(CUDA_FOUND)
   # always set the includedir when cuda is available
@@ -8,69 +8,33 @@ if(CUDA_FOUND)
 endif(CUDA_FOUND)
 
 if(USE_CUDA)
-  find_package(CUDA REQUIRED)
-  # Find CUDA doesn't find all the libraries we need, add the extra ones
-  find_library(CUDA_CUDA_LIBRARIES cuda
-    PATHS ${CUDA_TOOLKIT_ROOT_DIR}
-    PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs)
-  find_library(CUDA_NVRTC_LIBRARIES nvrtc
-    PATHS ${CUDA_TOOLKIT_ROOT_DIR}
-    PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs)
-  if(CUDA_CUDA_LIBRARIES)
-    set(CUDA_CUDA_LIBRARY ${CUDA_CUDA_LIBRARIES})
+  if(NOT CUDA_FOUND)
+    message(FATAL_ERROR "Cannot find CUDA, USE_CUDA=" ${USE_CUDA})
   endif()
   message(STATUS "Build with CUDA support")
   file(GLOB RUNTIME_CUDA_SRCS src/runtime/cuda/*.cc)
-  list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDART_LIBRARY})
-  list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDA_LIBRARY})
   list(APPEND RUNTIME_SRCS ${RUNTIME_CUDA_SRCS})
   list(APPEND COMPILER_SRCS src/codegen/opt/build_cuda_on.cc)
 
-  if(MSVC)
-    find_library(CUDA_NVRTC_LIB nvrtc
-      ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
-      ${CUDA_TOOLKIT_ROOT_DIR}/lib/win32)
-    list(APPEND TVM_LINKER_LIBS ${CUDA_NVRTC_LIB})
-  else(MSVC)
-    find_library(CUDA_NVRTC_LIB nvrtc
-      ${CUDA_TOOLKIT_ROOT_DIR}/lib64
-      ${CUDA_TOOLKIT_ROOT_DIR}/lib)
-    list(APPEND TVM_LINKER_LIBS ${CUDA_NVRTC_LIB})
-  endif(MSVC)
+  list(APPEND TVM_LINKER_LIBS ${CUDA_NVRTC_LIBRARY})
+  list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDART_LIBRARY})
+  list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDA_LIBRARY})
+  list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_NVRTC_LIBRARY})
 
   if(USE_CUDNN)
     message(STATUS "Build with cuDNN support")
     file(GLOB CONTRIB_CUDNN_SRCS src/contrib/cudnn/*.cc)
     list(APPEND RUNTIME_SRCS ${CONTRIB_CUDNN_SRCS})
-    if(MSVC)
-      find_library(CUDA_CUDNN_LIB cudnn
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib/win32)
-      list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDNN_LIB})
-    else(MSVC)
-      find_library(CUDA_CUDNN_LIB cudnn
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib64
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib)
-      list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDNN_LIB})
-    endif(MSVC)
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDNN_LIBRARY})
   endif(USE_CUDNN)
 
   if(USE_CUBLAS)
     message(STATUS "Build with cuBLAS support")
     file(GLOB CONTRIB_CUBLAS_SRCS src/contrib/cublas/*.cc)
     list(APPEND RUNTIME_SRCS ${CONTRIB_CUBLAS_SRCS})
-    if(MSVC)
-      find_library(CUDA_CUBLAS_LIB cublas
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib/win32)
-      list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUBLAS_LIB})
-    else(MSVC)
-      find_library(CUDA_CUBLAS_LIB cublas
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib64
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib)
-      list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUBLAS_LIB})
-    endif(MSVC)
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUBLAS_LIBRARY})
   endif(USE_CUBLAS)
+
 else(USE_CUDA)
   list(APPEND COMPILER_SRCS src/codegen/opt/build_cuda_off.cc)
 endif(USE_CUDA)

cmake/modules/LLVM.cmake

@@ -2,12 +2,7 @@
 add_definitions(-DDMLC_USE_FOPEN64=0)
 
 if(NOT USE_LLVM STREQUAL "OFF")
-  if(NOT USE_LLVM STREQUAL "ON")
-    set(LLVM_CONFIG "${USE_LLVM}")
-  else()
-    set(LLVM_CONFIG "")
-  endif()
-  find_llvm()
+  find_llvm(${USE_LLVM})
   include_directories(${LLVM_INCLUDE_DIRS})
   add_definitions(${LLVM_DEFINITIONS})
   message(STATUS "Build with LLVM " ${LLVM_PACKAGE_VERSION})

cmake/modules/Metal.cmake

@@ -5,7 +5,6 @@ if(USE_METAL)
   file(GLOB RUNTIME_METAL_SRCS src/runtime/metal/*.mm)
   list(APPEND TVM_RUNTIME_LINKER_LIBS ${METAL_LIB} ${FOUNDATION_LIB})
   list(APPEND RUNTIME_SRCS ${RUNTIME_METAL_SRCS})
-  list(APPEND COMPILER_SRCS src/codegen/opt/build_metal_on.cc)
 
   if(USE_MPS)
     file(GLOB MPS_CONTRIB_SRC src/contrib/mps/*.mm)

cmake/modules/OpenCL.cmake

@@ -13,7 +13,6 @@ if(USE_OPENCL)
   file(GLOB RUNTIME_OPENCL_SRCS src/runtime/opencl/*.cc)
   list(APPEND TVM_RUNTIME_LINKER_LIBS ${OpenCL_LIBRARIES})
   list(APPEND RUNTIME_SRCS ${RUNTIME_OPENCL_SRCS})
-  list(APPEND COMPILER_SRCS src/codegen/opt/build_opencl_on.cc)
 else()
   list(APPEND COMPILER_SRCS src/codegen/opt/build_opencl_off.cc)
 endif(USE_OPENCL)

cmake/modules/OpenGL.cmake

@@ -13,7 +13,6 @@ if(USE_OPENGL)
   file(GLOB RUNTIME_OPENGL_SRCS src/runtime/opengl/*.cc)
   list(APPEND TVM_RUNTIME_LINKER_LIBS ${OpenGL_LIBRARIES} glfw)
   list(APPEND RUNTIME_SRCS ${RUNTIME_OPENGL_SRCS})
-  list(APPEND COMPILER_SRCS src/codegen/opt/build_opengl_on.cc)
 else(USE_OPENGL)
   list(APPEND COMPILER_SRCS src/codegen/opt/build_opengl_off.cc)
 endif(USE_OPENGL)

cmake/modules/ROCM.cmake

@@ -1,34 +1,36 @@
 # ROCM Module
-if(NOT ROCM_PATH STREQUAL "")
-  include_directories(${ROCM_PATH}/include)
-  set(ROCM_LIB_PATH ${ROCM_PATH}/lib)
-else()
-  set(ROCM_LIB_PATH /lib)
-endif()
+find_rocm(${USE_ROCM})
+
+if(ROCM_FOUND)
+  # always set the includedir
+  # avoid global retrigger of cmake
+  include_directories(${ROCM_INCLUDE_DIRS})
+  add_definitions(-D__HIP_PLATFORM_HCC__=1)
+endif(ROCM_FOUND)
+
 
 if(USE_ROCM)
+  if(NOT ROCM_FOUND)
+    message(FATAL_ERROR "Cannot find ROCM, USE_ROCM=" ${USE_ROCM})
+  endif()
   message(STATUS "Build with ROCM support")
-  find_library(ROCM_LIBS hip_hcc ${ROCM_LIB_PATH})
   file(GLOB RUNTIME_ROCM_SRCS src/runtime/rocm/*.cc)
-  list(APPEND TVM_RUNTIME_LINKER_LIBS ${ROCM_LIBS})
   list(APPEND RUNTIME_SRCS ${RUNTIME_ROCM_SRCS})
-  add_definitions(-DTVM_ROCM_RUNTIME=1 -D__HIP_PLATFORM_HCC__=1)
+  list(APPEND TVM_RUNTIME_LINKER_LIBS ${ROCM_HIPHCC_LIBRARY})
 
   if(USE_MIOPEN)
     message(STATUS "Build with MIOpen support")
     file(GLOB MIOPEN_CONTRIB_SRCS src/contrib/miopen/*.cc)
     list(APPEND RUNTIME_SRCS ${MIOPEN_CONTRIB_SRCS})
-    find_library(MIOPEN_LIBS MIOpen ${ROCM_LIB_PATH})
-    list(APPEND TVM_RUNTIME_LINKER_LIBS ${MIOPEN_LIBS})
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${ROCM_MIOPEN_LIBRARY})
   endif(USE_MIOPEN)
 
   if(USE_ROCBLAS)
     message(STATUS "Build with RocBLAS support")
     file(GLOB ROCBLAS_CONTRIB_SRCS src/contrib/rocblas/*.cc)
     list(APPEND RUNTIME_SRCS ${ROCBLAS_CONTRIB_SRCS})
-    find_library(ROCBLAS_LIBS rocblas ${ROCM_LIB_PATH})
-    list(APPEND TVM_RUNTIME_LINKER_LIBS ${ROCBLAS_LIBS})
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${ROCM_ROCBLAS_LIBRARY})
   endif(USE_ROCBLAS)
 else(USE_ROCM)
-  add_definitions(-DTVM_ROCM_RUNTIME=0)
+  list(APPEND COMPILER_SRCS src/codegen/opt/build_rocm_off.cc)
 endif(USE_ROCM)

cmake/modules/Vulkan.cmake

@@ -1,31 +1,22 @@
 # Be compatible with older version of CMake
-if(NOT $ENV{VULKAN_SDK} STREQUAL "")
-  set(Vulkan_INCLUDE_DIRS $ENV{VULKAN_SDK}/include)
-  set(Vulkan_FOUND ON)
-else()
-  find_package(Vulkan QUIET)
-endif()
+find_vulkan(${USE_VULKAN})
 
 if(Vulkan_FOUND)
-  # always set the includedir when cuda is available
+  # always set the includedir
   # avoid global retrigger of cmake
   include_directories(${Vulkan_INCLUDE_DIRS})
 endif(Vulkan_FOUND)
 
 if(USE_VULKAN)
-  if(NOT $ENV{VULKAN_SDK} STREQUAL "")
-    find_library(Vulkan_LIBRARY vulkan $ENV{VULKAN_SDK}/lib)
-  else()
-    find_package(Vulkan REQUIRED)
+  if(NOT Vulkan_FOUND)
+    message(FATAL_ERROR "Cannot find Vulkan, USE_VULKAN=" ${USE_VULKAN})
   endif()
   message(STATUS "Build with VULKAN support")
   file(GLOB RUNTIME_VULKAN_SRCS src/runtime/vulkan/*.cc)
   file(GLOB COMPILER_VULKAN_SRCS src/codegen/spirv/*.cc)
-  list(APPEND TVM_RUNTIME_LINKER_LIBS ${Vulkan_LIBRARY})
   list(APPEND RUNTIME_SRCS ${RUNTIME_VULKAN_SRCS})
   list(APPEND COMPILER_SRCS ${COMPILER_VULKAN_SRCS})
-  get_filename_component(VULKAN_LIB_PATH ${Vulkan_LIBRARY} DIRECTORY)
-  find_library(SPIRV_TOOLS_LIB SPIRV-Tools
-               ${VULKAN_LIB_PATH}/spirv-tools)
-  list(APPEND TVM_LINKER_LIBS ${SPIRV_TOOLS_LIB})
+
+  list(APPEND TVM_LINKER_LIBS ${Vulkan_SPIRV_TOOLS_LIBRARY})
+  list(APPEND TVM_RUNTIME_LINKER_LIBS ${Vulkan_LIBRARY})
 endif(USE_VULKAN)

cmake/util/FindCUDA.cmake

@@ -0,0 +1,68 @@
+#######################################################
+# Enhanced version of find CUDA.
+#
+# Usage:
+#   find_cuda(${USE_CUDA})
+#
+# - When USE_CUDA=ON, use auto search
+# - When USE_CUDA=/path/to/cuda-path, use the cuda path
+#
+# Provide variables:
+#
+# - CUDA_FOUND
+# - CUDA_INCLUDE_DIRS
+# - CUDA_TOOLKIT_ROOT_DIR
+# - CUDA_CUDA_LIBRARY
+# - CUDA_CUDART_LIBRARY
+# - CUDA_NVRTC_LIBRARY
+# - CUDA_CUDNN_LIBRARY
+# - CUDA_CUBLAS_LIBRARY
+#
+macro(find_cuda use_cuda)
+  set(__use_cuda ${use_cuda})
+  if(__use_cuda STREQUAL "ON")
+    find_package(CUDA QUIET)
+  elseif(IS_DIRECTORY ${__use_cuda})
+    set(CUDA_TOOLKIT_ROOT_DIR ${__use_cuda})
+    message(STATUS "Custom CUDA_PATH=" ${CUDA_TOOLKIT_ROOT_DIR})
+    set(CUDA_INCLUDE_DIRS ${CUDA_TOOLKIT_ROOT_DIR}/include)
+    set(CUDA_FOUND TRUE)
+    if(MSVC)
+      find_library(CUDA_CUDAT_LIBRARY cudart
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib64
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib)
+    else(MSVC)
+      find_library(CUDA_CUDAT_LIBRARY cudart
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib64
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib)
+    endif(MSVC)
+  endif()
+
+  # additional libraries
+  if(CUDA_FOUND)
+    if(MSVC)
+      find_library(CUDA_NVRTC_LIBRARY cuda
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib/win32)
+      find_library(CUDA_NVRTC_LIBRARY nvrtc
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib/win32)
+      find_library(CUDA_CUDNN_LIBRARY cudnn
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib/win32)
+    else(MSVC)
+      find_library(_CUDA_CUDA_LIBRARY cuda
+        PATHS ${CUDA_TOOLKIT_ROOT_DIR}
+        PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs)
+      if(_CUDA_CUDA_LIBRARY)
+        set(CUDA_CUDA_LIBRARY ${_CUDA_CUDA_LIBRARY})
+      endif()
+      find_library(CUDA_NVRTC_LIBRARY nvrtc
+        PATHS ${CUDA_TOOLKIT_ROOT_DIR}
+        PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs)
+      find_library(CUDA_CUDNN_LIBRARY cudnn
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib64
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib)
+    endif(MSVC)
+  endif(CUDA_FOUND)
+endmacro(find_cuda)

cmake/util/FindLLVM.cmake

@@ -1,24 +1,26 @@
 #######################################################
-# Enhanced version of find llvm that allows set of LLVM_CONFIG
-# When LLVM_CONFIG_PATH is AUTO,
-# it defaults to system find llvm
+# Enhanced version of find llvm.
 #
 # Usage:
-#   find_llvm(LLVM_CONFIG_PATH)
+#   find_llvm(${USE_LLVM})
+#
+# - When USE_LLVM=ON, use auto search
+# - When USE_LLVM=/path/to/llvm-config, use corresponding config
 #
 # Provide variables:
 # - LLVM_INCLUDE_DIRS
 # - LLVM_LIBS
 # - LLVM_DEFINITIONS
-# - LLVM_VERSION_CONCAT
+# - TVM_LLVM_VERISON
 #
-macro(find_llvm)
-  if(LLVM_CONFIG STREQUAL "")
+macro(find_llvm use_llvm)
+  set(LLVM_CONFIG ${use_llvm})
+  if(LLVM_CONFIG STREQUAL "ON")
     find_package(LLVM REQUIRED CONFIG)
     llvm_map_components_to_libnames(LLVM_LIBS all)
     list(REMOVE_ITEM LLVM_LIBS LTO)
     set(TVM_LLVM_VERSION ${LLVM_VERSION_MAJOR}${LLVM_VERSION_MINOR})
-  else()
+  elseif(NOT LLVM_CONFIG STREQUAL "OFF")
     # use llvm config
     message(STATUS "Use llvm-config=" ${LLVM_CONFIG})
     execute_process(COMMAND ${LLVM_CONFIG} --includedir