From 8994dc477a99a3d4c1c1ba78dd42e5357ef31947 Mon Sep 17 00:00:00 2001
From: Allison Piper <alliepiper16@gmail.com>
Date: Tue, 19 Nov 2024 14:07:40 -0500
Subject: [PATCH] Add `thrust_create_target` `DISPATCH` option. (#2844)

* Remove the THRUST_DISPATCH_TYPE header tests.

These will end up specifying conflicting flags when `THRUST_DISPATCH_TYPE` is set to something other than `Dynamic`.

* Add `DISPATCH` option to `thrust_create_target`.

```
thrust_create_target(TargetName
  DISPATCH [Dynamic|Force32bit|Force64bit]
```

* Skip 64-bit offset tests when forcing 32-bit dispatch.

* Add 32/64-bit dispatch jobs to nightly CI.

* Add 32-bit dispatch to pull_request workflow.
---
 ci/matrix.yaml                                |  6 ++-
 lib/cmake/thrust/README.md                    | 45 +++++++++++++----
 lib/cmake/thrust/thrust-config.cmake          | 50 ++++++++++++++++---
 thrust/CMakeLists.txt                         |  2 +-
 thrust/cmake/ThrustBuildCompilerTargets.cmake |  6 ---
 thrust/cmake/ThrustBuildTargetList.cmake      |  1 +
 thrust/cmake/ThrustHeaderTesting.cmake        |  7 ---
 thrust/testing/copy.cu                        |  4 ++
 thrust/testing/count.cu                       |  2 +
 thrust/testing/cuda/adjacent_difference.cu    |  2 +
 thrust/testing/cuda/partition.cu              |  3 ++
 thrust/testing/cuda/reduce_by_key.cu          |  2 +
 thrust/testing/cuda/sort.cu                   |  7 ++-
 thrust/testing/inner_product.cu               |  2 +
 thrust/testing/max_element.cu                 |  2 +
 thrust/testing/min_element.cu                 |  2 +
 thrust/testing/minmax_element.cu              |  2 +
 thrust/testing/reduce.cu                      |  2 +
 thrust/testing/scan.cu                        |  4 ++
 thrust/testing/set_difference.cu              |  3 ++
 thrust/testing/set_intersection.cu            |  2 +
 thrust/testing/unique_by_key.cu               |  4 ++
 22 files changed, 125 insertions(+), 35 deletions(-)

diff --git a/ci/matrix.yaml b/ci/matrix.yaml
index 2586fab09e6..ae5cbfee083 100644
--- a/ci/matrix.yaml
+++ b/ci/matrix.yaml
@@ -28,6 +28,8 @@ workflows:
     - {jobs: ['build'], std: 'all', ctk: '12.5', cxx: 'nvhpc'}
     - {jobs: ['build'], std: 'all', cxx: ['gcc', 'clang'], cpu: 'arm64'}
     - {jobs: ['build'], std: 'all', cxx: ['gcc'], sm: '90a'}
+    # Test Thrust 32-bit-only dispatch here, since it's most likely to break. 64-bit-only is tested in nightly.
+    - {jobs: ['test_gpu'],  project: 'thrust', cmake_options: '-DTHRUST_DISPATCH_TYPE=Force32bit'}
     # default_projects: clang-cuda
     - {jobs: ['build'], std: 'all', cudacxx: 'clang', cxx: 'clang'}
     - {jobs: ['build'], project: 'libcudacxx', std: 'all', cudacxx: 'clang', cxx: 'clang', sm: '90'}
@@ -58,13 +60,15 @@ workflows:
     - {jobs: ['infra'], project: 'cccl', ctk: 'curr', cxx: ['gcc',   'clang']}
 
   nightly:
-    # Increased test coverage compared to nightlies:
+    # Increased test coverage compared to pull_request:
     - {jobs: ['test'],  std: 'all', cxx: ['gcc13', 'clang18', 'msvc2022']}
     - {jobs: ['test'],  project: 'cudax', ctk: ['12.0', 'curr'], std: 'all', cxx: ['gcc12']}
     - {jobs: ['test'],  project: 'cudax', ctk: ['12.0'        ], std: 'all', cxx: ['clang14']}
     - {jobs: ['test'],  project: 'cudax', ctk: [        'curr'], std: 'all', cxx: ['clang18']}
     # Edge-case jobs
     - {jobs: ['limited'], project: 'cub', std: 17}
+    - {jobs: ['test_gpu'],  project: 'thrust', cmake_options: '-DTHRUST_DISPATCH_TYPE=Force32bit'}
+    - {jobs: ['test_gpu'],  project: 'thrust', cmake_options: '-DTHRUST_DISPATCH_TYPE=Force64bit'}
 
 #  # These are waiting on the NVKS nodes:
 #    - {jobs: ['test'],  ctk: '11.1', gpu: 'v100',     sm: 'gpu', cxx: 'gcc6',    std: [11]}
diff --git a/lib/cmake/thrust/README.md b/lib/cmake/thrust/README.md
index 09cbc6debb5..165be15662c 100644
--- a/lib/cmake/thrust/README.md
+++ b/lib/cmake/thrust/README.md
@@ -38,6 +38,22 @@ If using Thrust from the CCCL sources, this would be
 $ cmake . -DThrust_DIR=<CCCL git repo root>/thrust/thrust/cmake/
 ```
 
+#### Large Array (64-bit offseet) Handling: `DISPATCH`
+
+The `DISPATCH` option allows users to select the tradeoff of compile-time / binary-size
+vs. performance vs. scalability when given large inputs that require 64-bit offset types.
+This currently only applies when DEVICE=CUDA.
+
+- `Dynamic` May compile each kernel twice, once for 32-bit offsets and again for 64-bit
+  offsets, and choose dynamically using the input size at runtime.
+  This significantly increases compile-time and binary-size, but provides optimal performance
+  for small input sizes while also supporting 64-bit indexed workloads.
+- `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 may degrade runtime performance
+  for 32-bit indexed workloads.
+
 #### TBB / OpenMP
 
 To explicitly specify host/device systems, `HOST` and `DEVICE` arguments can be
@@ -56,17 +72,21 @@ host system, but will find and use TBB or OpenMP for the device system.
 
 To allow a Thrust target to be configurable easily via `cmake-gui` or
 `ccmake`, pass the `FROM_OPTIONS` flag to `thrust_create_target`. This will add
-`THRUST_HOST_SYSTEM` and `THRUST_DEVICE_SYSTEM` options to the CMake cache that
-allow selection from the systems supported by this version of Thrust.
+`THRUST_HOST_SYSTEM`, `THRUST_DEVICE_SYSTEM`, and `THRUST_DISPATCH_TYPE` options
+to the CMake cache that allow selection from the systems supported by this version
+of Thrust.
 
 ```cmake
 thrust_create_target(Thrust FROM_OPTIONS
   [HOST_OPTION <option name>]
   [DEVICE_OPTION <option name>]
+  [DISPATCH_OPTION <option name>]
   [HOST_OPTION_DOC <doc string>]
   [DEVICE_OPTION_DOC <doc string>]
+  [DISPATCH_OPTION_DOC <doc string>]
   [HOST <default host system name>]
   [DEVICE <default device system name>]
+  [DISPATCH <default dispatch type>]
   [ADVANCED]
 )
 ```
@@ -74,15 +94,18 @@ thrust_create_target(Thrust FROM_OPTIONS
 The optional arguments have sensible defaults, but may be configured per
 `thrust_create_target` call:
 
-| Argument            | Default                 | Description                     |
-|---------------------|-------------------------|---------------------------------|
-| `HOST_OPTION`       | `THRUST_HOST_SYSTEM`    | Name of cache option for host   |
-| `DEVICE_OPTION`     | `THRUST_DEVICE_SYSTEM`  | Name of cache option for device |
-| `HOST_OPTION_DOC`   | Thrust's host system.   | Docstring for host option       |
-| `DEVICE_OPTION_DOC` | Thrust's device system. | Docstring for device option     |
-| `HOST`              | `CPP`                   | Default host system             |
-| `DEVICE`            | `CUDA`                  | Default device system           |
-| `ADVANCED`          | *N/A*                   | Mark cache options advanced     |
+| Argument              | Default                 | Description                       |
+|-----------------------|-------------------------|-----------------------------------|
+| `HOST_OPTION`         | `THRUST_HOST_SYSTEM`    | Name of cache option for host     |
+| `DEVICE_OPTION`       | `THRUST_DEVICE_SYSTEM`  | Name of cache option for device   |
+| `DISPATCH_OPTION`     | `THRUST_DISPATCH_TYPE`  | Name of cache option for dispatch |
+| `HOST_OPTION_DOC`     | Thrust's host system.   | Docstring for host option         |
+| `DEVICE_OPTION_DOC`   | Thrust's device system. | Docstring for device option       |
+| `DISPATCH_OPTION_DOC` | Thrust's dispatch type. | Docstring for dispatch option     |
+| `HOST`                | `CPP`                   | Default host system               |
+| `DEVICE`              | `CUDA`                  | Default device system             |
+| `DISPATCH`            | `Dispatch`                  | Default dispatch type             |
+| `ADVANCED`            | *N/A*                   | Mark cache options advanced       |
 
 ### Specifying Thrust Version Requirements
 
diff --git a/lib/cmake/thrust/thrust-config.cmake b/lib/cmake/thrust/thrust-config.cmake
index ecfee61fa1c..dd63fd8f5b5 100644
--- a/lib/cmake/thrust/thrust-config.cmake
+++ b/lib/cmake/thrust/thrust-config.cmake
@@ -3,8 +3,6 @@
 #
 # Provided by NVIDIA under the same license as the associated Thrust library.
 #
-# Reply-To: Allison Vacanti <alliepiper16@gmail.com>
-#
 # *****************************************************************************
 # **     The following is a short reference to using Thrust from CMake.      **
 # ** For more details, see the README.md in the same directory as this file. **
@@ -30,10 +28,13 @@
 # thrust_create_target(TargetName FROM_OPTIONS
 #   [HOST_OPTION <option_name>]      # Optionally rename the host system option
 #   [DEVICE_OPTION <option_name>]    # Optionally rename the device system option
+#   [DISPATCH_OPTION <option_name>]  # Optionally rename the dispatch system option
 #   [HOST_OPTION_DOC <doc_string>]   # Optionally change the cache label
 #   [DEVICE_OPTION_DOC <doc_string>] # Optionally change the cache label
+#   [DISPATCH_OPTION_DOC <doc_str>]  # Optionally change the cache label
 #   [HOST <default system>]          # Optionally change the default backend
 #   [DEVICE <default system>]        # Optionally change the default backend
+#   [DISPATCH <default dispatch>]    # Optionally change the default dispatch
 #   [ADVANCED]                       # Optionally mark options as advanced
 #   [GLOBAL]                         # Optionally mark the target as GLOBAL
 # )
@@ -59,6 +60,11 @@
 #   IGNORE_CUB_VERSION            # Skip configure-time and compile-time CUB version checks
 # )
 #
+# # DISPATCH options (See README):
+# thrust_create_target(TargetName DISPATCH Dynamic)
+# thrust_create_target(TargetName DISPATCH Force32bit)
+# thrust_create_target(TargetName DISPATCH Force64bit)
+#
 # # Test if a particular system has been loaded. ${var_name} is set to TRUE or
 # # FALSE to indicate if "system" is found.
 # thrust_is_system_found(<system> <var_name>)
@@ -100,6 +106,11 @@ set(THRUST_DEVICE_SYSTEM_OPTIONS
   CACHE INTERNAL "Valid Thrust device systems"
   FORCE
 )
+set(THRUST_DISPATCH_TYPE_OPTIONS
+  Dynamic Force32bit Force64bit
+  CACHE INTERNAL "Valid Thrust dispatch types"
+  FORCE
+)
 
 # Workaround cmake issue #20670 https://gitlab.kitware.com/cmake/cmake/-/issues/20670
 # Legacy all-caps THRUST variables:
@@ -137,6 +148,9 @@ function(thrust_create_target target_name)
     HOST
     HOST_OPTION
     HOST_OPTION_DOC
+    DISPATCH
+    DISPATCH_OPTION
+    DISPATCH_OPTION_DOC
   )
   cmake_parse_arguments(TCT "${options}" "${keys}" "" ${ARGN})
   if (TCT_UNPARSED_ARGUMENTS)
@@ -158,10 +172,13 @@ function(thrust_create_target target_name)
 
   _thrust_set_if_undefined(TCT_HOST CPP)
   _thrust_set_if_undefined(TCT_DEVICE CUDA)
+  _thrust_set_if_undefined(TCT_DISPATCH Dynamic)
   _thrust_set_if_undefined(TCT_HOST_OPTION THRUST_HOST_SYSTEM)
   _thrust_set_if_undefined(TCT_DEVICE_OPTION THRUST_DEVICE_SYSTEM)
-  _thrust_set_if_undefined(TCT_HOST_OPTION_DOC "Thrust host system.")
-  _thrust_set_if_undefined(TCT_DEVICE_OPTION_DOC "Thrust device system.")
+  _thrust_set_if_undefined(TCT_DISPATCH_OPTION THRUST_DISPATCH_TYPE)
+  _thrust_set_if_undefined(TCT_HOST_OPTION_DOC "Thrust host system: ${THRUST_HOST_SYSTEM_OPTIONS}")
+  _thrust_set_if_undefined(TCT_DEVICE_OPTION_DOC "Thrust device system: ${THRUST_DEVICE_SYSTEM_OPTIONS}")
+  _thrust_set_if_undefined(TCT_DISPATCH_OPTION_DOC "Thrust dispatch type: ${THRUST_DISPATCH_TYPE_OPTIONS}")
 
   if (NOT TCT_HOST IN_LIST THRUST_HOST_SYSTEM_OPTIONS)
     message(FATAL_ERROR
@@ -175,18 +192,26 @@ function(thrust_create_target target_name)
     )
   endif()
 
+  if (NOT TCT_DISPATCH IN_LIST THRUST_DISPATCH_TYPE_OPTIONS)
+    message(FATAL_ERROR
+      "Requested DISPATCH=${TCT_DISPATCH}; must be one of ${THRUST_DISPATCH_TYPE_OPTIONS}"
+    )
+  endif()
+
   if (TCT_FROM_OPTIONS)
     _thrust_create_cache_options(
-      ${TCT_HOST} ${TCT_DEVICE}
-      ${TCT_HOST_OPTION} ${TCT_DEVICE_OPTION}
-      ${TCT_HOST_OPTION_DOC} ${TCT_DEVICE_OPTION_DOC}
+      ${TCT_HOST} ${TCT_DEVICE} ${TCT_DISPATCH}
+      ${TCT_HOST_OPTION} ${TCT_DEVICE_OPTION} ${TCT_DISPATCH_OPTION}
+      ${TCT_HOST_OPTION_DOC} ${TCT_DEVICE_OPTION_DOC} ${TCT_DISPATCH_OPTION_DOC}
       ${TCT_ADVANCED}
     )
     set(TCT_HOST ${${TCT_HOST_OPTION}})
     set(TCT_DEVICE ${${TCT_DEVICE_OPTION}})
+    set(TCT_DISPATCH ${${TCT_DISPATCH_OPTION}})
     thrust_debug("Current option settings:" internal)
     thrust_debug("  - ${TCT_HOST_OPTION}=${TCT_HOST}" internal)
     thrust_debug("  - ${TCT_DEVICE_OPTION}=${TCT_DEVICE}" internal)
+    thrust_debug("  - ${TCT_DISPATCH_OPTION}=${TCT_DISPATCH}" internal)
   endif()
 
   _thrust_find_backend(${TCT_HOST} REQUIRED)
@@ -206,6 +231,12 @@ function(thrust_create_target target_name)
     Thrust::${TCT_DEVICE}::Device
   )
 
+  if (${TCT_DISPATCH} STREQUAL "Force32bit")
+    target_compile_definitions(${target_name} INTERFACE "THRUST_FORCE_32_BIT_OFFSET_TYPE")
+  elseif(${TCT_DISPATCH} STREQUAL "Force64bit")
+    target_compile_definitions(${target_name} INTERFACE "THRUST_FORCE_64_BIT_OFFSET_TYPE")
+  endif()
+
   # This would be nice to enforce, but breaks when using old cmake + new
   # compiler, since cmake doesn't know what features the new compiler version
   # supports.
@@ -416,14 +447,17 @@ function(_thrust_declare_interface_alias alias_name ugly_name)
 endfunction()
 
 # Create cache options for selecting the user/device systems with ccmake/cmake-gui.
-function(_thrust_create_cache_options host device host_option device_option host_doc device_doc advanced)
+function(_thrust_create_cache_options host device dispatch host_option device_option dispatch_option host_doc device_doc dispatch_doc advanced)
   thrust_debug("Creating system cache options: (advanced=${advanced})" internal)
   thrust_debug("  - Host Option=${host_option} Default=${host} Doc='${host_doc}'" internal)
   thrust_debug("  - Device Option=${device_option} Default=${device} Doc='${device_doc}'" internal)
+  thrust_debug("  - Dispatch Option=${dispatch_option} Default=${dispatch} Doc='${dispatch_doc}'" internal)
   set(${host_option} ${host} CACHE STRING "${host_doc}")
   set_property(CACHE ${host_option} PROPERTY STRINGS ${THRUST_HOST_SYSTEM_OPTIONS})
   set(${device_option} ${device} CACHE STRING "${device_doc}")
   set_property(CACHE ${device_option} PROPERTY STRINGS ${THRUST_DEVICE_SYSTEM_OPTIONS})
+  set(${dispatch_option} ${dispatch} CACHE STRING "${dispatch_doc}")
+  set_property(CACHE ${dispatch_option} PROPERTY STRINGS ${THRUST_DISPATCH_TYPE_OPTIONS})
   if (advanced)
     mark_as_advanced(${host_option} ${device_option})
   endif()
diff --git a/thrust/CMakeLists.txt b/thrust/CMakeLists.txt
index dbd265795dd..64e6ba2ef3f 100644
--- a/thrust/CMakeLists.txt
+++ b/thrust/CMakeLists.txt
@@ -30,7 +30,7 @@ option(THRUST_ENABLE_TESTING "Build Thrust testing suite." "ON")
 option(THRUST_ENABLE_EXAMPLES "Build Thrust examples." "ON")
 
 # 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(THRUST_DISPATCH_TYPE "Dynamic" CACHE STRING "Select Thrust offset type dispatch.")
 set_property(CACHE THRUST_DISPATCH_TYPE PROPERTY STRINGS "Dynamic" "Force32bit" "Force64bit")
 
 # Check if we're actually building anything before continuing. If not, no need
diff --git a/thrust/cmake/ThrustBuildCompilerTargets.cmake b/thrust/cmake/ThrustBuildCompilerTargets.cmake
index be8f40cb7f8..71d0c5dc51e 100644
--- a/thrust/cmake/ThrustBuildCompilerTargets.cmake
+++ b/thrust/cmake/ThrustBuildCompilerTargets.cmake
@@ -26,12 +26,6 @@ function(thrust_build_compiler_targets)
     append_option_if_available("/wd4146" cxx_compile_options)
   endif()
 
-  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()
-
   cccl_build_compiler_interface(thrust.compiler_interface
     "${cuda_compile_options}"
     "${cxx_compile_options}"
diff --git a/thrust/cmake/ThrustBuildTargetList.cmake b/thrust/cmake/ThrustBuildTargetList.cmake
index f1e9cbcaf2b..c5aea28992b 100644
--- a/thrust/cmake/ThrustBuildTargetList.cmake
+++ b/thrust/cmake/ThrustBuildTargetList.cmake
@@ -206,6 +206,7 @@ function(_thrust_build_target_list_multiconfig)
           thrust_create_target(${target_name}
             HOST ${host}
             DEVICE ${device}
+            DISPATCH ${THRUST_DISPATCH_TYPE}
             ${THRUST_TARGET_FLAGS}
           )
 
diff --git a/thrust/cmake/ThrustHeaderTesting.cmake b/thrust/cmake/ThrustHeaderTesting.cmake
index d90521c8d7d..ea041350e38 100644
--- a/thrust/cmake/ThrustHeaderTesting.cmake
+++ b/thrust/cmake/ThrustHeaderTesting.cmake
@@ -154,13 +154,6 @@ foreach(thrust_target IN LISTS THRUST_TARGETS)
     "CUB_WRAPPED_NAMESPACE=wrapped_cub")
   thrust_add_header_test(${thrust_target} wrap "${header_definitions}")
 
-  # We need to ensure that the different dispatch mechanisms work
-  set(header_definitions "THRUST_FORCE_32_BIT_OFFSET_TYPE")
-  thrust_add_header_test(${thrust_target} offset_32 "${header_definitions}")
-
-  set(header_definitions "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
diff --git a/thrust/testing/copy.cu b/thrust/testing/copy.cu
index dc44affc109..d4808142b1b 100644
--- a/thrust/testing/copy.cu
+++ b/thrust/testing/copy.cu
@@ -668,6 +668,8 @@ void TestCopyIfStencilDispatchImplicit()
 }
 DECLARE_UNITTEST(TestCopyIfStencilDispatchImplicit);
 
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
+
 struct only_set_when_expected_it
 {
   long long expected;
@@ -752,3 +754,5 @@ void TestCopyWithBigIndexes()
   TestCopyWithBigIndexesHelper(33);
 }
 DECLARE_UNITTEST(TestCopyWithBigIndexes);
+
+#endif
diff --git a/thrust/testing/count.cu b/thrust/testing/count.cu
index 129790b401e..9a994822c8d 100644
--- a/thrust/testing/count.cu
+++ b/thrust/testing/count.cu
@@ -119,8 +119,10 @@ void TestCountWithBigIndexesHelper(int magnitude)
 void TestCountWithBigIndexes()
 {
   TestCountWithBigIndexesHelper(30);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestCountWithBigIndexesHelper(31);
   TestCountWithBigIndexesHelper(32);
   TestCountWithBigIndexesHelper(33);
+#endif
 }
 DECLARE_UNITTEST(TestCountWithBigIndexes);
diff --git a/thrust/testing/cuda/adjacent_difference.cu b/thrust/testing/cuda/adjacent_difference.cu
index 1da14c07a6c..4b1eba4aba4 100644
--- a/thrust/testing/cuda/adjacent_difference.cu
+++ b/thrust/testing/cuda/adjacent_difference.cu
@@ -153,8 +153,10 @@ void TestAdjacentDifferenceWithBigIndexesHelper(int magnitude)
 void TestAdjacentDifferenceWithBigIndexes()
 {
   TestAdjacentDifferenceWithBigIndexesHelper(30);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestAdjacentDifferenceWithBigIndexesHelper(31);
   TestAdjacentDifferenceWithBigIndexesHelper(32);
   TestAdjacentDifferenceWithBigIndexesHelper(33);
+#endif
 }
 DECLARE_UNITTEST(TestAdjacentDifferenceWithBigIndexes);
diff --git a/thrust/testing/cuda/partition.cu b/thrust/testing/cuda/partition.cu
index 57199d03625..de2ba691dd1 100644
--- a/thrust/testing/cuda/partition.cu
+++ b/thrust/testing/cuda/partition.cu
@@ -677,9 +677,12 @@ void TestPartitionIfWithMagnitude(int magnitude)
 void TestPartitionIfWithLargeNumberOfItems()
 {
   TestPartitionIfWithMagnitude(30);
+  // These require 64-bit dispatches even when magnitude < 32.
+#  ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestPartitionIfWithMagnitude(31);
   TestPartitionIfWithMagnitude(32);
   TestPartitionIfWithMagnitude(33);
+#  endif
 }
 DECLARE_UNITTEST(TestPartitionIfWithLargeNumberOfItems);
 #endif
diff --git a/thrust/testing/cuda/reduce_by_key.cu b/thrust/testing/cuda/reduce_by_key.cu
index f6efd4788a3..6d160a6920c 100644
--- a/thrust/testing/cuda/reduce_by_key.cu
+++ b/thrust/testing/cuda/reduce_by_key.cu
@@ -433,8 +433,10 @@ void TestReduceByKeyWithBigIndexesHelper(int magnitude)
 void TestReduceByKeyWithBigIndexes()
 {
   TestReduceByKeyWithBigIndexesHelper(30);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestReduceByKeyWithBigIndexesHelper(31);
   TestReduceByKeyWithBigIndexesHelper(32);
   TestReduceByKeyWithBigIndexesHelper(33);
+#endif
 }
 DECLARE_UNITTEST(TestReduceByKeyWithBigIndexes);
diff --git a/thrust/testing/cuda/sort.cu b/thrust/testing/cuda/sort.cu
index 876fbb0d31e..db5e67eb202 100644
--- a/thrust/testing/cuda/sort.cu
+++ b/thrust/testing/cuda/sort.cu
@@ -251,9 +251,14 @@ void TestSortWithMagnitude(int magnitude)
 
 void TestSortWithLargeNumberOfItems()
 {
-  TestSortWithMagnitude(39);
+  TestSortWithMagnitude(30);
+  // These still require 64-bit dispatches when magnitude < 32.
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
+  TestSortWithMagnitude(31);
   TestSortWithMagnitude(32);
   TestSortWithMagnitude(33);
+  TestSortWithMagnitude(39);
+#endif
 }
 DECLARE_UNITTEST(TestSortWithLargeNumberOfItems);
 
diff --git a/thrust/testing/inner_product.cu b/thrust/testing/inner_product.cu
index eb98d9e1287..5f21fe4d24e 100644
--- a/thrust/testing/inner_product.cu
+++ b/thrust/testing/inner_product.cu
@@ -131,9 +131,11 @@ void TestInnerProductWithBigIndexesHelper(int magnitude)
 void TestInnerProductWithBigIndexes()
 {
   TestInnerProductWithBigIndexesHelper(30);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestInnerProductWithBigIndexesHelper(31);
   TestInnerProductWithBigIndexesHelper(32);
   TestInnerProductWithBigIndexesHelper(33);
+#endif
 }
 DECLARE_UNITTEST(TestInnerProductWithBigIndexes);
 
diff --git a/thrust/testing/max_element.cu b/thrust/testing/max_element.cu
index a490f3950c6..184f97af48c 100644
--- a/thrust/testing/max_element.cu
+++ b/thrust/testing/max_element.cu
@@ -104,8 +104,10 @@ void TestMaxElementWithBigIndexesHelper(int magnitude)
 void TestMaxElementWithBigIndexes()
 {
   TestMaxElementWithBigIndexesHelper(30);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestMaxElementWithBigIndexesHelper(31);
   TestMaxElementWithBigIndexesHelper(32);
   TestMaxElementWithBigIndexesHelper(33);
+#endif
 }
 DECLARE_UNITTEST(TestMaxElementWithBigIndexes);
diff --git a/thrust/testing/min_element.cu b/thrust/testing/min_element.cu
index 86856548870..cf2685aa3a7 100644
--- a/thrust/testing/min_element.cu
+++ b/thrust/testing/min_element.cu
@@ -102,8 +102,10 @@ void TestMinElementWithBigIndexesHelper(int magnitude)
 void TestMinElementWithBigIndexes()
 {
   TestMinElementWithBigIndexesHelper(30);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestMinElementWithBigIndexesHelper(31);
   TestMinElementWithBigIndexesHelper(32);
   TestMinElementWithBigIndexesHelper(33);
+#endif
 }
 DECLARE_UNITTEST(TestMinElementWithBigIndexes);
diff --git a/thrust/testing/minmax_element.cu b/thrust/testing/minmax_element.cu
index 59d4f023adb..3f6fefaca06 100644
--- a/thrust/testing/minmax_element.cu
+++ b/thrust/testing/minmax_element.cu
@@ -116,8 +116,10 @@ void TestMinMaxElementWithBigIndexesHelper(int magnitude)
 void TestMinMaxElementWithBigIndexes()
 {
   TestMinMaxElementWithBigIndexesHelper(30);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestMinMaxElementWithBigIndexesHelper(31);
   TestMinMaxElementWithBigIndexesHelper(32);
   TestMinMaxElementWithBigIndexesHelper(33);
+#endif
 }
 DECLARE_UNITTEST(TestMinMaxElementWithBigIndexes);
diff --git a/thrust/testing/reduce.cu b/thrust/testing/reduce.cu
index 24dea4288d0..e0865cce64e 100644
--- a/thrust/testing/reduce.cu
+++ b/thrust/testing/reduce.cu
@@ -200,8 +200,10 @@ void TestReduceWithBigIndexesHelper(int magnitude)
 void TestReduceWithBigIndexes()
 {
   TestReduceWithBigIndexesHelper(30);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestReduceWithBigIndexesHelper(31);
   TestReduceWithBigIndexesHelper(32);
   TestReduceWithBigIndexesHelper(33);
+#endif
 }
 DECLARE_UNITTEST(TestReduceWithBigIndexes);
diff --git a/thrust/testing/scan.cu b/thrust/testing/scan.cu
index 171868ee602..6a98cb99e61 100644
--- a/thrust/testing/scan.cu
+++ b/thrust/testing/scan.cu
@@ -617,8 +617,10 @@ void TestInclusiveScanWithBigIndexes()
 {
   TestInclusiveScanWithBigIndexesHelper(30);
   TestInclusiveScanWithBigIndexesHelper(31);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestInclusiveScanWithBigIndexesHelper(32);
   TestInclusiveScanWithBigIndexesHelper(33);
+#endif
 }
 
 DECLARE_UNITTEST(TestInclusiveScanWithBigIndexes);
@@ -646,8 +648,10 @@ void TestExclusiveScanWithBigIndexes()
 {
   TestExclusiveScanWithBigIndexesHelper(30);
   TestExclusiveScanWithBigIndexesHelper(31);
+#ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestExclusiveScanWithBigIndexesHelper(32);
   TestExclusiveScanWithBigIndexesHelper(33);
+#endif
 }
 
 DECLARE_UNITTEST(TestExclusiveScanWithBigIndexes);
diff --git a/thrust/testing/set_difference.cu b/thrust/testing/set_difference.cu
index 0986ff777f8..cdb538d384c 100644
--- a/thrust/testing/set_difference.cu
+++ b/thrust/testing/set_difference.cu
@@ -192,10 +192,13 @@ void TestSetDifferenceWithBigIndexesHelper(int magnitude)
 
 void TestSetDifferenceWithBigIndexes()
 {
+#  ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestSetDifferenceWithBigIndexesHelper(30);
   TestSetDifferenceWithBigIndexesHelper(31);
   TestSetDifferenceWithBigIndexesHelper(32);
   TestSetDifferenceWithBigIndexesHelper(33);
+#  endif
 }
 DECLARE_UNITTEST(TestSetDifferenceWithBigIndexes);
+
 #endif
diff --git a/thrust/testing/set_intersection.cu b/thrust/testing/set_intersection.cu
index c6d9f9e119a..392e23b7337 100644
--- a/thrust/testing/set_intersection.cu
+++ b/thrust/testing/set_intersection.cu
@@ -227,10 +227,12 @@ void TestSetDifferenceWithBigIndexesHelper(int magnitude)
 
 void TestSetDifferenceWithBigIndexes()
 {
+#  ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
   TestSetDifferenceWithBigIndexesHelper(30);
   TestSetDifferenceWithBigIndexesHelper(31);
   TestSetDifferenceWithBigIndexesHelper(32);
   TestSetDifferenceWithBigIndexesHelper(33);
+#  endif
 }
 DECLARE_UNITTEST(TestSetDifferenceWithBigIndexes);
 #endif
diff --git a/thrust/testing/unique_by_key.cu b/thrust/testing/unique_by_key.cu
index 942c0bcdd6f..4f93310633c 100644
--- a/thrust/testing/unique_by_key.cu
+++ b/thrust/testing/unique_by_key.cu
@@ -391,6 +391,8 @@ VariableUnitTest<TestUniqueCopyByKeyToDiscardIterator, IntegralTypes> TestUnique
 // OpenMP has issues with these tests, NVIDIA/cccl#1715
 #if THRUST_DEVICE_SYSTEM != THRUST_DEVICE_SYSTEM_OMP
 
+#  ifndef THRUST_FORCE_32_BIT_OFFSET_TYPE
+
 template <typename K>
 struct TestUniqueCopyByKeyLargeInput
 {
@@ -448,6 +450,8 @@ struct TestUniqueCopyByKeyLargeOutCount
 };
 SimpleUnitTest<TestUniqueCopyByKeyLargeOutCount, IntegralTypes> TestUniqueCopyByKeyLargeOutCountInstance;
 
+#  endif // THRUST_FORCE_32_BIT_OFFSET_TYPE
+
 #endif // non-OpenMP backend
 
 // This test fails only on GCC 6