Test updates of CCCL (thrust, cub, libcudacxx) to 2.1.0.

conda/recipes/libcudf/meta.yaml

@@ -46,7 +46,7 @@ requirements:
     - ninja
     - sysroot_{{ target_platform }} {{ sysroot_version }}
   host:
-    - librmm ={{ minor_version }}
+    # - librmm ={{ minor_version }}
     - libkvikio ={{ minor_version }}
     {% if cuda_major == "11" %}
     - cudatoolkit
@@ -100,7 +100,7 @@ outputs:
         {% endif %}
         - cuda-version {{ cuda_spec }}
         - nvcomp {{ nvcomp_version }}
-        - librmm ={{ minor_version }}
+        # - librmm ={{ minor_version }}
         - libkvikio ={{ minor_version }}
         - libarrow {{ libarrow_version }}
         - dlpack {{ dlpack_version }}

cpp/benchmarks/common/generate_input.cu

@@ -53,6 +53,8 @@
 #include <thrust/transform.h>
 #include <thrust/tuple.h>
 
+#include <cuda/functional>
+
 #include <algorithm>
 #include <cstdint>
 #include <memory>
@@ -247,12 +249,12 @@ struct random_value_fn<T, std::enable_if_t<cudf::is_chrono<T>()>> {
       sec.end(),
       ns.begin(),
       result.begin(),
-      [] __device__(int64_t sec_value, int64_t nanoseconds_value) {
+      cuda::proclaim_return_type<T>([] __device__(int64_t sec_value, int64_t nanoseconds_value) {
         auto const timestamp_ns =
           cudf::duration_s{sec_value} + cudf::duration_ns{nanoseconds_value};
         // Return value in the type's precision
         return T(cuda::std::chrono::duration_cast<typename T::duration>(timestamp_ns));
-      });
+      }));
     return result;
   }
 };
@@ -367,12 +369,13 @@ rmm::device_uvector<cudf::size_type> sample_indices_with_run_length(cudf::size_t
     // This is gather.
     auto avg_repeated_sample_indices_iterator = thrust::make_transform_iterator(
       thrust::make_counting_iterator(0),
-      [rb              = run_lens.begin(),
-       re              = run_lens.end(),
-       samples_indices = samples_indices.begin()] __device__(cudf::size_type i) {
-        auto sample_idx = thrust::upper_bound(thrust::seq, rb, re, i) - rb;
-        return samples_indices[sample_idx];
-      });
+      cuda::proclaim_return_type<cudf::size_type>(
+        [rb              = run_lens.begin(),
+         re              = run_lens.end(),
+         samples_indices = samples_indices.begin()] __device__(cudf::size_type i) {
+          auto sample_idx = thrust::upper_bound(thrust::seq, rb, re, i) - rb;
+          return samples_indices[sample_idx];
+        }));
     rmm::device_uvector<cudf::size_type> repeated_sample_indices(num_rows,
                                                                  cudf::get_default_stream());
     thrust::copy(thrust::device,
@@ -513,7 +516,7 @@ std::unique_ptr<cudf::column> create_random_utf8_string_column(data_profile cons
     lengths.end(),
     null_mask.begin(),
     lengths.begin(),
-    [] __device__(auto) { return 0; },
+    cuda::proclaim_return_type<cudf::size_type>([] __device__(auto) { return 0; }),
     thrust::logical_not<bool>{});
   auto valid_lengths = thrust::make_transform_iterator(
     thrust::make_zip_iterator(thrust::make_tuple(lengths.begin(), null_mask.begin())),

cpp/cmake/thirdparty/patches/thrust_disable_64bit_dispatching.diff

@@ -1,29 +1,25 @@
 diff --git a/thrust/system/cuda/detail/dispatch.h b/thrust/system/cuda/detail/dispatch.h
-index d0e3f94..76774b0 100644
+index d0e3f94e..5c32a9c6 100644
 --- a/thrust/system/cuda/detail/dispatch.h
 +++ b/thrust/system/cuda/detail/dispatch.h
-@@ -32,9 +32,8 @@
+@@ -32,8 +32,7 @@
          status = call arguments; \
      } \
      else { \
 -        auto THRUST_PP_CAT2(count, _fixed) = static_cast<thrust::detail::int64_t>(count); \
 -        status = call arguments; \
--    }
-+       throw std::runtime_error("THRUST_INDEX_TYPE_DISPATCH 64-bit count is unsupported in libcudf"); \
-+    }
++        throw std::runtime_error("THRUST_INDEX_TYPE_DISPATCH 64-bit count is unsupported in libcudf"); \
+     }
 
  /**
-  * Dispatch between 32-bit and 64-bit index based versions of the same algorithm
-@@ -52,10 +51,8 @@
+@@ -52,9 +51,7 @@
          status = call arguments; \
      } \
      else { \
 -        auto THRUST_PP_CAT2(count1, _fixed) = static_cast<thrust::detail::int64_t>(count1); \
 -        auto THRUST_PP_CAT2(count2, _fixed) = static_cast<thrust::detail::int64_t>(count2); \
 -        status = call arguments; \
--    }
-+       throw std::runtime_error("THRUST_DOUBLE_INDEX_TYPE_DISPATCH 64-bit count is unsupported in libcudf"); \
-+    }
++        throw std::runtime_error("THRUST_DOUBLE_INDEX_TYPE_DISPATCH 64-bit count is unsupported in libcudf"); \
+     }
  /**
   * Dispatch between 32-bit and 64-bit index based versions of the same algorithm
-  * implementation. This version allows using different token sequences for callables

cpp/cmake/thirdparty/patches/thrust_faster_scan_compile_times.diff

@@ -1,8 +1,8 @@
 diff --git a/dependencies/cub/cub/device/dispatch/dispatch_radix_sort.cuh b/dependencies/cub/cub/device/dispatch/dispatch_radix_sort.cuh
-index b188c75f..3f36656f 100644
+index 953b24b0..3a895323 100644
 --- a/dependencies/cub/cub/device/dispatch/dispatch_radix_sort.cuh
 +++ b/dependencies/cub/cub/device/dispatch/dispatch_radix_sort.cuh
-@@ -736,7 +736,7 @@ struct DeviceRadixSortPolicy
+@@ -738,7 +738,7 @@ struct DeviceRadixSortPolicy
 
 
      /// SM60 (GP100)
@@ -12,28 +12,28 @@ index b188c75f..3f36656f 100644
          enum {
              PRIMARY_RADIX_BITS      = (sizeof(KeyT) > 1) ? 7 : 5,    // 6.9B 32b keys/s (Quadro P100)
 diff --git a/dependencies/cub/cub/device/dispatch/dispatch_reduce.cuh b/dependencies/cub/cub/device/dispatch/dispatch_reduce.cuh
-index e0470ccb..6a0c2ed6 100644
+index de182a93..e2ba1f4b 100644
 --- a/dependencies/cub/cub/device/dispatch/dispatch_reduce.cuh
 +++ b/dependencies/cub/cub/device/dispatch/dispatch_reduce.cuh
-@@ -280,7 +280,7 @@ struct DeviceReducePolicy
-     };
+@@ -423,7 +423,7 @@ struct DeviceReducePolicy
+   };
 
-     /// SM60
--    struct Policy600 : ChainedPolicy<600, Policy600, Policy350>
-+    struct Policy600 : ChainedPolicy<600, Policy600, Policy600>
-     {
-         // ReducePolicy (P100: 591 GB/s @ 64M 4B items; 583 GB/s @ 256M 1B items)
-         typedef AgentReducePolicy<
+   /// SM60
+-  struct Policy600 : ChainedPolicy<600, Policy600, Policy350>
++  struct Policy600 : ChainedPolicy<600, Policy600, Policy600>
+   {
+     static constexpr int threads_per_block  = 256;
+     static constexpr int items_per_thread   = 16;
 diff --git a/dependencies/cub/cub/device/dispatch/dispatch_scan.cuh b/dependencies/cub/cub/device/dispatch/dispatch_scan.cuh
-index c2d04588..ac2d10e0 100644
+index 0df89b7c..6c499d87 100644
 --- a/dependencies/cub/cub/device/dispatch/dispatch_scan.cuh
 +++ b/dependencies/cub/cub/device/dispatch/dispatch_scan.cuh
-@@ -177,7 +177,7 @@ struct DeviceScanPolicy
-     };
+@@ -240,7 +240,7 @@ struct DeviceScanPolicy
+   };
 
-     /// SM600
--    struct Policy600 : ChainedPolicy<600, Policy600, Policy520>
-+    struct Policy600 : ChainedPolicy<600, Policy600, Policy600>
-     {
-         typedef AgentScanPolicy<
-                 128, 15,                                        ///< Threads per block, items per thread
+   /// SM600
+-  struct Policy600 : ChainedPolicy<600, Policy600, Policy520>
++  struct Policy600 : ChainedPolicy<600, Policy600, Policy600>
+   {
+     typedef AgentScanPolicy<128,
+                             15, ///< Threads per block, items per thread

cpp/cmake/thirdparty/patches/thrust_override.json

@@ -4,14 +4,14 @@
     "Thrust" : {
       "patches" : [
         {
-          "file" : "Thrust/install_rules.diff",
-          "issue" : "Thrust 1.X installs incorrect files [https://github.com/NVIDIA/thrust/issues/1790]",
-          "fixed_in" : "2.0.0"
+          "file" : "Thrust/reroot_support.diff",
+          "issue" : "Support conda-forge usage of CMake rerooting [https://github.com/NVIDIA/thrust/pull/1969]",
+          "fixed_in" : "2.2"
         },
         {
-          "file" : "${current_json_dir}/thrust_transform_iter_with_reduce_by_key.diff",
-          "issue" : "Support transform_output_iterator as output of reduce by key [https://github.com/NVIDIA/thrust/pull/1805]",
-          "fixed_in" : "2.1"
+          "file" : "Thrust/thrust_libcudacxx_2.1.0.diff",
+          "issue" : "Use libcudacxx 2.1.0 in Thrust config.",
+          "fixed_in" : "2.2"
         },
         {
           "file" : "${current_json_dir}/thrust_disable_64bit_dispatching.diff",
@@ -27,11 +27,6 @@
           "file" : "${current_json_dir}/thrust_faster_scan_compile_times.diff",
           "issue" : "Improve Thrust scan compile times by reducing the number of kernels generated [https://github.com/rapidsai/cudf/pull/8183]",
           "fixed_in" : ""
-        },
-        {
-          "file" : "${current_json_dir}/cub_segmented_sort_with_bool_key.diff",
-          "issue" : "Fix an error in CUB DeviceSegmentedSort when the keys are bool type [https://github.com/NVIDIA/cub/issues/594]",
-          "fixed_in" : "2.1"
         }
       ]
     }

cpp/include/cudf/detail/null_mask.cuh

@@ -37,6 +37,8 @@
 #include <thrust/transform.h>
 #include <thrust/tuple.h>
 
+#include <cuda/functional>
+
 #include <algorithm>
 #include <iterator>
 #include <optional>
@@ -330,20 +332,21 @@ rmm::device_uvector<size_type> segmented_count_bits(bitmask_type const* bitmask,
     // set bits from the length of the segment.
     auto segments_begin =
       thrust::make_zip_iterator(first_bit_indices_begin, last_bit_indices_begin);
-    auto segment_length_iterator =
-      thrust::transform_iterator(segments_begin, [] __device__(auto const& segment) {
+    auto segment_length_iterator = thrust::transform_iterator(
+      segments_begin, cuda::proclaim_return_type<size_type>([] __device__(auto const& segment) {
         auto const begin = thrust::get<0>(segment);
         auto const end   = thrust::get<1>(segment);
         return end - begin;
-      });
+      }));
     thrust::transform(rmm::exec_policy(stream),
                       segment_length_iterator,
                       segment_length_iterator + num_ranges,
                       d_bit_counts.data(),
                       d_bit_counts.data(),
-                      [] __device__(auto segment_size, auto segment_bit_count) {
-                        return segment_size - segment_bit_count;
-                      });
+                      cuda::proclaim_return_type<size_type>(
+                        [] __device__(auto segment_size, auto segment_bit_count) {
+                          return segment_size - segment_bit_count;
+                        }));
   }
 
   CUDF_CHECK_CUDA(stream.value());
@@ -541,12 +544,12 @@ std::pair<rmm::device_buffer, size_type> segmented_null_mask_reduction(
 {
   auto const segments_begin =
     thrust::make_zip_iterator(first_bit_indices_begin, last_bit_indices_begin);
-  auto const segment_length_iterator =
-    thrust::make_transform_iterator(segments_begin, [] __device__(auto const& segment) {
+  auto const segment_length_iterator = thrust::make_transform_iterator(
+    segments_begin, cuda::proclaim_return_type<size_type>([] __device__(auto const& segment) {
       auto const begin = thrust::get<0>(segment);
       auto const end   = thrust::get<1>(segment);
       return end - begin;
-    });
+    }));
 
   auto const num_segments =
     static_cast<size_type>(std::distance(first_bit_indices_begin, first_bit_indices_end));
@@ -555,9 +558,9 @@ std::pair<rmm::device_buffer, size_type> segmented_null_mask_reduction(
     return cudf::detail::valid_if(
       segment_length_iterator,
       segment_length_iterator + num_segments,
-      [valid_initial_value] __device__(auto const& length) {
+      cuda::proclaim_return_type<bool>([valid_initial_value] __device__(auto const& length) {
         return valid_initial_value.value_or(length > 0);
-      },
+      }),
       stream,
       mr);
   }
@@ -575,13 +578,14 @@ std::pair<rmm::device_buffer, size_type> segmented_null_mask_reduction(
   return cudf::detail::valid_if(
     length_and_valid_count,
     length_and_valid_count + num_segments,
-    [null_handling, valid_initial_value] __device__(auto const& length_and_valid_count) {
-      auto const length      = thrust::get<0>(length_and_valid_count);
-      auto const valid_count = thrust::get<1>(length_and_valid_count);
-      return (null_handling == null_policy::EXCLUDE)
-               ? (valid_initial_value.value_or(false) || valid_count > 0)
-               : (valid_initial_value.value_or(length > 0) && valid_count == length);
-    },
+    cuda::proclaim_return_type<bool>(
+      [null_handling, valid_initial_value] __device__(auto const& length_and_valid_count) {
+        auto const length      = thrust::get<0>(length_and_valid_count);
+        auto const valid_count = thrust::get<1>(length_and_valid_count);
+        return (null_handling == null_policy::EXCLUDE)
+                 ? (valid_initial_value.value_or(false) || valid_count > 0)
+                 : (valid_initial_value.value_or(length > 0) && valid_count == length);
+      }),
     stream,
     mr);
 }

cpp/include/cudf/detail/sizes_to_offsets_iterator.cuh

@@ -27,6 +27,8 @@
 #include <thrust/distance.h>
 #include <thrust/scan.h>
 
+#include <cuda/functional>
+
 #include <stdexcept>
 
 namespace cudf {
@@ -311,9 +313,10 @@ std::pair<std::unique_ptr<column>, size_type> make_offsets_child_column(
   // using exclusive-scan technically requires count+1 input values even though
   // the final input value is never used.
   // The input iterator is wrapped here to allow the last value to be safely read.
-  auto map_fn = [begin, count] __device__(size_type idx) -> size_type {
-    return idx < count ? static_cast<size_type>(begin[idx]) : size_type{0};
-  };
+  auto map_fn =
+    cuda::proclaim_return_type<size_type>([begin, count] __device__(size_type idx) -> size_type {
+      return idx < count ? static_cast<size_type>(begin[idx]) : size_type{0};
+    });
   auto input_itr = cudf::detail::make_counting_transform_iterator(0, map_fn);
   // Use the sizes-to-offsets iterator to compute the total number of elements
   auto const total_elements = sizes_to_offsets(input_itr, input_itr + count + 1, d_offsets, stream);

cpp/include/cudf/detail/utilities/device_operators.cuh

@@ -55,6 +55,34 @@ CUDF_HOST_DEVICE inline auto max(LHS const& lhs, RHS const& rhs)
 {
   return std::max(lhs, rhs);
 }
+
+/**
+ * @brief Functor that casts another functor's result to a specified type.
+ *
+ * CUB 2.0.0 reductions require that the binary operator returns the same type
+ * as the initial value type, so we wrap binary operators with this when used
+ * by CUB.
+ */
+template <typename ResultType, typename F>
+struct cast_functor_fn {
+  F f;
+
+  template <typename... Ts>
+  CUDF_HOST_DEVICE inline ResultType operator()(Ts&&... args)
+  {
+    return static_cast<ResultType>(f(std::forward<Ts>(args)...));
+  }
+};
+
+/**
+ * @brief Function creating a casting functor.
+ */
+template <typename ResultType, typename F>
+cast_functor_fn<ResultType, F> cast_functor(F&& f)
+{
+  return {std::forward<F>(f)};
+}
+
 }  // namespace detail
 
 /**