CUDA: Determine CUDA toolkit location for NVCC

Similar to how we already do for Clang. Avoids a lot of redundant work in
FindCUDAToolkit.

4 files changed, 134 insertions, 126 deletions

diff --git a/Help/release/3.18.rst b/Help/release/3.18.rst
index 027ef5c..db5f62a 100644
--- a/Help/release/3.18.rst
+++ b/Help/release/3.18.rst
@@ -194,9 +194,13 @@ Modules
 
 .. _`SWIG-Fortran`: https://github.com/swig-fortran/swig
 
-* The :module:`FindCUDAToolkit` module gained the variable
-  ``CUDAToolkit_LIBRARY_ROOT``, which is the directory containing the ``nvvm``
-  directory and ``version.txt``.
+* The :module:`FindCUDAToolkit` module:
+
+  * gained the variable
+    ``CUDAToolkit_LIBRARY_ROOT``, which is the directory containing the
+    ``nvvm`` directory and ``version.txt``.
+
+  * uses toolkit and library root found during ``CUDA`` compiler detection.
 
 Generator Expressions
 ---------------------
diff --git a/Modules/CMakeCUDACompiler.cmake.in b/Modules/CMakeCUDACompiler.cmake.in
index 64fb469..704ad09 100644
--- a/Modules/CMakeCUDACompiler.cmake.in
+++ b/Modules/CMakeCUDACompiler.cmake.in
@@ -43,6 +43,9 @@ if(CMAKE_CUDA_LIBRARY_ARCHITECTURE)
   set(CMAKE_LIBRARY_ARCHITECTURE "@CMAKE_CUDA_LIBRARY_ARCHITECTURE@")
 endif()
 
+set(CMAKE_CUDA_COMPILER_TOOLKIT_ROOT "@CMAKE_CUDA_COMPILER_TOOLKIT_ROOT@")
+set(CMAKE_CUDA_COMPILER_LIBRARY_ROOT "@CMAKE_CUDA_COMPILER_LIBRARY_ROOT@")
+
 set(CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES "@CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES@")
 
 set(CMAKE_CUDA_HOST_IMPLICIT_LINK_LIBRARIES "@CMAKE_CUDA_HOST_IMPLICIT_LINK_LIBRARIES@")
diff --git a/Modules/CMakeDetermineCUDACompiler.cmake b/Modules/CMakeDetermineCUDACompiler.cmake
index 3411b80..9dcebe8 100644
--- a/Modules/CMakeDetermineCUDACompiler.cmake
+++ b/Modules/CMakeDetermineCUDACompiler.cmake
@@ -65,130 +65,140 @@ if(NOT CMAKE_CUDA_COMPILER_ID_RUN)
   set(CMAKE_CXX_COMPILER_ID_TOOL_MATCH_INDEX 2)
   set(CMAKE_CUDA_COMPILER_ID_FLAGS_ALWAYS "-v")
 
-  # We determine the vendor so we can skip doing extra work not necessary for a given compiler.
-  # E.g. skip finding the CUDA toolkit if the compiler isn't Clang.
+  # We determine the vendor to help with find the toolkit and use the right flags for detection right away.
   include(${CMAKE_ROOT}/Modules/CMakeDetermineCompilerId.cmake)
   CMAKE_DETERMINE_COMPILER_ID_VENDOR(CUDA "--version")
 
+  # Find the CUDA toolkit. We store the CMAKE_CUDA_COMPILER_TOOLKIT_ROOT and CMAKE_CUDA_COMPILER_LIBRARY_ROOT
+  # in CMakeCUDACompiler.cmake, so FindCUDAToolkit can avoid searching on future runs and the toolkit stays the same.
+  # This is very similar to FindCUDAToolkit, but somewhat simplified since we can issue fatal errors
+  # if we fail to find things we need and we don't need to account for searching the libraries.
+
+  # For NVCC we can easily deduce the SDK binary directory from the compiler path.
   if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA")
-    set(nvcc_test_flags "--keep --keep-dir tmp")
-    if(CMAKE_CUDA_HOST_COMPILER)
-      string(APPEND nvcc_test_flags " -ccbin=${CMAKE_CUDA_HOST_COMPILER}")
-    endif()
-  elseif(CMAKE_CUDA_COMPILER_ID STREQUAL "Clang")
-    # We search for the CUDA toolkit and compute _CUDA_ROOT_DIR and _CUDA_LIBRARY_ROOT so we can
-    # pass the toolkit path to Clang.
-    # This is very similar to FindCUDAToolkit, but simplified due to not having to account for NVCC.
-    # There are differences in searching to get equivalent behaviour to FindCUDAToolkit.
-
-    # Search using CUDAToolkit_ROOT and then CUDA_PATH for equivalence with FindCUDAToolkit.
-    # In FindCUDAToolkit CUDAToolkit_ROOT is searched automatically due to being in a find_package().
-    # First we search candidate non-default paths to give them priority.
-    find_program(_CUDA_NVCC_EXECUTABLE
+    get_filename_component(_CUDA_BIN_DIR "${CMAKE_CUDA_COMPILER}" DIRECTORY)
+    find_program(CUDAToolkit_NVCC_EXECUTABLE
       NAMES nvcc nvcc.exe
-      PATHS ${CUDAToolkit_ROOT}
-      ENV CUDAToolkit_ROOT
-      ENV CUDA_PATH
-      PATH_SUFFIXES bin
+      PATHS ${_CUDA_BIN_DIR}
       NO_DEFAULT_PATH
     )
+    unset(_CUDA_BIN_DIR)
+  endif()
 
-    # If we didn't find NVCC, then try the default paths.
-    find_program(_CUDA_NVCC_EXECUTABLE
-      NAMES nvcc nvcc.exe
-      PATH_SUFFIXES bin
-    )
+  # Search using CUDAToolkit_ROOT and then CUDA_PATH for equivalence with FindCUDAToolkit.
+  # In FindCUDAToolkit CUDAToolkit_ROOT is searched automatically due to being in a find_package().
+  # First we search candidate non-default paths to give them priority.
+  find_program(_CUDA_NVCC_EXECUTABLE
+    NAMES nvcc nvcc.exe
+    PATHS ${CUDAToolkit_ROOT}
+    ENV CUDAToolkit_ROOT
+    ENV CUDA_PATH
+    PATH_SUFFIXES bin
+    NO_DEFAULT_PATH
+  )
+
+  # If we didn't find NVCC, then try the default paths.
+  find_program(_CUDA_NVCC_EXECUTABLE
+    NAMES nvcc nvcc.exe
+    PATH_SUFFIXES bin
+  )
 
-    # If the user specified CUDAToolkit_ROOT but nvcc could not be found, this is an error.
-    if(NOT _CUDA_NVCC_EXECUTABLE AND (DEFINED CUDAToolkit_ROOT OR DEFINED ENV{CUDAToolkit_ROOT}))
-      set(fail_base "Could not find nvcc executable in path specified by")
+  # If the user specified CUDAToolkit_ROOT but nvcc could not be found, this is an error.
+  if(NOT _CUDA_NVCC_EXECUTABLE AND (DEFINED CUDAToolkit_ROOT OR DEFINED ENV{CUDAToolkit_ROOT}))
+    set(fail_base "Could not find nvcc executable in path specified by")
 
-      if(DEFINED CUDAToolkit_ROOT)
-        message(FATAL_ERROR "${fail_base} CUDAToolkit_ROOT=${CUDAToolkit_ROOT}")
-      elseif(DEFINED ENV{CUDAToolkit_ROOT})
-        message(FATAL_ERROR "${fail_base} environment variable CUDAToolkit_ROOT=$ENV{CUDAToolkit_ROOT}")
-      endif()
+    if(DEFINED CUDAToolkit_ROOT)
+      message(FATAL_ERROR "${fail_base} CUDAToolkit_ROOT=${CUDAToolkit_ROOT}")
+    elseif(DEFINED ENV{CUDAToolkit_ROOT})
+      message(FATAL_ERROR "${fail_base} environment variable CUDAToolkit_ROOT=$ENV{CUDAToolkit_ROOT}")
     endif()
+  endif()
 
-    # CUDAToolkit_ROOT cmake/env variable not specified, try platform defaults.
-    #
-    # - Linux: /usr/local/cuda-X.Y
-    # - macOS: /Developer/NVIDIA/CUDA-X.Y
-    # - Windows: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vX.Y
-    #
-    # We will also search the default symlink location /usr/local/cuda first since
-    # if CUDAToolkit_ROOT is not specified, it is assumed that the symlinked
-    # directory is the desired location.
-    if(NOT _CUDA_NVCC_EXECUTABLE)
-      if(UNIX)
-        if(NOT APPLE)
-          set(platform_base "/usr/local/cuda-")
-        else()
-          set(platform_base "/Developer/NVIDIA/CUDA-")
-        endif()
+  # CUDAToolkit_ROOT cmake/env variable not specified, try platform defaults.
+  #
+  # - Linux: /usr/local/cuda-X.Y
+  # - macOS: /Developer/NVIDIA/CUDA-X.Y
+  # - Windows: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vX.Y
+  #
+  # We will also search the default symlink location /usr/local/cuda first since
+  # if CUDAToolkit_ROOT is not specified, it is assumed that the symlinked
+  # directory is the desired location.
+  if(NOT _CUDA_NVCC_EXECUTABLE)
+    if(UNIX)
+      if(NOT APPLE)
+        set(platform_base "/usr/local/cuda-")
       else()
-        set(platform_base "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v")
+        set(platform_base "/Developer/NVIDIA/CUDA-")
       endif()
+    else()
+      set(platform_base "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v")
+    endif()
 
-      # Build out a descending list of possible cuda installations, e.g.
-      file(GLOB possible_paths "${platform_base}*")
-      # Iterate the glob results and create a descending list.
-      set(versions)
-      foreach(p ${possible_paths})
-        # Extract version number from end of string
-        string(REGEX MATCH "[0-9][0-9]?\\.[0-9]$" p_version ${p})
-        if(IS_DIRECTORY ${p} AND p_version)
-          list(APPEND versions ${p_version})
-        endif()
-      endforeach()
-
-      # Sort numerically in descending order, so we try the newest versions first.
-      list(SORT versions COMPARE NATURAL ORDER DESCENDING)
+    # Build out a descending list of possible cuda installations, e.g.
+    file(GLOB possible_paths "${platform_base}*")
+    # Iterate the glob results and create a descending list.
+    set(versions)
+    foreach(p ${possible_paths})
+      # Extract version number from end of string
+      string(REGEX MATCH "[0-9][0-9]?\\.[0-9]$" p_version ${p})
+      if(IS_DIRECTORY ${p} AND p_version)
+        list(APPEND versions ${p_version})
+      endif()
+    endforeach()
 
-      # With a descending list of versions, populate possible paths to search.
-      set(search_paths)
-      foreach(v ${versions})
-        list(APPEND search_paths "${platform_base}${v}")
-      endforeach()
+    # Sort numerically in descending order, so we try the newest versions first.
+    list(SORT versions COMPARE NATURAL ORDER DESCENDING)
 
-      # Force the global default /usr/local/cuda to the front on Unix.
-      if(UNIX)
-        list(INSERT search_paths 0 "/usr/local/cuda")
-      endif()
+    # With a descending list of versions, populate possible paths to search.
+    set(search_paths)
+    foreach(v ${versions})
+      list(APPEND search_paths "${platform_base}${v}")
+    endforeach()
 
-      # Now search for nvcc again using the platform default search paths.
-      find_program(_CUDA_NVCC_EXECUTABLE
-        NAMES nvcc nvcc.exe
-        PATHS ${search_paths}
-        PATH_SUFFIXES bin
-      )
-
-      # We are done with these variables now, cleanup.
-      unset(platform_base)
-      unset(possible_paths)
-      unset(versions)
-      unset(search_paths)
-
-      if(NOT _CUDA_NVCC_EXECUTABLE)
-        message(FATAL_ERROR "Could not find nvcc, please set CUDAToolkit_ROOT.")
-      endif()
+    # Force the global default /usr/local/cuda to the front on Unix.
+    if(UNIX)
+      list(INSERT search_paths 0 "/usr/local/cuda")
     endif()
 
-    get_filename_component(_CUDA_ROOT_DIR "${_CUDA_NVCC_EXECUTABLE}" DIRECTORY)
-    get_filename_component(_CUDA_ROOT_DIR "${_CUDA_ROOT_DIR}" DIRECTORY ABSOLUTE)
-
-    # _CUDA_LIBRARY_ROOT contains the device library and version file.
-    # In a non-scattered installation this is equivalent to _CUDA_ROOT_DIR.
-    # We first check for a non-scattered installation to prefer it over a scattered installation.
-    if(EXISTS "${_CUDA_ROOT_DIR}/version.txt")
-      set(_CUDA_LIBRARY_ROOT "${_CUDA_ROOT_DIR}")
-    elseif(CMAKE_SYSROOT_LINK AND EXISTS "${CMAKE_SYSROOT_LINK}/usr/lib/cuda/version.txt")
-      set(_CUDA_LIBRARY_ROOT "${CMAKE_SYSROOT_LINK}/usr/lib/cuda")
-    elseif(EXISTS "${CMAKE_SYSROOT}/usr/lib/cuda/version.txt")
-      set(_CUDA_LIBRARY_ROOT "${CMAKE_SYSROOT}/usr/lib/cuda")
+    # Now search for nvcc again using the platform default search paths.
+    find_program(_CUDA_NVCC_EXECUTABLE
+      NAMES nvcc nvcc.exe
+      PATHS ${search_paths}
+      PATH_SUFFIXES bin
+    )
+
+    # We are done with these variables now, cleanup.
+    unset(platform_base)
+    unset(possible_paths)
+    unset(versions)
+    unset(search_paths)
+
+    if(NOT _CUDA_NVCC_EXECUTABLE)
+      message(FATAL_ERROR "Could not find nvcc, please set CUDAToolkit_ROOT.")
     endif()
+  endif()
 
-    set(clang_test_flags "--cuda-path=\"${_CUDA_LIBRARY_ROOT}\"")
+  get_filename_component(CMAKE_CUDA_COMPILER_TOOLKIT_ROOT "${_CUDA_NVCC_EXECUTABLE}" DIRECTORY)
+  get_filename_component(CMAKE_CUDA_COMPILER_TOOLKIT_ROOT "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}" DIRECTORY ABSOLUTE)
+
+  # CMAKE_CUDA_COMPILER_LIBRARY_ROOT contains the device library and version file.
+  # In a non-scattered installation this is equivalent to CMAKE_CUDA_COMPILER_TOOLKIT_ROOT.
+  # We first check for a non-scattered installation to prefer it over a scattered installation.
+  if(EXISTS "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}/version.txt")
+    set(CMAKE_CUDA_COMPILER_LIBRARY_ROOT "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}")
+  elseif(CMAKE_SYSROOT_LINK AND EXISTS "${CMAKE_SYSROOT_LINK}/usr/lib/cuda/version.txt")
+    set(CMAKE_CUDA_COMPILER_LIBRARY_ROOT "${CMAKE_SYSROOT_LINK}/usr/lib/cuda")
+  elseif(EXISTS "${CMAKE_SYSROOT}/usr/lib/cuda/version.txt")
+    set(CMAKE_CUDA_COMPILER_LIBRARY_ROOT "${CMAKE_SYSROOT}/usr/lib/cuda")
+  endif()
+
+  if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA")
+    set(nvcc_test_flags "--keep --keep-dir tmp")
+    if(CMAKE_CUDA_HOST_COMPILER)
+      string(APPEND nvcc_test_flags " -ccbin=${CMAKE_CUDA_HOST_COMPILER}")
+    endif()
+  elseif(CMAKE_CUDA_COMPILER_ID STREQUAL "Clang")
+    set(clang_test_flags "--cuda-path=\"${CMAKE_CUDA_COMPILER_LIBRARY_ROOT}\"")
     if(CMAKE_CROSSCOMPILING)
       # Need to pass the host target and include directories if we're crosscompiling.
       string(APPEND clang_test_flags " --sysroot=\"${CMAKE_SYSROOT}\" --target=${CMAKE_CUDA_COMPILER_TARGET}")
@@ -295,11 +305,11 @@ elseif(CMAKE_CUDA_COMPILER_ID STREQUAL "Clang")
       set(_CUDA_TARGET_NAME "x86_64-linux")
     endif()
 
-    if(EXISTS "${_CUDA_ROOT_DIR}/targets/${_CUDA_TARGET_NAME}")
-      set(_CUDA_TARGET_DIR "${_CUDA_ROOT_DIR}/targets/${_CUDA_TARGET_NAME}")
+    if(EXISTS "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}/targets/${_CUDA_TARGET_NAME}")
+      set(_CUDA_TARGET_DIR "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}/targets/${_CUDA_TARGET_NAME}")
     endif()
   else()
-    set(_CUDA_TARGET_DIR "${_CUDA_ROOT_DIR}")
+    set(_CUDA_TARGET_DIR "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}")
   endif()
 
   # We can't use find_library() yet at this point, so try a few guesses.
@@ -328,13 +338,6 @@ elseif(CMAKE_CUDA_COMPILER_ID STREQUAL "Clang")
   set(CMAKE_CUDA_HOST_IMPLICIT_LINK_DIRECTORIES "${_CUDA_LIBRARY_DIR}")
   set(CMAKE_CUDA_HOST_IMPLICIT_LINK_LIBRARIES "")
   set(CMAKE_CUDA_HOST_IMPLICIT_LINK_FRAMEWORK_DIRECTORIES "")
-
-  # Don't leak variables unnecessarily to user code.
-  unset(_CUDA_INCLUDE_DIR CACHE)
-  unset(_CUDA_NVCC_EXECUTABLE CACHE)
-  unset(_CUDA_LIBRARY_DIR)
-  unset(_CUDA_TARGET_DIR)
-  unset(_CUDA_TARGET_NAME)
 elseif(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA")
   set(_nvcc_log "")
   string(REPLACE "\r" "" _nvcc_output_orig "${CMAKE_CUDA_COMPILER_PRODUCED_OUTPUT}")
@@ -477,15 +480,6 @@ else()
   set(_SET_CMAKE_CUDA_COMPILER_SYSROOT "")
 endif()
 
-if(CMAKE_CUDA_COMPILER_ID STREQUAL "Clang")
-  string(APPEND _SET_CMAKE_CUDA_COMPILER_SYSROOT
-    "\nset(CMAKE_CUDA_COMPILER_TOOLKIT_ROOT \"${_CUDA_ROOT_DIR}\")\n"
-    "set(CMAKE_CUDA_COMPILER_LIBRARY_ROOT \"${_CUDA_LIBRARY_ROOT}\")")
-
-  unset(_CUDA_LIBRARY_ROOT)
-  unset(_CUDA_ROOT_DIR)
-endif()
-
 # Determine CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES
 if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA")
   set(CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES)
@@ -561,7 +555,14 @@ endif()
 configure_file(${CMAKE_ROOT}/Modules/CMakeCUDACompiler.cmake.in
   ${CMAKE_PLATFORM_INFO_DIR}/CMakeCUDACompiler.cmake
   @ONLY
-  )
+)
+
+# Don't leak variables unnecessarily to user code.
+unset(_CUDA_INCLUDE_DIR CACHE)
+unset(_CUDA_NVCC_EXECUTABLE CACHE)
+unset(_CUDA_LIBRARY_DIR)
+unset(_CUDA_TARGET_DIR)
+unset(_CUDA_TARGET_NAME)
 
 set(CMAKE_CUDA_COMPILER_ENV_VAR "CUDACXX")
 set(CMAKE_CUDA_HOST_COMPILER_ENV_VAR "CUDAHOSTCXX")
diff --git a/Modules/FindCUDAToolkit.cmake b/Modules/FindCUDAToolkit.cmake
index 0c8a441..00d1a50 100644
--- a/Modules/FindCUDAToolkit.cmake
+++ b/Modules/FindCUDAToolkit.cmake
@@ -477,7 +477,7 @@ Result variables
 #
 ###############################################################################
 
-# On Clang the toolkit is found during compiler detection and stored in CMakeCUDACompiler.cmake as
+# The toolkit is located during compiler detection for CUDA and stored in CMakeCUDACompiler.cmake as
 # CMAKE_CUDA_COMPILER_TOOLKIT_ROOT and CMAKE_CUDA_COMPILER_LIBRARY_ROOT.
 # We compute the rest based on those here to avoid re-searching and to avoid finding a possibly
 # different installation.