From f9ce36f83a8803396d84eceab0e9b3a1706f8510 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Thu, 25 Jan 2024 14:11:48 -0800
Subject: [PATCH] Use new cuco_allocator

---
 .../cudf/detail/hash_reduce_by_row.cuh        |  8 +--
 cpp/include/cudf/detail/join.hpp              | 11 ++--
 .../cudf/hashing/detail/hash_allocator.cuh    | 62 -------------------
 cpp/src/hash/concurrent_unordered_map.cuh     |  6 +-
 cpp/src/io/json/json_tree.cu                  | 11 ++--
 cpp/src/join/hash_join.cu                     |  5 +-
 cpp/src/join/join_common_utils.hpp            | 13 ++--
 cpp/src/join/mixed_join.cu                    | 25 ++++----
 cpp/src/join/mixed_join_semi.cu               |  7 ++-
 cpp/src/reductions/histogram.cu               | 13 ++--
 cpp/src/search/contains_table.cu              |  4 +-
 cpp/src/stream_compaction/distinct.cu         |  4 +-
 cpp/src/stream_compaction/distinct_count.cu   | 16 ++---
 .../stream_compaction_common.hpp              | 10 +--
 cpp/src/text/bpe/byte_pair_encoding.cuh       | 11 ++--
 cpp/src/text/bpe/load_merge_pairs.cu          | 33 +++++-----
 cpp/src/text/vocabulary_tokenize.cu           | 12 ++--
 17 files changed, 85 insertions(+), 166 deletions(-)
 delete mode 100644 cpp/include/cudf/hashing/detail/hash_allocator.cuh

diff --git a/cpp/include/cudf/detail/hash_reduce_by_row.cuh b/cpp/include/cudf/detail/hash_reduce_by_row.cuh
index f63d1922950..006cb5142c9 100644
--- a/cpp/include/cudf/detail/hash_reduce_by_row.cuh
+++ b/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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include <cudf/hashing/detail/hash_allocator.cuh>
+#include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/table/experimental/row_operators.cuh>
 #include <cudf/types.hpp>
@@ -22,7 +22,6 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
-#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <thrust/for_each.h>
 #include <thrust/iterator/counting_iterator.h>
@@ -32,9 +31,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>;
+  cuco::static_map<size_type, size_type, cuda::thread_scope_device, cudf::detail::cuco_allocator>;
 
 /**
  * @brief The base struct for customized reduction functor to perform reduce-by-key with keys are
diff --git a/cpp/include/cudf/detail/join.hpp b/cpp/include/cudf/detail/join.hpp
index b69632c83ca..ad6269dae30 100644
--- a/cpp/include/cudf/detail/join.hpp
+++ b/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.
@@ -24,7 +24,6 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_uvector.hpp>
-#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <cuco/static_multimap.cuh>
 
@@ -33,9 +32,6 @@
 #include <optional>
 
 // Forward declaration
-template <typename T>
-class default_allocator;
-
 namespace cudf::experimental::row::equality {
 class preprocessed_table;
 }
@@ -43,6 +39,9 @@ class preprocessed_table;
 namespace cudf {
 namespace detail {
 
+// Forward declaration
+class cuco_allocator;
+
 constexpr int DEFAULT_JOIN_CG_SIZE = 2;
 
 enum class join_kind { INNER_JOIN, LEFT_JOIN, FULL_JOIN, LEFT_SEMI_JOIN, LEFT_ANTI_JOIN };
@@ -62,7 +61,7 @@ struct hash_join {
     cuco::static_multimap<hash_value_type,
                           cudf::size_type,
                           cuda::thread_scope_device,
-                          rmm::mr::stream_allocator_adaptor<default_allocator<char>>,
+                          cudf::detail::cuco_allocator,
                           cuco::double_hashing<DEFAULT_JOIN_CG_SIZE, Hasher, Hasher>>;
 
   hash_join()                            = delete;
diff --git a/cpp/include/cudf/hashing/detail/hash_allocator.cuh b/cpp/include/cudf/hashing/detail/hash_allocator.cuh
deleted file mode 100644
index 64a2a852ae4..00000000000
--- a/cpp/include/cudf/hashing/detail/hash_allocator.cuh
+++ /dev/null
@@ -1,62 +0,0 @@
-/*
- * Copyright (c) 2017-2023, 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
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#pragma once
-
-#include <new>
-
-#include <cudf/utilities/default_stream.hpp>
-
-#include <rmm/cuda_stream_view.hpp>
-#include <rmm/mr/device/device_memory_resource.hpp>
-#include <rmm/mr/device/managed_memory_resource.hpp>
-#include <rmm/mr/device/per_device_resource.hpp>
-
-template <class T>
-struct default_allocator {
-  using value_type                    = T;
-  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource();
-
-  default_allocator() = default;
-
-  template <class U>
-  constexpr default_allocator(default_allocator<U> const&) noexcept
-  {
-  }
-
-  T* allocate(std::size_t n, rmm::cuda_stream_view stream = cudf::get_default_stream()) const
-  {
-    return static_cast<T*>(mr->allocate(n * sizeof(T), stream));
-  }
-
-  void deallocate(T* p,
-                  std::size_t n,
-                  rmm::cuda_stream_view stream = cudf::get_default_stream()) const
-  {
-    mr->deallocate(p, n * sizeof(T), stream);
-  }
-};
-
-template <class T, class U>
-bool operator==(default_allocator<T> const&, default_allocator<U> const&)
-{
-  return true;
-}
-template <class T, class U>
-bool operator!=(default_allocator<T> const&, default_allocator<U> const&)
-{
-  return false;
-}
diff --git a/cpp/src/hash/concurrent_unordered_map.cuh b/cpp/src/hash/concurrent_unordered_map.cuh
index d773c2763df..adc87c2400e 100644
--- a/cpp/src/hash/concurrent_unordered_map.cuh
+++ b/cpp/src/hash/concurrent_unordered_map.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2024, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,12 +20,12 @@
 
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/hashing/detail/default_hash.cuh>
-#include <cudf/hashing/detail/hash_allocator.cuh>
 #include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/error.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <thrust/pair.h>
 
@@ -117,7 +117,7 @@ template <typename Key,
           typename Element,
           typename Hasher    = cudf::hashing::detail::default_hash<Key>,
           typename Equality  = equal_to<Key>,
-          typename Allocator = default_allocator<thrust::pair<Key, Element>>>
+          typename Allocator = rmm::mr::polymorphic_allocator<thrust::pair<Key, Element>>>
 class concurrent_unordered_map {
  public:
   using size_type      = size_t;
diff --git a/cpp/src/io/json/json_tree.cu b/cpp/src/io/json/json_tree.cu
index 275907c19c9..db9daf28c06 100644
--- a/cpp/src/io/json/json_tree.cu
+++ b/cpp/src/io/json/json_tree.cu
@@ -17,12 +17,12 @@
 #include "nested_json.hpp"
 #include <io/utilities/hostdevice_vector.hpp>
 
+#include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/scatter.cuh>
 #include <cudf/detail/utilities/algorithm.cuh>
 #include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/hashing/detail/default_hash.cuh>
-#include <cudf/hashing/detail/hash_allocator.cuh>
 #include <cudf/hashing/detail/hashing.hpp>
 #include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/utilities/error.hpp>
@@ -31,7 +31,6 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
-#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <cub/device/device_radix_sort.cuh>
 
@@ -511,7 +510,6 @@ rmm::device_uvector<size_type> hash_node_type_with_field_name(device_span<Symbol
                                                               rmm::cuda_stream_view stream)
 {
   CUDF_FUNC_RANGE();
-  using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
 
   auto const num_nodes  = d_tree.node_categories.size();
   auto const num_fields = thrust::count(rmm::exec_policy(stream),
@@ -555,7 +553,7 @@ rmm::device_uvector<size_type> hash_node_type_with_field_name(device_span<Symbol
                                    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},
+                                   cudf::detail::cuco_allocator{stream},
                                    stream.value()};
   key_set.insert_if_async(iter,
                           iter + num_nodes,
@@ -735,15 +733,14 @@ std::pair<rmm::device_uvector<size_type>, rmm::device_uvector<size_type>> hash_n
   };
 
   constexpr size_type empty_node_index_sentinel = -1;
-  using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
-  using hasher_type               = decltype(d_hashed_cache);
+  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},
+    cudf::detail::cuco_allocator{stream},
     stream.value()};
 
   // insert and convert node ids to unique set ids
diff --git a/cpp/src/join/hash_join.cu b/cpp/src/join/hash_join.cu
index 453257ab228..17616818a58 100644
--- a/cpp/src/join/hash_join.cu
+++ b/cpp/src/join/hash_join.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2020-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.
@@ -19,6 +19,7 @@
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/join.hpp>
 #include <cudf/detail/structs/utilities.hpp>
+#include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/join.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -368,7 +369,7 @@ hash_join<Hasher>::hash_join(cudf::table_view const& build,
                 cuco::empty_key{std::numeric_limits<hash_value_type>::max()},
                 cuco::empty_value{cudf::detail::JoinNoneValue},
                 stream.value(),
-                detail::hash_table_allocator_type{default_allocator<char>{}, stream}},
+                cudf::detail::cuco_allocator{stream}},
     _build{build},
     _preprocessed_build{
       cudf::experimental::row::equality::preprocessed_table::create(_build, stream)}
diff --git a/cpp/src/join/join_common_utils.hpp b/cpp/src/join/join_common_utils.hpp
index e96505e5ed6..b88a4fdef58 100644
--- a/cpp/src/join/join_common_utils.hpp
+++ b/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.
@@ -15,16 +15,13 @@
  */
 #pragma once
 
+#include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/detail/join.hpp>
 #include <cudf/hashing/detail/default_hash.cuh>
-#include <cudf/hashing/detail/hash_allocator.cuh>
-#include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/join.hpp>
 #include <cudf/table/row_operators.cuh>
 #include <cudf/table/table_view.hpp>
 
-#include <rmm/mr/device/polymorphic_allocator.hpp>
-
 #include <cuco/static_map.cuh>
 #include <cuco/static_multimap.cuh>
 
@@ -43,8 +40,6 @@ using pair_type = cuco::pair<hash_value_type, size_type>;
 
 using hash_type = cuco::murmurhash3_32<hash_value_type>;
 
-using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
-
 using multimap_type = cudf::hash_join::impl_type::map_type;
 
 // Multimap type used for mixed joins. TODO: This is a temporary alias used
@@ -53,11 +48,11 @@ using multimap_type = cudf::hash_join::impl_type::map_type;
 using mixed_multimap_type = cuco::static_multimap<hash_value_type,
                                                   size_type,
                                                   cuda::thread_scope_device,
-                                                  hash_table_allocator_type,
+                                                  cudf::detail::cuco_allocator,
                                                   cuco::double_hashing<1, hash_type, hash_type>>;
 
 using semi_map_type = cuco::
-  static_map<hash_value_type, size_type, cuda::thread_scope_device, hash_table_allocator_type>;
+  static_map<hash_value_type, size_type, cuda::thread_scope_device, cudf::detail::cuco_allocator>;
 
 using row_hash_legacy =
   cudf::row_hasher<cudf::hashing::detail::default_hash, cudf::nullate::DYNAMIC>;
diff --git a/cpp/src/join/mixed_join.cu b/cpp/src/join/mixed_join.cu
index 3d902bf93b2..6223114fcd0 100644
--- a/cpp/src/join/mixed_join.cu
+++ b/cpp/src/join/mixed_join.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.
@@ -21,6 +21,7 @@
 #include <cudf/ast/detail/expression_parser.hpp>
 #include <cudf/ast/expressions.hpp>
 #include <cudf/detail/utilities/cuda.cuh>
+#include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/join.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_device_view.cuh>
@@ -127,12 +128,11 @@ mixed_join(
   auto build_view = table_device_view::create(build, stream);
 
   // Don't use multimap_type because we want a CG size of 1.
-  mixed_multimap_type hash_table{
-    compute_hash_table_size(build.num_rows()),
-    cuco::empty_key{std::numeric_limits<hash_value_type>::max()},
-    cuco::empty_value{cudf::detail::JoinNoneValue},
-    stream.value(),
-    detail::hash_table_allocator_type{default_allocator<char>{}, stream}};
+  mixed_multimap_type hash_table{compute_hash_table_size(build.num_rows()),
+                                 cuco::empty_key{std::numeric_limits<hash_value_type>::max()},
+                                 cuco::empty_value{cudf::detail::JoinNoneValue},
+                                 stream.value(),
+                                 cudf::detail::cuco_allocator{stream}};
 
   // TODO: To add support for nested columns we will need to flatten in many
   // places. However, this probably isn't worth adding any time soon since we
@@ -393,12 +393,11 @@ compute_mixed_join_output_size(table_view const& left_equality,
   auto build_view = table_device_view::create(build, stream);
 
   // Don't use multimap_type because we want a CG size of 1.
-  mixed_multimap_type hash_table{
-    compute_hash_table_size(build.num_rows()),
-    cuco::empty_key{std::numeric_limits<hash_value_type>::max()},
-    cuco::empty_value{cudf::detail::JoinNoneValue},
-    stream.value(),
-    detail::hash_table_allocator_type{default_allocator<char>{}, stream}};
+  mixed_multimap_type hash_table{compute_hash_table_size(build.num_rows()),
+                                 cuco::empty_key{std::numeric_limits<hash_value_type>::max()},
+                                 cuco::empty_value{cudf::detail::JoinNoneValue},
+                                 stream.value(),
+                                 cudf::detail::cuco_allocator{stream}};
 
   // TODO: To add support for nested columns we will need to flatten in many
   // places. However, this probably isn't worth adding any time soon since we
diff --git a/cpp/src/join/mixed_join_semi.cu b/cpp/src/join/mixed_join_semi.cu
index f619ed0d558..edf6c32eadf 100644
--- a/cpp/src/join/mixed_join_semi.cu
+++ b/cpp/src/join/mixed_join_semi.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.
@@ -23,6 +23,7 @@
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/utilities/cuda.cuh>
+#include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/join.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_device_view.cuh>
@@ -172,7 +173,7 @@ std::unique_ptr<rmm::device_uvector<size_type>> mixed_join_semi(
   semi_map_type hash_table{compute_hash_table_size(build.num_rows()),
                            cuco::empty_key{std::numeric_limits<hash_value_type>::max()},
                            cuco::empty_value{cudf::detail::JoinNoneValue},
-                           detail::hash_table_allocator_type{default_allocator<char>{}, stream},
+                           cudf::detail::cuco_allocator{stream},
                            stream.value()};
 
   // Create hash table containing all keys found in right table
@@ -433,7 +434,7 @@ compute_mixed_join_output_size_semi(table_view const& left_equality,
   semi_map_type hash_table{compute_hash_table_size(build.num_rows()),
                            cuco::empty_key{std::numeric_limits<hash_value_type>::max()},
                            cuco::empty_value{cudf::detail::JoinNoneValue},
-                           detail::hash_table_allocator_type{default_allocator<char>{}, stream},
+                           cudf::detail::cuco_allocator{stream},
                            stream.value()};
 
   // Create hash table containing all keys found in right table
diff --git a/cpp/src/reductions/histogram.cu b/cpp/src/reductions/histogram.cu
index 218e2e57420..42ef266a684 100644
--- a/cpp/src/reductions/histogram.cu
+++ b/cpp/src/reductions/histogram.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-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.
@@ -163,12 +163,11 @@ compute_row_frequencies(table_view const& input,
                "Nested types are not yet supported in histogram aggregation.",
                std::invalid_argument);
 
-  auto map = cudf::detail::hash_map_type{
-    compute_hash_table_size(input.num_rows()),
-    cuco::empty_key{-1},
-    cuco::empty_value{std::numeric_limits<size_type>::min()},
-    cudf::detail::hash_table_allocator_type{default_allocator<char>{}, stream},
-    stream.value()};
+  auto map = cudf::detail::hash_map_type{compute_hash_table_size(input.num_rows()),
+                                         cuco::empty_key{-1},
+                                         cuco::empty_value{std::numeric_limits<size_type>::min()},
+                                         cudf::detail::cuco_allocator{stream},
+                                         stream.value()};
 
   auto const preprocessed_input =
     cudf::experimental::row::hash::preprocessed_table::create(input, stream);
diff --git a/cpp/src/search/contains_table.cu b/cpp/src/search/contains_table.cu
index b8ece03c4a0..ce069abcb78 100644
--- a/cpp/src/search/contains_table.cu
+++ b/cpp/src/search/contains_table.cu
@@ -16,7 +16,9 @@
 
 #include <join/join_common_utils.cuh>
 
+#include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/detail/null_mask.hpp>
+#include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/table/experimental/row_operators.cuh>
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
@@ -231,7 +233,7 @@ rmm::device_uvector<bool> contains(table_view const& haystack,
         cuco::empty_key{lhs_index_type{-1}},
         d_equal,
         probing_scheme,
-        detail::hash_table_allocator_type{default_allocator<lhs_index_type>{}, stream},
+        cudf::detail::cuco_allocator{stream},
         stream.value()};
 
       if (haystack_has_nulls && compare_nulls == null_equality::UNEQUAL) {
diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu
index b867df1565a..e73bab1345e 100644
--- a/cpp/src/stream_compaction/distinct.cu
+++ b/cpp/src/stream_compaction/distinct.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.
@@ -54,7 +54,7 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
   auto map = hash_map_type{compute_hash_table_size(input.num_rows()),
                            cuco::empty_key{-1},
                            cuco::empty_value{std::numeric_limits<size_type>::min()},
-                           detail::hash_table_allocator_type{default_allocator<char>{}, stream},
+                           cudf::detail::cuco_allocator{stream},
                            stream.value()};
 
   auto const preprocessed_input =
diff --git a/cpp/src/stream_compaction/distinct_count.cu b/cpp/src/stream_compaction/distinct_count.cu
index ac4811ad279..507bad777eb 100644
--- a/cpp/src/stream_compaction/distinct_count.cu
+++ b/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,13 @@ 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::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},
+                                     cudf::detail::cuco_allocator{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
diff --git a/cpp/src/stream_compaction/stream_compaction_common.hpp b/cpp/src/stream_compaction/stream_compaction_common.hpp
index 18c531e3e69..ceb62d1d059 100644
--- a/cpp/src/stream_compaction/stream_compaction_common.hpp
+++ b/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.
@@ -15,13 +15,11 @@
  */
 #pragma once
 
-#include <cudf/hashing/detail/hash_allocator.cuh>
+#include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/table/row_operators.cuh>
 #include <cudf/table/table_device_view.cuh>
 
-#include <rmm/mr/device/polymorphic_allocator.hpp>
-
 #include <cuco/static_map.cuh>
 
 #include <limits>
@@ -29,10 +27,8 @@
 namespace cudf {
 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>;
+  cuco::static_map<size_type, size_type, cuda::thread_scope_device, cudf::detail::cuco_allocator>;
 
 }  // namespace detail
 }  // namespace cudf
diff --git a/cpp/src/text/bpe/byte_pair_encoding.cuh b/cpp/src/text/bpe/byte_pair_encoding.cuh
index 2a170317909..1a3f8eadea0 100644
--- a/cpp/src/text/bpe/byte_pair_encoding.cuh
+++ b/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.
@@ -20,14 +20,13 @@
 
 #include <cudf/column/column.hpp>
 #include <cudf/column/column_device_view.cuh>
-#include <cudf/hashing/detail/hash_allocator.cuh>
+#include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/hashing/detail/hashing.hpp>
 #include <cudf/hashing/detail/murmurhash3_x86_32.cuh>
 #include <cudf/strings/string_view.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
-#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <cuco/static_map.cuh>
 
@@ -46,8 +45,6 @@ using string_hasher_type = cudf::hashing::detail::MurmurHash3_x86_32<cudf::strin
 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>>;
-
 /**
  * @brief Hasher function used for building and using the cuco static-map
  *
@@ -109,7 +106,7 @@ using merge_pairs_map_type = cuco::experimental::static_map<cudf::size_type,
                                                             cuda::thread_scope_device,
                                                             bpe_equal,
                                                             bpe_probe_scheme,
-                                                            hash_table_allocator_type>;
+                                                            cudf::detail::cuco_allocator>;
 
 /**
  * @brief Hasher function used for building and using the cuco static-map
@@ -166,7 +163,7 @@ using mp_table_map_type = cuco::experimental::static_map<cudf::size_type,
                                                          cuda::thread_scope_device,
                                                          mp_equal,
                                                          mp_probe_scheme,
-                                                         hash_table_allocator_type>;
+                                                         cudf::detail::cuco_allocator>;
 
 }  // namespace detail
 
diff --git a/cpp/src/text/bpe/load_merge_pairs.cu b/cpp/src/text/bpe/load_merge_pairs.cu
index c07d929e98a..3b630886b3e 100644
--- a/cpp/src/text/bpe/load_merge_pairs.cu
+++ b/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.
@@ -42,14 +42,14 @@ namespace {
 std::unique_ptr<detail::merge_pairs_map_type> initialize_merge_pairs_map(
   cudf::column_device_view const& input, rmm::cuda_stream_view stream)
 {
-  auto merge_pairs_map = std::make_unique<merge_pairs_map_type>(
-    static_cast<size_t>(input.size()),
-    cuco::empty_key{-1},
-    cuco::empty_value{-1},
-    bpe_equal{input},
-    bpe_probe_scheme{bpe_hasher{input}},
-    hash_table_allocator_type{default_allocator<char>{}, stream},
-    stream.value());
+  auto merge_pairs_map =
+    std::make_unique<merge_pairs_map_type>(static_cast<size_t>(input.size()),
+                                           cuco::empty_key{-1},
+                                           cuco::empty_value{-1},
+                                           bpe_equal{input},
+                                           bpe_probe_scheme{bpe_hasher{input}},
+                                           cudf::detail::cuco_allocator{stream},
+                                           stream.value());
 
   auto iter = cudf::detail::make_counting_transform_iterator(
     0,
@@ -64,14 +64,13 @@ std::unique_ptr<detail::merge_pairs_map_type> initialize_merge_pairs_map(
 std::unique_ptr<detail::mp_table_map_type> initialize_mp_table_map(
   cudf::column_device_view const& input, rmm::cuda_stream_view stream)
 {
-  auto mp_table_map = std::make_unique<mp_table_map_type>(
-    static_cast<size_t>(input.size()),
-    cuco::empty_key{-1},
-    cuco::empty_value{-1},
-    mp_equal{input},
-    mp_probe_scheme{mp_hasher{input}},
-    hash_table_allocator_type{default_allocator<char>{}, stream},
-    stream.value());
+  auto mp_table_map = std::make_unique<mp_table_map_type>(static_cast<size_t>(input.size()),
+                                                          cuco::empty_key{-1},
+                                                          cuco::empty_value{-1},
+                                                          mp_equal{input},
+                                                          mp_probe_scheme{mp_hasher{input}},
+                                                          cudf::detail::cuco_allocator{stream},
+                                                          stream.value());
 
   auto iter = cudf::detail::make_counting_transform_iterator(
     0,
diff --git a/cpp/src/text/vocabulary_tokenize.cu b/cpp/src/text/vocabulary_tokenize.cu
index a9e8d4d9a24..80f275dba7d 100644
--- a/cpp/src/text/vocabulary_tokenize.cu
+++ b/cpp/src/text/vocabulary_tokenize.cu
@@ -21,13 +21,13 @@
 #include <cudf/column/column.hpp>
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/column/column_factories.hpp>
+#include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/detail/get_value.cuh>
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/sizes_to_offsets_iterator.cuh>
 #include <cudf/detail/utilities/cuda.cuh>
-#include <cudf/hashing/detail/hash_allocator.cuh>
 #include <cudf/hashing/detail/murmurhash3_x86_32.cuh>
 #include <cudf/strings/string_view.cuh>
 #include <cudf/strings/strings_column_view.hpp>
@@ -35,7 +35,6 @@
 #include <cudf/utilities/error.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
-#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <cuco/static_map.cuh>
 
@@ -93,15 +92,14 @@ 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,
+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>;
+                                                           cudf::detail::cuco_allocator>;
 }  // namespace
 }  // namespace detail
 
@@ -150,7 +148,7 @@ 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}},
-    detail::hash_table_allocator_type{default_allocator<char>{}, stream},
+    cudf::detail::cuco_allocator{stream},
     stream.value());
 
   // the row index is the token id (value for each key in the map)