Use cuco_allocator instead of the deprecated default_allocator

rapidsai · Jan 25, 2024 · 4d20609 · 4d20609
1 parent 5e4dcfa
commit 4d20609
Show file tree

Hide file tree

Showing 16 changed files with 85 additions and 105 deletions.
diff --git 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,15 +14,14 @@
  * 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>
 
 #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
@@ -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,16 +32,16 @@
 #include <optional>
 
 // Forward declaration
-template <typename T>
-class default_allocator;
-
 namespace cudf::experimental::row::equality {
 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/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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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.
@@ -17,6 +17,7 @@
 #include <join/join_common_utils.cuh>
 
 #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 +232,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
@@ -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 =