Fix FindNVTX.cmake (#3421)

* Fix FindNVTX.cmake
* Set HINTS to find_path() rather than NVTX_ROOT
* Add CUDAToolkit_* HINTS in FindNCCL.cmake
* Fix lib paths for DDL

Signed-off-by: Max H. Gerlach <git@maxgerlach.de>

CMakeLists.txt

@@ -229,7 +229,7 @@ endif()
 if(HOROVOD_GPU_ALLREDUCE STREQUAL "D")
     message(DEPRECATION "DDL backend has been deprecated. Please, start using the NCCL backend by building Horovod with "
                         "'HOROVOD_GPU_OPERATIONS=NCCL'. Will be removed in v0.21.0.")
-    list(APPEND LINKER_LIBS "${CUDAToolkit_LIBRARY_ROOT}/lib/libddl.so" "${CUDAToolkit_LIBRARY_ROOT}/lib/libddl_pack.so")
+    list(APPEND LINKER_LIBS "${CUDAToolkit_LIBRARY_DIR}/libddl.so" "${CUDAToolkit_LIBRARY_DIR}/libddl_pack.so")
     list(APPEND SOURCES "${PROJECT_SOURCE_DIR}/horovod/common/mpi/ddl_mpi_context_manager.cc"
                         "${PROJECT_SOURCE_DIR}/horovod/common/ops/ddl_operations.cc")
     add_definitions(-DHAVE_DDL=1)

cmake/Modules/FindNCCL.cmake

@@ -11,7 +11,7 @@
 #  NCCL_LIBRARIES
 #  NCCL_MAJOR_VERSION
 #
-# The path hints include CUDAToolkit_LIBRARY_ROOT seeing as some folks
+# The path hints include CUDAToolkit_* seeing as some folks
 # install NCCL in the same location as the CUDA toolkit.
 
 set(HOROVOD_NCCL_HOME $ENV{HOROVOD_NCCL_HOME} CACHE PATH "Folder contains NVIDIA NCCL")
@@ -24,7 +24,7 @@ list(APPEND CMAKE_PREFIX_PATH ${NCCL_ROOT})
 
 find_path(NCCL_INCLUDE_DIR
         NAMES nccl.h
-        HINTS ${HOROVOD_NCCL_INCLUDE})
+        HINTS ${HOROVOD_NCCL_INCLUDE} ${CUDAToolkit_INCLUDE_DIRS})
 
 set(HOROVOD_NCCL_LINK $ENV{HOROVOD_NCCL_LINK})
 if (HOROVOD_NCCL_LINK STREQUAL "SHARED")
@@ -36,7 +36,7 @@ endif()
 
 find_library(NCCL_LIBRARY
         NAMES ${NCCL_LIBNAME}
-        HINTS ${HOROVOD_NCCL_LIB})
+        HINTS ${HOROVOD_NCCL_LIB} ${CUDAToolkit_LIBRARY_DIR})
 
 include(FindPackageHandleStandardArgs)
 find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIR NCCL_LIBRARY)

cmake/Modules/FindNVTX.cmake

@@ -1,6 +1,6 @@
 # Try to find NVTX
 #
-# NVTX comes with the CUDA toolkit so we use that root dir to search for the header-only variation of NVTX.
+# NVTX comes with the CUDA toolkit so we use those include dirs to search for the header-only variation of NVTX.
 # Alternatively an explicit path can be given via the variable HOROVOD_NVTX_INCLUDE
 #
 # The following are set after configuration is done:
@@ -10,13 +10,9 @@
 
 set(HOROVOD_NVTX_INCLUDE $ENV{HOROVOD_NVTX_INCLUDE} CACHE PATH "Folder containing NVIDIA NVTX3 headers")
 
-list(APPEND NVTX_ROOT ${CUDAToolkit_LIBRARY_ROOT})
-# Compatible layer for CMake <3.12:
-list(APPEND CMAKE_PREFIX_PATH ${NVTX_ROOT})
-
 find_path(NVTX_INCLUDE_DIR
           NAMES nvtx3/nvToolsExt.h
-          HINTS ${HOROVOD_NVTX_INCLUDE})
+          HINTS ${HOROVOD_NVTX_INCLUDE} ${CUDAToolkit_INCLUDE_DIRS})
 
 include(FindPackageHandleStandardArgs)
 find_package_handle_standard_args(NVTX DEFAULT_MSG NVTX_INCLUDE_DIR)