From e9068eb2929b4a3d233f50f0931c7e4b4a5d4642 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Tue, 16 Nov 2021 16:07:46 -0500
Subject: [PATCH 1/4] Use stream alloactor adaptor for hash join table

---
 cpp/cmake/thirdparty/get_cucollections.cmake | 4 ++--
 cpp/src/join/hash_join.cu                    | 3 ++-
 cpp/src/join/join_common_utils.hpp           | 5 ++++-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/cpp/cmake/thirdparty/get_cucollections.cmake b/cpp/cmake/thirdparty/get_cucollections.cmake
index b6cb9757ae8..53c74dc3bcf 100644
--- a/cpp/cmake/thirdparty/get_cucollections.cmake
+++ b/cpp/cmake/thirdparty/get_cucollections.cmake
@@ -20,8 +20,8 @@ function(find_and_configure_cucollections)
     # cuCollections doesn't have a version yet
     cuco 0.0
     GLOBAL_TARGETS cuco::cuco
-    CPM_ARGS GITHUB_REPOSITORY NVIDIA/cuCollections
-    GIT_TAG f0eecb203590f1f4ac4a9f1700229f4434ac64dc
+    CPM_ARGS GITHUB_REPOSITORY PointKernel/cuCollections
+    GIT_TAG allocator-without-stream
     OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF"
   )
 
diff --git a/cpp/src/join/hash_join.cu b/cpp/src/join/hash_join.cu
index e4bd1938ecc..a540dde1041 100644
--- a/cpp/src/join/hash_join.cu
+++ b/cpp/src/join/hash_join.cu
@@ -293,7 +293,8 @@ hash_join::hash_join_impl::hash_join_impl(cudf::table_view const& build,
     _hash_table{compute_hash_table_size(build.num_rows()),
                 std::numeric_limits<hash_value_type>::max(),
                 cudf::detail::JoinNoneValue,
-                stream.value()}
+                stream.value(),
+                detail::hash_table_allocator_type{default_allocator<char>{}, stream}}
 {
   CUDF_FUNC_RANGE();
   CUDF_EXPECTS(0 != build.num_columns(), "Hash join build table is empty");
diff --git a/cpp/src/join/join_common_utils.hpp b/cpp/src/join/join_common_utils.hpp
index d6eb5e93a98..4231bf46017 100644
--- a/cpp/src/join/join_common_utils.hpp
+++ b/cpp/src/join/join_common_utils.hpp
@@ -22,6 +22,7 @@
 #include <hash/concurrent_unordered_multimap.cuh>
 
 #include <cuco/static_multimap.cuh>
+#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <limits>
 
@@ -38,11 +39,13 @@ using pair_type = cuco::pair_type<hash_value_type, size_type>;
 
 using hash_type = cuco::detail::MurmurHash3_32<hash_value_type>;
 
+using hash_table_allocator_type = rmm::mr::stream_allocator_adaptor<default_allocator<char>>;
+
 using multimap_type =
   cuco::static_multimap<hash_value_type,
                         size_type,
                         cuda::thread_scope_device,
-                        default_allocator<char>,
+                        hash_table_allocator_type,
                         cuco::double_hashing<DEFAULT_JOIN_CG_SIZE, hash_type, hash_type>>;
 
 using row_hash = cudf::row_hasher<default_hash>;

From 91c5130f15ee44b3bbb6dff87dcb42c8439f0de5 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Thu, 2 Dec 2021 14:02:33 -0500
Subject: [PATCH 2/4] Update git tag

---
 cpp/cmake/thirdparty/get_cucollections.cmake | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/cmake/thirdparty/get_cucollections.cmake b/cpp/cmake/thirdparty/get_cucollections.cmake
index 53c74dc3bcf..b58bdb55de3 100644
--- a/cpp/cmake/thirdparty/get_cucollections.cmake
+++ b/cpp/cmake/thirdparty/get_cucollections.cmake
@@ -20,8 +20,8 @@ function(find_and_configure_cucollections)
     # cuCollections doesn't have a version yet
     cuco 0.0
     GLOBAL_TARGETS cuco::cuco
-    CPM_ARGS GITHUB_REPOSITORY PointKernel/cuCollections
-    GIT_TAG allocator-without-stream
+    CPM_ARGS GITHUB_REPOSITORY NVIDIA/cuCollections
+    GIT_TAG 6433e8ad7571f14cc5384051b049029c60dd1ce0
     OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF"
   )
 

From fb003969e0760d6ef3aa16d58f1a2ed8d5e99c54 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Fri, 3 Dec 2021 09:57:04 -0500
Subject: [PATCH 3/4] Header cleanups

---
 cpp/src/join/join_common_utils.hpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/cpp/src/join/join_common_utils.hpp b/cpp/src/join/join_common_utils.hpp
index 4231bf46017..a0fe1563edf 100644
--- a/cpp/src/join/join_common_utils.hpp
+++ b/cpp/src/join/join_common_utils.hpp
@@ -19,10 +19,9 @@
 #include <cudf/table/row_operators.cuh>
 #include <cudf/table/table_view.hpp>
 
-#include <hash/concurrent_unordered_multimap.cuh>
+#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <cuco/static_multimap.cuh>
-#include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <limits>
 

From 4a2aba772428e314ae8e8879c94bb386f86ebeef Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Fri, 3 Dec 2021 11:01:21 -0500
Subject: [PATCH 4/4] Include necessary headers

---
 cpp/src/join/join_common_utils.hpp | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/cpp/src/join/join_common_utils.hpp b/cpp/src/join/join_common_utils.hpp
index a0fe1563edf..8e76def72e1 100644
--- a/cpp/src/join/join_common_utils.hpp
+++ b/cpp/src/join/join_common_utils.hpp
@@ -15,10 +15,14 @@
  */
 #pragma once
 
+#include <cudf/detail/utilities/device_atomics.cuh>
 #include <cudf/detail/utilities/hash_functions.cuh>
 #include <cudf/table/row_operators.cuh>
 #include <cudf/table/table_view.hpp>
 
+#include <hash/hash_allocator.cuh>
+#include <hash/helper_functions.cuh>
+
 #include <rmm/mr/device/polymorphic_allocator.hpp>
 
 #include <cuco/static_multimap.cuh>