Update cudf for compatibility with the latest cuco

cpp/cmake/thirdparty/get_cucollections.cmake

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2021-2022, NVIDIA CORPORATION.
+# Copyright (c) 2021-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 # in compliance with the License. You may obtain a copy of the License at
@@ -14,12 +14,14 @@
 
 # This function finds cuCollections and performs any additional configuration.
 function(find_and_configure_cucollections)
-  include(${rapids-cmake-dir}/cpm/cuco.cmake)
-  if(BUILD_SHARED_LIBS)
-    rapids_cpm_cuco(BUILD_EXPORT_SET cudf-exports)
-  else()
-    rapids_cpm_cuco(BUILD_EXPORT_SET cudf-exports INSTALL_EXPORT_SET cudf-exports)
-  endif()
+  rapids_cpm_find(
+    cuco 0.0.1
+    GLOBAL_TARGETS cuco::cuco
+    CPM_ARGS
+    GIT_REPOSITORY https://github.com/NVIDIA/cuCollections.git
+    GIT_TAG fabd885327dc1662e1a63d9b8870ebf888893c09
+    OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF"
+  )
 endfunction()
 
 find_and_configure_cucollections()

cpp/include/cudf/detail/hash_reduce_by_row.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,8 +33,8 @@
 namespace cudf::detail {
 
 using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
-using hash_map_type =
-  cuco::static_map<size_type, size_type, cuda::thread_scope_device, hash_table_allocator_type>;
+using hash_map_type             = cuco::legacy::
+  static_map<size_type, size_type, cuda::thread_scope_device, hash_table_allocator_type>;
 
 /**
  * @brief The base struct for customized reduction functor to perform reduce-by-key with keys are

cpp/include/cudf/detail/join.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -63,7 +63,7 @@ struct hash_join {
                           cudf::size_type,
                           cuda::thread_scope_device,
                           rmm::mr::stream_allocator_adaptor<default_allocator<char>>,
-                          cuco::double_hashing<DEFAULT_JOIN_CG_SIZE, Hasher, Hasher>>;
+                          cuco::legacy::double_hashing<DEFAULT_JOIN_CG_SIZE, Hasher, Hasher>>;
 
   hash_join()                            = delete;
   ~hash_join()                           = default;

cpp/src/io/json/json_tree.cu

@@ -550,21 +550,22 @@ rmm::device_uvector<size_type> hash_node_type_with_field_name(device_span<Symbol
   using hasher_type                             = decltype(d_hasher);
   constexpr size_type empty_node_index_sentinel = -1;
   auto key_set =
-    cuco::experimental::static_set{cuco::experimental::extent{compute_hash_table_size(
-                                     num_fields, 40)},  // 40% occupancy in hash map
-                                   cuco::empty_key{empty_node_index_sentinel},
-                                   d_equal,
-                                   cuco::experimental::linear_probing<1, hasher_type>{d_hasher},
-                                   hash_table_allocator_type{default_allocator<char>{}, stream},
-                                   stream.value()};
+    cuco::static_set{cuco::extent{compute_hash_table_size(num_fields, 40)},  // 40% occupancy
+                     cuco::empty_key{empty_node_index_sentinel},
+                     d_equal,
+                     cuco::linear_probing<1, hasher_type>{d_hasher},
+                     {},
+                     {},
+                     hash_table_allocator_type{default_allocator<char>{}, stream},
+                     stream.value()};
   key_set.insert_if_async(iter,
                           iter + num_nodes,
                           thrust::counting_iterator<size_type>(0),  // stencil
                           is_field_name_node,
                           stream.value());
 
   auto const get_hash_value =
-    [key_set = key_set.ref(cuco::experimental::op::find)] __device__(auto node_id) -> size_type {
+    [key_set = key_set.ref(cuco::op::find)] __device__(auto node_id) -> size_type {
     auto const it = key_set.find(node_id);
     return (it == key_set.end()) ? size_type{0} : *it;
   };
@@ -738,13 +739,14 @@ std::pair<rmm::device_uvector<size_type>, rmm::device_uvector<size_type>> hash_n
   using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
   using hasher_type               = decltype(d_hashed_cache);
 
-  auto key_set = cuco::experimental::static_set{
-    cuco::experimental::extent{compute_hash_table_size(num_nodes)},
-    cuco::empty_key<cudf::size_type>{empty_node_index_sentinel},
-    d_equal,
-    cuco::experimental::linear_probing<1, hasher_type>{d_hashed_cache},
-    hash_table_allocator_type{default_allocator<char>{}, stream},
-    stream.value()};
+  auto key_set = cuco::static_set{cuco::extent{compute_hash_table_size(num_nodes)},
+                                  cuco::empty_key<cudf::size_type>{empty_node_index_sentinel},
+                                  d_equal,
+                                  cuco::linear_probing<1, hasher_type>{d_hashed_cache},
+                                  {},
+                                  {},
+                                  hash_table_allocator_type{default_allocator<char>{}, stream},
+                                  stream.value()};
 
   // insert and convert node ids to unique set ids
   auto nodes_itr         = thrust::make_counting_iterator<size_type>(0);

cpp/src/io/orc/orc_gpu.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -44,7 +44,7 @@ using cudf::detail::host_2dspan;
 auto constexpr KEY_SENTINEL   = size_type{-1};
 auto constexpr VALUE_SENTINEL = size_type{-1};
 
-using map_type = cuco::static_map<size_type, size_type>;
+using map_type = cuco::legacy::static_map<size_type, size_type>;
 
 /**
  * @brief The alias of `map_type::pair_atomic_type` class.

cpp/src/io/parquet/parquet_gpu.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,7 +28,7 @@ namespace cudf::io::parquet::detail {
 auto constexpr KEY_SENTINEL   = size_type{-1};
 auto constexpr VALUE_SENTINEL = size_type{-1};
 
-using map_type = cuco::static_map<size_type, size_type>;
+using map_type = cuco::legacy::static_map<size_type, size_type>;
 
 /**
  * @brief The alias of `map_type::pair_atomic_type` class.

cpp/src/join/join_common_utils.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -50,13 +50,14 @@ using multimap_type = cudf::hash_join::impl_type::map_type;
 // Multimap type used for mixed joins. TODO: This is a temporary alias used
 // until the mixed joins are converted to using CGs properly. Right now it's
 // using a cooperative group of size 1.
-using mixed_multimap_type = cuco::static_multimap<hash_value_type,
-                                                  size_type,
-                                                  cuda::thread_scope_device,
-                                                  hash_table_allocator_type,
-                                                  cuco::double_hashing<1, hash_type, hash_type>>;
-
-using semi_map_type = cuco::
+using mixed_multimap_type =
+  cuco::static_multimap<hash_value_type,
+                        size_type,
+                        cuda::thread_scope_device,
+                        hash_table_allocator_type,
+                        cuco::legacy::double_hashing<1, hash_type, hash_type>>;
+
+using semi_map_type = cuco::legacy::
   static_map<hash_value_type, size_type, cuda::thread_scope_device, hash_table_allocator_type>;
 
 using row_hash_legacy =

cpp/src/search/contains_table.cu

@@ -156,9 +156,9 @@ void dispatch_nan_comparator(
   // Distinguish probing scheme CG sizes between nested and flat types for better performance
   auto const probing_scheme = [&]() {
     if constexpr (HasNested) {
-      return cuco::experimental::linear_probing<4, Hasher>{d_hasher};
+      return cuco::linear_probing<4, Hasher>{d_hasher};
     } else {
-      return cuco::experimental::linear_probing<1, Hasher>{d_hasher};
+      return cuco::linear_probing<1, Hasher>{d_hasher};
     }
   }();
 
@@ -226,11 +226,13 @@ rmm::device_uvector<bool> contains(table_view const& haystack,
     [&](auto const& d_self_equal, auto const& d_two_table_equal, auto const& probing_scheme) {
       auto const d_equal = comparator_adapter{d_self_equal, d_two_table_equal};
 
-      auto set = cuco::experimental::static_set{
-        cuco::experimental::extent{compute_hash_table_size(haystack.num_rows())},
+      auto set = cuco::static_set{
+        cuco::extent{compute_hash_table_size(haystack.num_rows())},
         cuco::empty_key{lhs_index_type{-1}},
         d_equal,
         probing_scheme,
+        {},
+        {},
         detail::hash_table_allocator_type{default_allocator<lhs_index_type>{}, stream},
         stream.value()};
 

cpp/src/stream_compaction/distinct_count.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -141,13 +141,15 @@ cudf::size_type distinct_count(table_view const& keys,
 
   auto const comparator_helper = [&](auto const row_equal) {
     using hasher_type = decltype(hash_key);
-    auto key_set      = cuco::experimental::static_set{
-      cuco::experimental::extent{compute_hash_table_size(num_rows)},
-      cuco::empty_key<cudf::size_type>{-1},
-      row_equal,
-      cuco::experimental::linear_probing<1, hasher_type>{hash_key},
-      detail::hash_table_allocator_type{default_allocator<char>{}, stream},
-      stream.value()};
+    auto key_set =
+      cuco::static_set{cuco::extent{compute_hash_table_size(num_rows)},
+                       cuco::empty_key<cudf::size_type>{-1},
+                       row_equal,
+                       cuco::linear_probing<1, hasher_type>{hash_key},
+                       {},
+                       {},
+                       detail::hash_table_allocator_type{default_allocator<char>{}, stream},
+                       stream.value()};
 
     auto const iter = thrust::counting_iterator<cudf::size_type>(0);
     // when nulls are equal, we skip hashing any row that has a null

cpp/src/stream_compaction/stream_compaction_common.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -31,8 +31,8 @@ namespace detail {
 
 using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
 
-using hash_map_type =
-  cuco::static_map<size_type, size_type, cuda::thread_scope_device, hash_table_allocator_type>;
+using hash_map_type = cuco::legacy::
+  static_map<size_type, size_type, cuda::thread_scope_device, hash_table_allocator_type>;
 
 }  // namespace detail
 }  // namespace cudf

cpp/src/text/bpe/byte_pair_encoding.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -47,6 +47,7 @@ using hash_value_type    = string_hasher_type::result_type;
 using merge_pair_type    = thrust::pair<cudf::string_view, cudf::string_view>;
 
 using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
+using cuco_storage              = cuco::storage<1>;
 
 /**
  * @brief Hasher function used for building and using the cuco static-map
@@ -101,15 +102,16 @@ struct bpe_equal {
   }
 };
 
-using bpe_probe_scheme = cuco::experimental::linear_probing<1, bpe_hasher>;
+using bpe_probe_scheme = cuco::linear_probing<1, bpe_hasher>;
 
-using merge_pairs_map_type = cuco::experimental::static_map<cudf::size_type,
-                                                            cudf::size_type,
-                                                            cuco::experimental::extent<std::size_t>,
-                                                            cuda::thread_scope_device,
-                                                            bpe_equal,
-                                                            bpe_probe_scheme,
-                                                            hash_table_allocator_type>;
+using merge_pairs_map_type = cuco::static_map<cudf::size_type,
+                                              cudf::size_type,
+                                              cuco::extent<std::size_t>,
+                                              cuda::thread_scope_device,
+                                              bpe_equal,
+                                              bpe_probe_scheme,
+                                              hash_table_allocator_type,
+                                              cuco_storage>;
 
 /**
  * @brief Hasher function used for building and using the cuco static-map
@@ -158,15 +160,16 @@ struct mp_equal {
   }
 };
 
-using mp_probe_scheme = cuco::experimental::linear_probing<1, mp_hasher>;
+using mp_probe_scheme = cuco::linear_probing<1, mp_hasher>;
 
-using mp_table_map_type = cuco::experimental::static_map<cudf::size_type,
-                                                         cudf::size_type,
-                                                         cuco::experimental::extent<std::size_t>,
-                                                         cuda::thread_scope_device,
-                                                         mp_equal,
-                                                         mp_probe_scheme,
-                                                         hash_table_allocator_type>;
+using mp_table_map_type = cuco::static_map<cudf::size_type,
+                                           cudf::size_type,
+                                           cuco::extent<std::size_t>,
+                                           cuda::thread_scope_device,
+                                           mp_equal,
+                                           mp_probe_scheme,
+                                           hash_table_allocator_type,
+                                           cuco_storage>;
 
 }  // namespace detail
 
@@ -188,8 +191,8 @@ struct bpe_merge_pairs::bpe_merge_pairs_impl {
                        std::unique_ptr<detail::mp_table_map_type>&& mp_table_map);
 
   auto const get_merge_pairs() const { return *d_merge_pairs; }
-  auto get_merge_pairs_ref() const { return merge_pairs_map->ref(cuco::experimental::op::find); }
-  auto get_mp_table_ref() const { return mp_table_map->ref(cuco::experimental::op::find); }
+  auto get_merge_pairs_ref() const { return merge_pairs_map->ref(cuco::op::find); }
+  auto get_mp_table_ref() const { return mp_table_map->ref(cuco::op::find); }
 };
 
 }  // namespace nvtext

cpp/src/text/bpe/load_merge_pairs.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -48,6 +48,8 @@ std::unique_ptr<detail::merge_pairs_map_type> initialize_merge_pairs_map(
     cuco::empty_value{-1},
     bpe_equal{input},
     bpe_probe_scheme{bpe_hasher{input}},
+    cuco::thread_scope_device,
+    cuco_storage{},
     hash_table_allocator_type{default_allocator<char>{}, stream},
     stream.value());
 
@@ -70,6 +72,8 @@ std::unique_ptr<detail::mp_table_map_type> initialize_mp_table_map(
     cuco::empty_value{-1},
     mp_equal{input},
     mp_probe_scheme{mp_hasher{input}},
+    cuco::thread_scope_device,
+    cuco_storage{},
     hash_table_allocator_type{default_allocator<char>{}, stream},
     stream.value());
 

cpp/src/text/vocabulary_tokenize.cu

@@ -94,14 +94,16 @@ struct vocab_equal {
 };
 
 using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
-using probe_scheme              = cuco::experimental::linear_probing<1, vocab_hasher>;
-using vocabulary_map_type       = cuco::experimental::static_map<cudf::size_type,
-                                                           cudf::size_type,
-                                                           cuco::experimental::extent<std::size_t>,
-                                                           cuda::thread_scope_device,
-                                                           vocab_equal,
-                                                           probe_scheme,
-                                                           hash_table_allocator_type>;
+using probe_scheme              = cuco::linear_probing<1, vocab_hasher>;
+using cuco_storage              = cuco::storage<1>;
+using vocabulary_map_type       = cuco::static_map<cudf::size_type,
+                                             cudf::size_type,
+                                             cuco::extent<std::size_t>,
+                                             cuda::thread_scope_device,
+                                             vocab_equal,
+                                             probe_scheme,
+                                             hash_table_allocator_type,
+                                             cuco_storage>;
 }  // namespace
 }  // namespace detail
 
@@ -116,7 +118,7 @@ struct tokenize_vocabulary::tokenize_vocabulary_impl {
   col_device_view const d_vocabulary;
   std::unique_ptr<detail::vocabulary_map_type> vocabulary_map;
 
-  auto get_map_ref() const { return vocabulary_map->ref(cuco::experimental::op::find); }
+  auto get_map_ref() const { return vocabulary_map->ref(cuco::op::find); }
 
   tokenize_vocabulary_impl(std::unique_ptr<cudf::column>&& vocab,
                            col_device_view&& d_vocab,
@@ -150,6 +152,8 @@ tokenize_vocabulary::tokenize_vocabulary(cudf::strings_column_view const& input,
     cuco::empty_value{-1},
     detail::vocab_equal{*d_vocabulary},
     detail::probe_scheme{detail::vocab_hasher{*d_vocabulary}},
+    cuco::thread_scope_device,
+    detail::cuco_storage{},
     detail::hash_table_allocator_type{default_allocator<char>{}, stream},
     stream.value());