WIP: Prep for libcudacxx dependency.

NVIDIA · Feb 5, 2022 · d51b751 · d51b751
1 parent d40f50e
commit d51b751
Show file tree

Hide file tree

Showing 6 changed files with 138 additions and 58 deletions.
diff --git a/cmake/ThrustInstallRules.cmake b/cmake/ThrustInstallRules.cmake
@@ -23,26 +23,35 @@ configure_file("${Thrust_SOURCE_DIR}/thrust/cmake/thrust-header-search.cmake.in"
 install(FILES "${Thrust_BINARY_DIR}/thrust/cmake/thrust-header-search.cmake"
   DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/thrust")
 
-# Depending on how Thrust is configured, CUB's CMake scripts may or may not be
-# included, so maintain a set of CUB install rules in both projects. By default
-# CUB headers are installed alongside Thrust -- this may be disabled by turning
-# off THRUST_INSTALL_CUB_HEADERS.
-option(THRUST_INSTALL_CUB_HEADERS "Include cub headers when installing." ON)
+# Depending on how Thrust is configured, libcudacxx and CUB's CMake scripts may
+# or may not be include()'d, so force include their install rules when requested.
+# By default, these projects are installed alongside Thrust. This is controlled by
+# THRUST_INSTALL_CUB_HEADERS and THRUST_INSTALL_LIBCUDACXX_HEADERS.
+option(THRUST_INSTALL_CUB_HEADERS "Include CUB headers when installing." ON)
 if (THRUST_INSTALL_CUB_HEADERS)
-  install(DIRECTORY "${Thrust_SOURCE_DIR}/dependencies/cub/cub"
-    DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}"
-    FILES_MATCHING
-      PATTERN "*.cuh"
-  )
+  # Use a function to limit scope of the CUB_*_DIR vars:
+  function(_thrust_install_cub_headers)
+    # Fake these for the logic in CUBInstallRules.cmake:
+    set(CUB_SOURCE_DIR "${Thrust_SOURCE_DIR}/dependencies/cub/")
+    set(CUB_BINARY_DIR "${Thrust_BINARY_DIR}/cub-config/")
+    set(CUB_ENABLE_INSTALL_RULES ON)
+    set(CUB_IN_THRUST OFF)
+    include("${Thrust_SOURCE_DIR}/dependencies/cub/cmake/CubInstallRules.cmake")
+  endfunction()
 
-  # Need to configure a file to store THRUST_INSTALL_HEADER_INFIX
-  install(DIRECTORY "${Thrust_SOURCE_DIR}/dependencies/cub/cub/cmake/"
-    DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/cub"
-    PATTERN cub-header-search EXCLUDE
-  )
-  configure_file("${Thrust_SOURCE_DIR}/dependencies/cub/cub/cmake/cub-header-search.cmake.in"
-    "${Thrust_BINARY_DIR}/dependencies/cub/cub/cmake/cub-header-search.cmake"
-    @ONLY)
-  install(FILES "${Thrust_BINARY_DIR}/dependencies/cub/cub/cmake/cub-header-search.cmake"
-    DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/cub")
+  _thrust_install_cub_headers()
+endif()
+
+option(THRUST_INSTALL_LIBCUDACXX_HEADERS "Include libcudacxx headers when installing." ON)
+if (THRUST_INSTALL_LIBCUDACXX_HEADERS)
+  # Use a function to limit scope of the libcudacxx_*_DIR vars:
+  function(_thrust_install_libcudacxx_headers)
+    # Fake these for the logic in libcudacxxInstallRules.cmake:
+    set(libcudacxx_SOURCE_DIR "${Thrust_SOURCE_DIR}/dependencies/libcudacxx/")
+    set(libcudacxx_BINARY_DIR "${Thrust_BINARY_DIR}/libcudacxx-config/")
+    set(libcudacxx_ENABLE_INSTALL_RULES ON)
+    include("${Thrust_SOURCE_DIR}/dependencies/libcudacxx/cmake/libcudacxxInstallRules.cmake")
+  endfunction()
+
+  _thrust_install_libcudacxx_headers()
 endif()
diff --git a/dependencies/cub b/dependencies/cub
diff --git a/dependencies/libcudacxx b/dependencies/libcudacxx
diff --git a/thrust/cmake/thrust-config.cmake b/thrust/cmake/thrust-config.cmake
@@ -37,14 +37,15 @@
 #   [ADVANCED]                       # Optionally mark options as advanced
 # )
 #
-# # Use a custom TBB, CUB, and/or OMP
+# # Use a custom TBB, CUB, libcudacxx, and/or OMP
 # # (Note that once set, these cannot be changed. This includes COMPONENT
 # # preloading and lazy lookups in thrust_create_target)
 # find_package(Thrust REQUIRED)
 # thrust_set_CUB_target(MyCUBTarget)  # MyXXXTarget contains an existing
 # thrust_set_TBB_target(MyTBBTarget)  # interface to XXX for Thrust to use.
+# thrust_set_libcudacxx_target(MyLibcudacxxTarget)
 # thrust_set_OMP_target(MyOMPTarget)
-# thrust_create_target(ThrustWithMyCUB DEVICE CUDA)
+# thrust_create_target(ThrustWithMyCUBAndLibcudacxx DEVICE CUDA)
 # thrust_create_target(ThrustWithMyTBB DEVICE TBB)
 # thrust_create_target(ThrustWithMyOMP DEVICE OMP)
 #
@@ -77,6 +78,9 @@
 
 cmake_minimum_required(VERSION 3.15)
 
+# Minimum libcudacxx version:
+set(thrust_libcudacxx_version 999.0.0) # FIXME: Change this to a real version once we have one that works.
+
 ################################################################################
 # User variables and APIs. Users can rely on these:
 #
@@ -85,19 +89,21 @@ cmake_minimum_required(VERSION 3.15)
 set(THRUST_HOST_SYSTEM_OPTIONS
   CPP OMP TBB
   CACHE INTERNAL "Valid Thrust host systems."
+  FORCE
 )
 set(THRUST_DEVICE_SYSTEM_OPTIONS
   CUDA CPP OMP TBB
   CACHE INTERNAL "Valid Thrust device systems"
+  FORCE
 )
 
 # Workaround cmake issue #20670 https://gitlab.kitware.com/cmake/cmake/-/issues/20670
-set(THRUST_VERSION ${${CMAKE_FIND_PACKAGE_NAME}_VERSION} CACHE INTERNAL "")
-set(THRUST_VERSION_MAJOR ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_MAJOR} CACHE INTERNAL "")
-set(THRUST_VERSION_MINOR ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_MINOR} CACHE INTERNAL "")
-set(THRUST_VERSION_PATCH ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_PATCH} CACHE INTERNAL "")
-set(THRUST_VERSION_TWEAK ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_TWEAK} CACHE INTERNAL "")
-set(THRUST_VERSION_COUNT ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_COUNT} CACHE INTERNAL "")
+set(THRUST_VERSION ${${CMAKE_FIND_PACKAGE_NAME}_VERSION} CACHE INTERNAL "" FORCE)
+set(THRUST_VERSION_MAJOR ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_MAJOR} CACHE INTERNAL "" FORCE)
+set(THRUST_VERSION_MINOR ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_MINOR} CACHE INTERNAL "" FORCE)
+set(THRUST_VERSION_PATCH ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_PATCH} CACHE INTERNAL "" FORCE)
+set(THRUST_VERSION_TWEAK ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_TWEAK} CACHE INTERNAL "" FORCE)
+set(THRUST_VERSION_COUNT ${${CMAKE_FIND_PACKAGE_NAME}_VERSION_COUNT} CACHE INTERNAL "" FORCE)
 
 function(thrust_create_target target_name)
   thrust_debug("Assembling target ${target_name}. Options: ${ARGN}" internal)
@@ -109,21 +115,21 @@ function(thrust_create_target target_name)
     IGNORE_DEPRECATED_COMPILER
     IGNORE_DEPRECATED_CPP_11
     IGNORE_DEPRECATED_CPP_DIALECT
-    )
+  )
   set(keys
     DEVICE
     DEVICE_OPTION
     DEVICE_OPTION_DOC
     HOST
     HOST_OPTION
     HOST_OPTION_DOC
-    )
+  )
   cmake_parse_arguments(TCT "${options}" "${keys}" "" ${ARGN})
   if (TCT_UNPARSED_ARGUMENTS)
     message(AUTHOR_WARNING
       "Unrecognized arguments passed to thrust_create_target: "
       ${TCT_UNPARSED_ARGUMENTS}
-      )
+    )
   endif()
 
   # Check that the main Thrust internal target is available
@@ -133,7 +139,7 @@ function(thrust_create_target target_name)
     message(AUTHOR_WARNING
       "The `thrust_create_target` function was called outside the scope of the "
       "thrust targets. Call find_package again to recreate targets."
-      )
+    )
   endif()
 
   _thrust_set_if_undefined(TCT_HOST CPP)
@@ -145,12 +151,14 @@ function(thrust_create_target target_name)
 
   if (NOT TCT_HOST IN_LIST THRUST_HOST_SYSTEM_OPTIONS)
     message(FATAL_ERROR
-      "Requested HOST=${TCT_HOST}; must be one of ${THRUST_HOST_SYSTEM_OPTIONS}")
+      "Requested HOST=${TCT_HOST}; must be one of ${THRUST_HOST_SYSTEM_OPTIONS}"
+    )
   endif()
 
   if (NOT TCT_DEVICE IN_LIST THRUST_DEVICE_SYSTEM_OPTIONS)
     message(FATAL_ERROR
-      "Requested DEVICE=${TCT_DEVICE}; must be one of ${THRUST_DEVICE_SYSTEM_OPTIONS}")
+      "Requested DEVICE=${TCT_DEVICE}; must be one of ${THRUST_DEVICE_SYSTEM_OPTIONS}"
+    )
   endif()
 
   if (TCT_FROM_OPTIONS)
@@ -172,7 +180,7 @@ function(thrust_create_target target_name)
 
   # We can just create an INTERFACE IMPORTED target here instead of going
   # through _thrust_declare_interface_alias as long as we aren't hanging any
-  # Thrust/CUB include paths on ${target_name}.
+  # Thrust/CUB include paths directly on ${target_name}.
   add_library(${target_name} INTERFACE IMPORTED)
   target_link_libraries(${target_name}
     INTERFACE
@@ -346,14 +354,17 @@ function(thrust_debug_internal_targets)
 
   _thrust_debug_backend_targets(CPP "Thrust ${THRUST_VERSION}")
 
-  _thrust_debug_backend_targets(CUDA "CUB ${THRUST_CUB_VERSION}")
-  thrust_debug_target(CUB::CUB "${THRUST_CUB_VERSION}")
+  _thrust_debug_backend_targets(OMP "${THRUST_OMP_VERSION}")
+  thrust_debug_target(OpenMP::OpenMP_CXX "${THRUST_OMP_VERSION}")
 
   _thrust_debug_backend_targets(TBB "${THRUST_TBB_VERSION}")
   thrust_debug_target(TBB:tbb "${THRUST_TBB_VERSION}")
 
-  _thrust_debug_backend_targets(OMP "${THRUST_OMP_VERSION}")
-  thrust_debug_target(OpenMP::OpenMP_CXX "${THRUST_OMP_VERSION}")
+  _thrust_debug_backend_targets(CUDA
+    "CUB ${THRUST_CUB_VERSION} libcudacxx ${THRUST_libcudacxx_VERSION}"
+  )
+  thrust_debug_target(CUB::CUB "${THRUST_CUB_VERSION}")
+  thrust_debug_target(libcudacxx::libcudacxx "${THRUST_libcudacxx_VERSION}")
 endfunction()
 
 ################################################################################
@@ -434,18 +445,37 @@ function(_thrust_setup_system backend)
   endif()
 endfunction()
 
-# Use the provided cub_target for the CUDA backend. If Thrust::CUDA already
+# Use the provided cub_target for the CUDA backend. If Thrust::CUB already
 # exists, this call has no effect.
 function(thrust_set_CUB_target cub_target)
-  if (NOT TARGET Thrust::CUDA)
+  if (NOT TARGET Thrust::CUB)
     thrust_debug("Setting CUB target to ${cub_target}" internal)
     # Workaround cmake issue #20670 https://gitlab.kitware.com/cmake/cmake/-/issues/20670
-    set(THRUST_CUB_VERSION ${CUB_VERSION} CACHE INTERNAL "CUB version used by Thrust")
-    _thrust_declare_interface_alias(Thrust::CUDA _Thrust_CUDA)
-    target_link_libraries(_Thrust_CUDA INTERFACE Thrust::Thrust ${cub_target})
+    set(THRUST_CUB_VERSION ${CUB_VERSION} CACHE INTERNAL
+      "CUB version used by Thrust"
+      FORCE
+    )
+    _thrust_declare_interface_alias(Thrust::CUB _Thrust_CUB)
+    target_link_libraries(_Thrust_CUB INTERFACE ${cub_target})
     thrust_debug_target(${cub_target} "${THRUST_CUB_VERSION}" internal)
-    thrust_debug_target(Thrust::CUDA "CUB ${THRUST_CUB_VERSION}" internal)
-    _thrust_setup_system(CUDA)
+    thrust_debug_target(Thrust::CUB "CUB ${THRUST_CUB_VERSION}" internal)
+  endif()
+endfunction()
+
+# Use the provided libcudacxx_target for the CUDA backend. If Thrust::libcudacxx
+# already exists, this call has no effect.
+function(thrust_set_libcudacxx_target libcudacxx_target)
+  if (NOT TARGET Thrust::libcudacxx)
+    thrust_debug("Setting libcudacxx target to ${libcudacxx_target}" internal)
+    # Workaround cmake issue #20670 https://gitlab.kitware.com/cmake/cmake/-/issues/20670
+    set(THRUST_libcudacxx_VERSION ${libcudacxx_VERSION} CACHE INTERNAL
+      "libcudacxx version used by Thrust"
+      FORCE
+    )
+    _thrust_declare_interface_alias(Thrust::libcudacxx _Thrust_libcudacxx)
+    target_link_libraries(_Thrust_libcudacxx INTERFACE ${libcudacxx_target})
+    thrust_debug_target(${libcudacxx_target} "${THRUST_libcudacxx_VERSION}" internal)
+    thrust_debug_target(Thrust::libcudacxx "libcudacxx ${THRUST_libcudacxx_VERSION}" internal)
   endif()
 endfunction()
 
@@ -455,7 +485,10 @@ function(thrust_set_TBB_target tbb_target)
   if (NOT TARGET Thrust::TBB)
     thrust_debug("Setting TBB target to ${tbb_target}" internal)
     # Workaround cmake issue #20670 https://gitlab.kitware.com/cmake/cmake/-/issues/20670
-    set(THRUST_TBB_VERSION ${TBB_VERSION} CACHE INTERNAL "TBB version used by Thrust")
+    set(THRUST_TBB_VERSION ${TBB_VERSION} CACHE INTERNAL
+      "TBB version used by Thrust"
+      FORCE
+    )
     _thrust_declare_interface_alias(Thrust::TBB _Thrust_TBB)
     target_link_libraries(_Thrust_TBB INTERFACE Thrust::Thrust ${tbb_target})
     thrust_debug_target(${tbb_target} "${THRUST_TBB_VERSION}" internal)
@@ -470,7 +503,10 @@ function(thrust_set_OMP_target omp_target)
   if (NOT TARGET Thrust::OMP)
     thrust_debug("Setting OMP target to ${omp_target}" internal)
     # Workaround cmake issue #20670 https://gitlab.kitware.com/cmake/cmake/-/issues/20670
-    set(THRUST_OMP_VERSION ${OpenMP_CXX_VERSION} CACHE INTERNAL "OpenMP version used by Thrust")
+    set(THRUST_OMP_VERSION ${OpenMP_CXX_VERSION} CACHE INTERNAL
+      "OpenMP version used by Thrust"
+      FORCE
+    )
     _thrust_declare_interface_alias(Thrust::OMP _Thrust_OMP)
     target_link_libraries(_Thrust_OMP INTERFACE Thrust::Thrust ${omp_target})
     thrust_debug_target(${omp_target} "${THRUST_OMP_VERSION}" internal)
@@ -495,7 +531,27 @@ endfunction()
 # #20670 -- otherwise variables like CUB_VERSION, etc won't be in the caller's
 # scope.
 macro(_thrust_find_CUDA required)
-  if (NOT TARGET Thrust::CUDA)
+  # Find libcudacxx first to ensure that CUB finds the same package.
+  if (NOT TARGET Thrust::libcudacxx)
+    thrust_debug("Searching for libcudacxx ${required}" internal)
+    find_package(libcudacxx 999.0.0 CONFIG # FIXME: Put this at the real version once ready
+      ${_THRUST_QUIET_FLAG}
+      ${required}
+      NO_DEFAULT_PATH # Only check the explicit HINTS below:
+      HINTS
+        "${_THRUST_INCLUDE_DIR}/dependencies/libcudacxx" # Source layout (GitHub)
+        "${_THRUST_INCLUDE_DIR}/../libcudacxx"           # Source layout (Perforce)
+        "${_THRUST_CMAKE_DIR}/.."                        # Install layout
+    )
+
+    if (TARGET libcudacxx::libcudacxx)
+      thrust_set_libcudacxx_target(libcudacxx::libcudacxx)
+    else()
+      thrust_debug("libcudacxx not found!" internal)
+    endif()
+  endif()
+
+  if (NOT TARGET Thrust::CUB)
     thrust_debug("Searching for CUB ${required}" internal)
     find_package(CUB ${THRUST_VERSION} CONFIG
       ${_THRUST_QUIET_FLAG}
@@ -513,6 +569,17 @@ macro(_thrust_find_CUDA required)
       thrust_debug("CUB not found!" internal)
     endif()
   endif()
+
+  if (NOT TARGET Thrust::CUDA)
+    _thrust_declare_interface_alias(Thrust::CUDA _Thrust_CUDA)
+    _thrust_setup_system(CUDA)
+    target_link_libraries(_Thrust_CUDA INTERFACE
+      Thrust::Thrust
+      Thrust::CUB
+      Thrust::libcudacxx
+    )
+    thrust_debug_target(Thrust::CUDA "" internal)
+  endif()
 endmacro()
 
 # This must be a macro instead of a function to ensure that backends passed to
@@ -619,21 +686,25 @@ endmacro()
 #
 
 if (${CMAKE_FIND_PACKAGE_NAME}_FIND_QUIETLY)
-  set(_THRUST_QUIET ON CACHE INTERNAL "Quiet mode enabled for Thrust find_package calls.")
-  set(_THRUST_QUIET_FLAG "QUIET" CACHE INTERNAL "")
+  set(_THRUST_QUIET ON CACHE INTERNAL "Quiet mode enabled for Thrust find_package calls." FORCE)
+  set(_THRUST_QUIET_FLAG "QUIET" CACHE INTERNAL "" FORCE)
 else()
   unset(_THRUST_QUIET CACHE)
   unset(_THRUST_QUIET_FLAG CACHE)
 endif()
 
-set(_THRUST_CMAKE_DIR "${CMAKE_CURRENT_LIST_DIR}" CACHE INTERNAL "Location of thrust-config.cmake")
+set(_THRUST_CMAKE_DIR "${CMAKE_CURRENT_LIST_DIR}" CACHE INTERNAL
+  "Location of thrust-config.cmake"
+  FORCE
+)
 
 # Internal target that actually holds the Thrust interface. Used by all other Thrust targets.
 if (NOT TARGET Thrust::Thrust)
   _thrust_declare_interface_alias(Thrust::Thrust _Thrust_Thrust)
   # Pull in the include dir detected by thrust-config-version.cmake
   set(_THRUST_INCLUDE_DIR "${_THRUST_VERSION_INCLUDE_DIR}"
     CACHE INTERNAL "Location of Thrust headers."
+    FORCE
   )
   unset(_THRUST_VERSION_INCLUDE_DIR CACHE) # Clear tmp variable from cache
   target_include_directories(_Thrust_Thrust INTERFACE "${_THRUST_INCLUDE_DIR}")

diff --git a/thrust/system/cuda/detail/core/agent_launcher.h b/thrust/system/cuda/detail/core/agent_launcher.h
@@ -514,9 +514,9 @@ namespace core {
     {
       if (debug_sync)
       {
-        NV_IF_TARGET(NV_IS_DEVICE,
-                     (cub::detail::device_synchronize();),
-                     (cudaStreamSynchronize(stream);));
+        NV_IF_TARGET(NV_IS_HOST,
+		     (cudaStreamSynchronize(stream);),
+                    (cub::detail::device_synchronize();));
       }
     }
 

diff --git a/thrust/system/cuda/detail/util.h b/thrust/system/cuda/detail/util.h
@@ -248,7 +248,7 @@ trivial_copy_device_to_device(Policy &    policy,
 inline void __host__ __device__
 terminate()
 {
-  NV_IF_TARGET(NV_IS_DEVICE, (asm("trap;");), (std::terminate();));
+  NV_IF_TARGET(NV_IS_HOST, (std::terminate();), (asm("trap;");));
 }
 
 __host__  __device__