Reduce HashCSR hash join peak memory usage

cpp/src/join/hash_join/dispatch.cuh

@@ -10,24 +10,24 @@
 #include <cudf/detail/row_operator/hashing.cuh>
 #include <cudf/detail/row_operator/primitive_row_operators.cuh>
 
-#include <cuco/pair.cuh>
+#include <cuda/std/utility>
 
 #include <memory>
 #include <utility>
 
 namespace cudf::detail {
 
 /**
- * @brief Equality comparator for cuco hash table probing with row-level equality.
+ * @brief Equality comparator for hash table probing with row-level equality.
  */
 template <typename Equal>
 class pair_equal {
  public:
   pair_equal(Equal check_row_equality) : _check_row_equality{std::move(check_row_equality)} {}
 
   __device__ __forceinline__ bool operator()(
-    cuco::pair<hash_value_type, size_type> const& lhs,
-    cuco::pair<hash_value_type, size_type> const& rhs) const noexcept
+    cuda::std::pair<hash_value_type, size_type> const& lhs,
+    cuda::std::pair<hash_value_type, size_type> const& rhs) const noexcept
   {
     using detail::row::lhs_index_type;
     using detail::row::rhs_index_type;
@@ -41,7 +41,7 @@ class pair_equal {
 };
 
 /**
- * @brief Equality comparator for cuco hash table probing with primitive row equality.
+ * @brief Equality comparator for hash table probing with primitive row equality.
  */
 class primitive_pair_equal {
  public:
@@ -51,8 +51,8 @@ class primitive_pair_equal {
   }
 
   __device__ __forceinline__ bool operator()(
-    cuco::pair<hash_value_type, size_type> const& lhs,
-    cuco::pair<hash_value_type, size_type> const& rhs) const noexcept
+    cuda::std::pair<hash_value_type, size_type> const& lhs,
+    cuda::std::pair<hash_value_type, size_type> const& rhs) const noexcept
   {
     return lhs.first == rhs.first and _check_row_equality(lhs.second, rhs.second);
   }

cpp/src/join/hash_join/hash_csr.cuh

@@ -9,64 +9,88 @@
 #include <cudf/hashing.hpp>
 #include <cudf/types.hpp>
 
-#include <cuco/pair.cuh>
 #include <cuda/atomic>
+#include <cuda/cmath>
 #include <cuda/std/cstdint>
+#include <cuda/std/limits>
+#include <cuda/std/utility>
 
 namespace cudf::detail {
 
-/// One open-addressed slot: the row hash and the index of the build row that claimed it.
-using hash_table_entry_type = cuco::pair<hash_value_type, size_type>;
+/// A row index and the high bits of its hash share one 32-bit hash table slot.
+using hash_table_slot_type = cuda::std::uint32_t;
 
-/// Where a build row landed: the slot it claimed and its rank among the rows sharing that slot.
-/// Computing the rank during the build lets retrieval index straight into the CSR without a
-/// second pass.
-using build_position_type = cuco::pair<cuda::std::uint32_t, size_type>;
-
-/// Device-side view of the open-addressed table, linearly probed with a power-of-two capacity.
+/// Device-side view of the open-addressed table. The low bits hold the representative build
+/// row, and the remaining bits hold a hash fingerprint. Fingerprint matches always undergo
+/// a full row comparison; reducing the fingerprint width cannot introduce false matches.
+/// For N build rows, bit_width(N) bits encode the row index. The all-ones row index is
+/// outside [0, N), leaving the all-ones slot available as the empty sentinel.
 struct hash_table_ref {
-  hash_table_entry_type* entries;
-  cuda::std::uint32_t capacity;  ///< Power of two, so the probe index is a mask instead of a modulo
+  hash_table_slot_type* slots;
+  cuda::std::uint32_t capacity;
+  cuda::std::uint32_t row_mask;
+  cuda::fast_mod_div<cuda::std::uint32_t> modulo;
 
-  __device__ cuda::std::uint32_t mask() const { return capacity - 1; }
+  template <typename Equal>
+  __device__ bool equal(cuda::std::pair<hash_value_type, size_type> key,
+                        hash_table_slot_type slot,
+                        Equal check_row_equality) const
+  {
+    return ((key.first ^ slot) & ~row_mask) == 0 &&
+           check_row_equality(key, {key.first, static_cast<size_type>(slot & row_mask)});
+  }
 
   template <typename Equal>
-  __device__ cuda::std::uint32_t insert(hash_table_entry_type key, Equal equal) const
+  __device__ size_type insert(cuda::std::pair<hash_value_type, size_type> key,
+                              Equal equal_rows) const
   {
+    auto const desired = (key.first & ~row_mask) | static_cast<cuda::std::uint32_t>(key.second);
+    auto slot          = key.first % modulo;
     for (cuda::std::uint32_t step = 0; step < capacity; ++step) {
-      auto const slot = (static_cast<cuda::std::uint32_t>(key.first) + step) & mask();
-      auto entry_ref =
-        cuda::atomic_ref<hash_table_entry_type, cuda::thread_scope_device>{entries[slot]};
-      auto old = hash_table_entry_type{hash_value_type{-1}, size_type{CUDF_SIZE_TYPE_SENTINEL}};
-      if (entry_ref.compare_exchange_strong(old, key, cuda::memory_order_relaxed)) { return slot; }
-      if (equal(key, old)) { return slot; }
+      auto slot_ref =
+        cuda::atomic_ref<hash_table_slot_type, cuda::thread_scope_device>{slots[slot]};
+      auto old = cuda::std::numeric_limits<hash_table_slot_type>::max();
+      if (slot_ref.compare_exchange_strong(old, desired, cuda::memory_order_relaxed)) {
+        return key.second;
+      }
+      if (equal(key, old, equal_rows)) { return static_cast<size_type>(old & row_mask); }
+      ++slot;
+      if (slot == capacity) { slot = 0; }
     }
-    return capacity;
+    return CUDF_SIZE_TYPE_SENTINEL;
   }
 
-  template <typename Equal>
-  __device__ cuda::std::uint32_t find(hash_table_entry_type key, Equal equal) const
+  template <bool IsBuild = false, typename Equal>
+  __device__ size_type find(cuda::std::pair<hash_value_type, size_type> key, Equal equal_rows) const
   {
+    auto slot = key.first % modulo;
     for (cuda::std::uint32_t step = 0; step < capacity; ++step) {
-      auto const slot    = (static_cast<cuda::std::uint32_t>(key.first) + step) & mask();
-      auto const current = entries[slot];
-      if (current.second == CUDF_SIZE_TYPE_SENTINEL) { return capacity; }
-      if (equal(key, current)) { return slot; }
+      auto const current = slots[slot];
+      if (current == cuda::std::numeric_limits<hash_table_slot_type>::max()) {
+        return CUDF_SIZE_TYPE_SENTINEL;
+      }
+      // Under null_equality::UNEQUAL a nested row containing nulls need not equal itself.
+      // The fill pass must still find the row that claimed this slot during construction.
+      if constexpr (IsBuild) {
+        if (static_cast<size_type>(current & row_mask) == key.second) { return key.second; }
+      }
+      if (equal(key, current, equal_rows)) { return static_cast<size_type>(current & row_mask); }
+      ++slot;
+      if (slot == capacity) { slot = 0; }
     }
-    return capacity;
+    return CUDF_SIZE_TYPE_SENTINEL;
   }
 };
 
+/// CSR segments are indexed by representative build row, including zero-length segments for
+/// rows that did not claim a hash table slot. This avoids an offset for every empty hash slot.
 struct csr_ref {
-  size_type const* cumulative_ends;
+  size_type const* offsets;
   size_type const* values;
 
-  __device__ size_type begin(size_type slot) const
-  {
-    return slot == 0 ? size_type{0} : cumulative_ends[slot - 1];
-  }
+  __device__ size_type begin(size_type row) const { return offsets[row]; }
 
-  __device__ size_type size(size_type slot) const { return cumulative_ends[slot] - begin(slot); }
+  __device__ size_type size(size_type row) const { return offsets[row + 1] - offsets[row]; }
 };
 
 }  // namespace cudf::detail