CUDA: Fix CMAKE_CUDA_COMPILER_TOOLKIT_ROOT detection in Visual Studio

When using a Visual Studio generator, we do not know the path to
`CMAKE_CUDA_COMPILER` until after the compiler id detection project
extracts it.  However, based on the VS toolchain configuration we know
that the compiler id will be NVIDIA, so use that to drive the CUDA
architecture flag logic.  After the main compiler id step we know
`CMAKE_CUDA_COMPILER`, so `CMAKE_CUDA_COMPILER_TOOLKIT_ROOT` can then be
computed.

1 files changed, 121 insertions, 108 deletions

diff --git a/Modules/CMakeDetermineCUDACompiler.cmake b/Modules/CMakeDetermineCUDACompiler.cmake
index 15eebf7..692ce20 100644
--- a/Modules/CMakeDetermineCUDACompiler.cmake
+++ b/Modules/CMakeDetermineCUDACompiler.cmake
@@ -53,128 +53,134 @@ if(NOT CMAKE_CUDA_COMPILER_ID_RUN)
 
   include(${CMAKE_ROOT}/Modules/CMakeDetermineCompilerId.cmake)
 
-  # We determine the vendor to help with find the toolkit and use the right flags for detection right away.
-  # The main compiler identification is still needed below to extract other information.
-  list(APPEND CMAKE_CUDA_COMPILER_ID_VENDORS NVIDIA Clang)
-  set(CMAKE_CUDA_COMPILER_ID_VENDOR_REGEX_NVIDIA "nvcc: NVIDIA \\(R\\) Cuda compiler driver")
-  set(CMAKE_CUDA_COMPILER_ID_VENDOR_REGEX_Clang "(clang version)")
-  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(_CUDA_NVCC_EXECUTABLE "${CMAKE_CUDA_COMPILER}")
+  if(${CMAKE_GENERATOR} MATCHES "Visual Studio")
+    # We will not know CMAKE_CUDA_COMPILER until the main compiler id step
+    # below extracts it, but we do know that the compiler id will be NVIDIA.
+    set(CMAKE_CUDA_COMPILER_ID "NVIDIA")
   else()
-    # 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(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()
+    # We determine the vendor to help with find the toolkit and use the right flags for detection right away.
+    # The main compiler identification is still needed below to extract other information.
+    list(APPEND CMAKE_CUDA_COMPILER_ID_VENDORS NVIDIA Clang)
+    set(CMAKE_CUDA_COMPILER_ID_VENDOR_REGEX_NVIDIA "nvcc: NVIDIA \\(R\\) Cuda compiler driver")
+    set(CMAKE_CUDA_COMPILER_ID_VENDOR_REGEX_Clang "(clang version)")
+    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(_CUDA_NVCC_EXECUTABLE "${CMAKE_CUDA_COMPILER}")
+    else()
+      # 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
+      )
 
-    # 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-")
+      # 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(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()
-      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})
+      # 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()
+        else()
+          set(platform_base "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v")
         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)
+        # Force the global default /usr/local/cuda to the front on Unix.
+        if(UNIX)
+          list(INSERT search_paths 0 "/usr/local/cuda")
+        endif()
 
-      if(NOT _CUDA_NVCC_EXECUTABLE)
-        message(FATAL_ERROR "Could not find nvcc, please set CUDAToolkit_ROOT.")
+        # 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()
     endif()
-  endif()
 
-  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)
-
-  # 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")
+    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)
+
+    # 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()
   endif()
 
   set(CMAKE_CUDA_COMPILER_ID_FLAGS_ALWAYS "-v")
@@ -236,6 +242,13 @@ if(NOT CMAKE_CUDA_COMPILER_ID_RUN)
 
   CMAKE_DETERMINE_COMPILER_ID(CUDA CUDAFLAGS CMakeCUDACompilerId.cu)
 
+  if(${CMAKE_GENERATOR} MATCHES "Visual Studio")
+    # Now that we have the path to nvcc, we can compute the toolkit root.
+    get_filename_component(CMAKE_CUDA_COMPILER_TOOLKIT_ROOT "${CMAKE_CUDA_COMPILER}" DIRECTORY)
+    get_filename_component(CMAKE_CUDA_COMPILER_TOOLKIT_ROOT "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}" DIRECTORY)
+    set(CMAKE_CUDA_COMPILER_LIBRARY_ROOT "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}")
+  endif()
+
   _cmake_find_compiler_sysroot(CUDA)
 endif()