Make the thrust dispatch mechanisms configurable

docs/thrust/cmake_options.rst

@@ -69,6 +69,19 @@ Generic CMake Options
    -  If true, installation rules will be generated for thrust. Default
       is ``ON``.
 
+-  ``THRUST_DISPATCH_TYPE={Dynamic, Force32bit, Force64bit}``
+
+   -  Allows the user to force Thrust to use a specific size for the offset type. Default
+      is ``Dynamic``.
+
+      -  ``Dynamic`` lets Thrust choose the index type based on input size, allowing
+         large inputs and optimal performance at the cost of increased compile time and binary size,
+         as Thrust will compile each kernel twice, once for 32 bit and once for 64 bit.
+      -  ``Force32bit`` forces Thrust to use a 32 bit offset type. This improves compile time and
+         binary size but limits the input size.
+      -  ``Force64bit`` forces Thrust to use a 64 bit offset type. This improves compile time and
+         binary size and allows large input sizes. However, it might degrade runtime performance.
+
 Single Config CMake Options
 ---------------------------
 

thrust/CMakeLists.txt

@@ -59,6 +59,10 @@ option(THRUST_ENABLE_TESTING "Build Thrust testing suite." "ON")
 option(THRUST_ENABLE_EXAMPLES "Build Thrust examples." "ON")
 option(THRUST_ENABLE_BENCHMARKS "Build Thrust runtime benchmarks." "${CCCL_ENABLE_BENCHMARKS}")
 
+# Allow the user to optionally select offset type dispatch to fixed 32 or 64 bit types
+set(THRUST_DISPATCH_TYPE "Dynamic" CACHE STRING "Select Thrust offset type dispatch." FORCE)
+set_property(CACHE THRUST_DISPATCH_TYPE PROPERTY STRINGS "Dynamic" "Force32bit" "Force64bit")
+
 # Check if we're actually building anything before continuing. If not, no need
 # to search for deps, etc. This is a common approach for packagers that just
 # need the install rules. See GH issue NVIDIA/thrust#1211.

thrust/cmake/ThrustBuildCompilerTargets.cmake

@@ -131,6 +131,12 @@ function(thrust_build_compiler_targets)
     )
   endforeach()
 
+  if (THRUST_DISPATCH_TYPE STREQUAL "Force32bit")
+    list(APPEND cxx_compile_definitions "THRUST_FORCE_32_BIT_OFFSET_TYPE")
+  elseif (THRUST_DISPATCH_TYPE STREQUAL "Force64bit")
+    list(APPEND cxx_compile_definitions "THRUST_FORCE_64_BIT_OFFSET_TYPE")
+  endif()
+
   foreach (cxx_definition IN LISTS cxx_compile_definitions)
     # Add these for both CUDA and CXX targets:
     target_compile_definitions(thrust.compiler_interface INTERFACE

thrust/cmake/ThrustHeaderTesting.cmake

@@ -146,6 +146,19 @@ foreach(thrust_target IN LISTS THRUST_TARGETS)
     "CUB_WRAPPED_NAMESPACE=wrapped_cub")
   thrust_add_header_test(${thrust_target} base "${header_definitions}")
 
+  # We need to ensure that the different dispatch mechanisms work
+  set(header_definitions
+    "THRUST_WRAPPED_NAMESPACE=wrapped_thrust"
+    "CUB_WRAPPED_NAMESPACE=wrapped_cub"
+    "THRUST_FORCE_32_BIT_OFFSET_TYPE")
+  thrust_add_header_test(${thrust_target} offset_32 "${header_definitions}")
+
+  set(header_definitions
+    "THRUST_WRAPPED_NAMESPACE=wrapped_thrust"
+    "CUB_WRAPPED_NAMESPACE=wrapped_cub"
+    "THRUST_FORCE_64_BIT_OFFSET_TYPE")
+  thrust_add_header_test(${thrust_target} offset_64 "${header_definitions}")
+
   thrust_get_target_property(config_device ${thrust_target} DEVICE)
   if ("CUDA" STREQUAL "${config_device}")
     # Check that BF16 support can be disabled