diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 3066fdf32e77..47377902f3a1 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -865,7 +865,6 @@ add_library( src/join/hash_join/full_join_match_context.cpp src/join/hash_join/full_join_retrieve.cu src/join/hash_join/full_join_size.cu - src/join/hash_join/full_join_size_impl.cu src/join/hash_join/hash_join.cu src/join/hash_join/inner_join_match_context.cpp src/join/hash_join/inner_join_retrieve.cu @@ -874,14 +873,10 @@ add_library( src/join/hash_join/left_join_retrieve.cu src/join/hash_join/left_join_size.cu src/join/hash_join/match_context.cu - src/join/hash_join/partitioned_count.cu - src/join/hash_join/partitioned_count_outer.cu src/join/hash_join/partitioned_full_join.cu src/join/hash_join/partitioned_inner_join.cu src/join/hash_join/partitioned_join_retrieve.cu src/join/hash_join/partitioned_left_join.cu - src/join/hash_join/partitioned_retrieve.cu - src/join/hash_join/partitioned_retrieve_outer.cu src/join/join.cu src/join/join_utils.cu src/join/key_remapping.cu diff --git a/cpp/include/cudf/detail/join/hash_join.hpp b/cpp/include/cudf/detail/join/hash_join.hpp index fc32d86c0e5d..9b2ed9674bf3 100644 --- a/cpp/include/cudf/detail/join/hash_join.hpp +++ b/cpp/include/cudf/detail/join/hash_join.hpp @@ -176,7 +176,7 @@ class hash_join { rmm::device_async_resource_ref mr) const; private: - bool const _is_empty; ///< true if `_hash_table` is empty + bool const _is_empty; ///< true if the build-side (right) table is empty bool const _has_nulls; ///< true if nulls are present in either right table or any left table cudf::null_equality const _nulls_equal; ///< whether to consider nulls as equal cudf::table_view _right; ///< input table to build the hash map diff --git a/cpp/src/join/hash_join/common.cuh b/cpp/src/join/hash_join/common.cuh index 517573412e65..e8fac29b3493 100644 --- a/cpp/src/join/hash_join/common.cuh +++ b/cpp/src/join/hash_join/common.cuh @@ -21,22 +21,9 @@ namespace cudf::detail { using hash_join_hasher = cudf::hashing::detail::MurmurHash3_x86_32; -using hash_table_t = typename cudf::detail::hash_join::impl::hash_table_t; bool is_trivial_join(table_view const& left, table_view const& right, join_kind join_type); void validate_hash_join_probe(table_view const& right, table_view const& left, bool has_nulls); -std::unique_ptr> make_join_match_counts( - table_view const& right, - std::shared_ptr const& preprocessed_right, - cudf::detail::hash_table_t const& hash_table, - bool is_empty, - bool has_nulls, - null_equality compare_nulls, - join_kind join, - table_view const& left, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - } // namespace cudf::detail diff --git a/cpp/src/join/hash_join/dispatch.cuh b/cpp/src/join/hash_join/dispatch.cuh index 2092817e2931..f68893c45cf2 100644 --- a/cpp/src/join/hash_join/dispatch.cuh +++ b/cpp/src/join/hash_join/dispatch.cuh @@ -1,12 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include "common.cuh" -#include #include #include #include @@ -41,17 +40,6 @@ class pair_equal { Equal _check_row_equality; }; -/** - * @brief Extracts the right-side row index from a cuco hash table slot. - */ -struct output_fn { - __device__ constexpr cudf::size_type operator()( - cuco::pair const& slot) const - { - return slot.second; - } -}; - /** * @brief Equality comparator for cuco hash table probing with primitive row equality. */ diff --git a/cpp/src/join/hash_join/full_join_size_impl.cu b/cpp/src/join/hash_join/full_join_size_impl.cu deleted file mode 100644 index 5c9403a020c6..000000000000 --- a/cpp/src/join/hash_join/full_join_size_impl.cu +++ /dev/null @@ -1,98 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "retrieve_impl.cuh" - -#include -#include - -#include -#include - -#include -#include -#include -#include - -#include - -namespace cudf::detail { - -namespace { -std::size_t compute_left_join_complement_size(cudf::device_span right_indices, - size_type left_table_row_count, - size_type right_table_row_count, - cuda::stream_ref stream) -{ - if (left_table_row_count == 0) { return right_table_row_count; } - - auto invalid_index_map = - std::make_unique>(right_table_row_count, stream); - thrust::uninitialized_fill( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - invalid_index_map->begin(), - invalid_index_map->end(), - int32_t{1}); - - valid_range valid(0, right_table_row_count); - - thrust::scatter_if(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::make_constant_iterator(0), - cuda::make_constant_iterator(0) + right_indices.size(), - right_indices.begin(), - right_indices.begin(), - invalid_index_map->begin(), - valid); - - return cudf::detail::count_if( - invalid_index_map->begin(), invalid_index_map->end(), cuda::std::identity{}, stream); -} -} // namespace - -std::size_t get_full_join_size( - cudf::table_view const& right_table, - cudf::table_view const& left_table, - std::shared_ptr const& preprocessed_right, - std::shared_ptr const& preprocessed_left, - cudf::detail::hash_table_t const& hash_table, - bool has_nulls, - null_equality compare_nulls, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - std::size_t join_size = compute_join_output_size(right_table, - left_table, - preprocessed_right, - preprocessed_left, - hash_table, - has_nulls, - compare_nulls, - stream); - - if (join_size == 0) { return join_size; } - - auto right_indices = std::make_unique>(join_size, stream, mr); - - auto const out_build_begin = - cuda::make_transform_output_iterator(right_indices->begin(), output_fn{}); - - retrieve_left_join_build_indices(right_table, - left_table, - preprocessed_right, - preprocessed_left, - hash_table, - has_nulls, - compare_nulls, - out_build_begin, - stream); - - auto const left_table_row_count = left_table.num_rows(); - auto const right_table_row_count = right_table.num_rows(); - - return join_size + compute_left_join_complement_size( - *right_indices, left_table_row_count, right_table_row_count, stream); -} - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/hash_csr.cuh b/cpp/src/join/hash_join/hash_csr.cuh new file mode 100644 index 000000000000..f44718eaa347 --- /dev/null +++ b/cpp/src/join/hash_join/hash_csr.cuh @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include + +#include + +namespace cudf::detail { + +using hash_csr_key_type = cuco::pair; +using hash_csr_build_position_type = cuco::pair; + +struct hash_csr_map_view { + hash_csr_key_type* entries; + std::uint32_t capacity; + std::uint32_t mask; + + template + __device__ std::uint32_t insert(hash_csr_key_type key, Equal equal) const + { + for (std::uint32_t step = 0; step < capacity; ++step) { + auto const slot = (static_cast(key.first) + step) & mask; + auto entry_ref = + cuda::atomic_ref{entries[slot]}; + auto old = hash_csr_key_type{hash_value_type{-1}, CUDF_SIZE_TYPE_SENTINEL}; + if (entry_ref.compare_exchange_strong(old, key, cuda::memory_order_relaxed)) { return slot; } + if (equal(key, old)) { return slot; } + } + return capacity; + } + + template + __device__ std::uint32_t find(hash_csr_key_type key, Equal equal) const + { + for (std::uint32_t step = 0; step < capacity; ++step) { + auto const slot = (static_cast(key.first) + step) & mask; + auto const current = entries[slot]; + if (current.second == CUDF_SIZE_TYPE_SENTINEL) { return capacity; } + if (equal(key, current)) { return slot; } + } + return capacity; + } +}; + +struct hash_csr_view { + size_type const* cumulative_ends; + 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 size(size_type slot) const { return cumulative_ends[slot] - begin(slot); } +}; + +} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/hash_csr_kernels.cuh b/cpp/src/join/hash_join/hash_csr_kernels.cuh new file mode 100644 index 000000000000..d9c06c5e4c6f --- /dev/null +++ b/cpp/src/join/hash_join/hash_csr_kernels.cuh @@ -0,0 +1,267 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "hash_csr.cuh" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +namespace cudf::detail { + +constexpr thread_index_type hash_csr_block_size = 256; +constexpr thread_index_type hash_csr_warps_per_block = + hash_csr_block_size / cudf::detail::warp_size; +constexpr thread_index_type hash_csr_outputs_per_lane = 32; + +template +CUDF_KERNEL void hash_csr_build_count_kernel(size_type num_rows, + bitmask_type const* valid_rows, + hash_csr_build_position_type* build_positions, + size_type* slot_counts, + hash_csr_map_view map, + Equal equal, + Hasher hasher) +{ + auto const stride = grid_1d::grid_stride(); + for (auto row = grid_1d::global_thread_id(); row < num_rows; row += stride) { + auto const index = static_cast(row); + if (valid_rows != nullptr && !cudf::bit_is_set(valid_rows, index)) { + build_positions[index] = {std::uint32_t{-1}, CUDF_SIZE_TYPE_SENTINEL}; + continue; + } + + auto const slot = map.insert(hash_csr_key_type{hasher(index), index}, equal); + if (slot == map.capacity) { + build_positions[index] = {std::uint32_t{-1}, CUDF_SIZE_TYPE_SENTINEL}; + continue; + } + auto slot_count_ref = cuda::atomic_ref{slot_counts[slot]}; + auto const rank = slot_count_ref.fetch_add(size_type{1}, cuda::memory_order_relaxed); + build_positions[index] = {slot, rank}; + } +} + +CUDF_KERNEL void hash_csr_build_fill_kernel(size_type num_rows, + hash_csr_build_position_type const* build_positions, + size_type const* cumulative_ends, + size_type* values) +{ + auto const stride = grid_1d::grid_stride(); + for (auto row = grid_1d::global_thread_id(); row < num_rows; row += stride) { + auto const index = static_cast(row); + auto const position = build_positions[index]; + if (position.first == std::uint32_t{-1}) { continue; } + auto const slot = position.first; + auto const rank = position.second; + auto const begin = slot == 0 ? size_type{0} : cumulative_ends[slot - 1]; + values[begin + rank] = index; + } +} + +template +CUDF_KERNEL void hash_csr_probe_count_kernel(size_type num_rows, + bitmask_type const* valid_rows, + size_type* probe_slots, + size_type* match_counts, + std::uint32_t* matched_slots, + unsigned long long* matched_build_rows, + hash_csr_map_view map, + hash_csr_view csr, + Equal equal, + Hasher hasher) +{ + auto const stride = grid_1d::grid_stride(); + for (auto row = grid_1d::global_thread_id(); row < num_rows; row += stride) { + auto const index = static_cast(row); + auto slot = map.capacity; + if (valid_rows == nullptr || cudf::bit_is_set(valid_rows, index)) { + slot = map.find(hash_csr_key_type{hasher(index), index}, equal); + } + + auto const found = slot != map.capacity; + auto const count = found ? csr.size(static_cast(slot)) : size_type{0}; + if (probe_slots != nullptr) { + probe_slots[index] = found ? static_cast(slot) : CUDF_SIZE_TYPE_SENTINEL; + } + if (match_counts != nullptr) { + match_counts[index] = IsOuter ? cuda::std::max(count, size_type{1}) : count; + } + + if (found && matched_slots != nullptr) { + auto matched_slot_ref = + cuda::atomic_ref{matched_slots[slot]}; + auto expected = std::uint32_t{0}; + if (matched_slot_ref.compare_exchange_strong( + expected, std::uint32_t{1}, cuda::memory_order_relaxed)) { + cuda::atomic_ref{*matched_build_rows} + .fetch_add(static_cast(count), cuda::memory_order_relaxed); + } + } + } +} + +template +void launch_hash_csr_build_count(size_type num_rows, + bitmask_type const* valid_rows, + hash_csr_build_position_type* build_positions, + size_type* slot_counts, + hash_csr_map_view map, + Equal equal, + Hasher hasher, + cuda::stream_ref stream) +{ + if (num_rows == 0) { return; } + auto const config = grid_1d{num_rows, hash_csr_block_size}; + hash_csr_build_count_kernel<<>>( + num_rows, valid_rows, build_positions, slot_counts, map, equal, hasher); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +[[maybe_unused]] static void launch_hash_csr_build_fill( + size_type num_rows, + hash_csr_build_position_type const* build_positions, + size_type const* cumulative_ends, + size_type* values, + cuda::stream_ref stream) +{ + if (num_rows == 0) { return; } + auto const config = grid_1d{num_rows, hash_csr_block_size}; + hash_csr_build_fill_kernel<<>>( + num_rows, build_positions, cumulative_ends, values); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +template +void launch_hash_csr_probe_count(size_type num_rows, + bitmask_type const* valid_rows, + size_type* probe_slots, + size_type* match_counts, + std::uint32_t* matched_slots, + unsigned long long* matched_build_rows, + hash_csr_map_view map, + hash_csr_view csr, + Equal equal, + Hasher hasher, + cuda::stream_ref stream) +{ + if (num_rows == 0) { return; } + auto const config = grid_1d{num_rows, hash_csr_block_size}; + hash_csr_probe_count_kernel + <<>>(num_rows, + valid_rows, + probe_slots, + match_counts, + matched_slots, + matched_build_rows, + map, + csr, + equal, + hasher); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +template +CUDF_KERNEL void hash_csr_retrieve_kernel(std::int64_t output_size, + size_type num_probe_rows, + std::int64_t outputs_per_warp, + std::int64_t const* offsets, + size_type const* probe_slots, + hash_csr_view csr, + size_type left_index_offset, + size_type* left_indices, + size_type* right_indices) +{ + auto const warp = cooperative_groups::tiled_partition( + cooperative_groups::this_thread_block()); + auto const lane_id = static_cast(warp.thread_rank()); + auto const warp_in_block = static_cast(threadIdx.x) / cudf::detail::warp_size; + auto const global_warp = + static_cast(blockIdx.x) * hash_csr_warps_per_block + warp_in_block; + auto const range_begin = outputs_per_warp * global_warp; + if (range_begin >= output_size) { return; } + auto const range_end = cuda::std::min(range_begin + outputs_per_warp, output_size); + + size_type endpoint_probe{}; + if (lane_id < 2) { + auto const endpoint = lane_id == 0 ? range_begin : range_end - 1; + endpoint_probe = static_cast( + cuda::std::upper_bound(offsets, offsets + num_probe_rows + 1, endpoint) - offsets - 1); + } + auto const first_probe = warp.shfl(endpoint_probe, 0); + auto const last_probe = warp.shfl(endpoint_probe, 1); + +#pragma unroll + for (thread_index_type item = 0; item < hash_csr_outputs_per_lane; ++item) { + auto const output_index = range_begin + lane_id + item * cudf::detail::warp_size; + if (output_index < range_end) { + auto const probe_row = + first_probe == last_probe + ? first_probe + : static_cast(cuda::std::upper_bound(offsets + first_probe, + offsets + last_probe + 2, + output_index) - + offsets - 1); + auto const slot = probe_slots[probe_row]; + left_indices[output_index] = probe_row + left_index_offset; + if constexpr (IsOuter) { + if (slot == CUDF_SIZE_TYPE_SENTINEL) { + right_indices[output_index] = JoinNoMatch; + continue; + } + } + auto const local_match = static_cast(output_index - offsets[probe_row]); + right_indices[output_index] = csr.values[csr.begin(slot) + local_match]; + } + } +} + +template +void launch_hash_csr_retrieve(std::int64_t output_size, + size_type num_probe_rows, + std::int64_t const* offsets, + size_type const* probe_slots, + hash_csr_view csr, + size_type left_index_offset, + size_type* left_indices, + size_type* right_indices, + cuda::stream_ref stream) +{ + if (output_size == 0) { return; } + auto const min_blocks = size_type{2} * cudf::detail::num_multiprocessors(); + constexpr auto outputs_per_block = + hash_csr_warps_per_block * cudf::detail::warp_size * hash_csr_outputs_per_lane; + auto const requested_blocks = + cudf::util::div_rounding_up_safe(output_size, static_cast(outputs_per_block)); + auto const num_blocks = + static_cast(cuda::std::max(requested_blocks, min_blocks)); + auto const num_warps = static_cast(num_blocks) * hash_csr_warps_per_block; + auto const outputs_per_warp = cudf::util::div_rounding_up_safe(output_size, num_warps); + + hash_csr_retrieve_kernel + <<>>(output_size, + num_probe_rows, + outputs_per_warp, + offsets, + probe_slots, + csr, + left_index_offset, + left_indices, + right_indices); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/hash_join.cu b/cpp/src/join/hash_join/hash_join.cu index 89d990bec73b..06fbf59dccd2 100644 --- a/cpp/src/join/hash_join/hash_join.cu +++ b/cpp/src/join/hash_join/hash_join.cu @@ -4,10 +4,11 @@ */ #include "common.cuh" +#include "dispatch.cuh" +#include "hash_csr_kernels.cuh" #include "join/join_common_utils.cuh" #include -#include #include #include #include @@ -19,12 +20,15 @@ #include #include -#include +#include -#include +#include +#include +#include #include #include +#include #include namespace cudf::detail { @@ -59,43 +63,19 @@ void validate_hash_join_probe(table_view const& right, table_view const& left, b } namespace { -void build_hash_join( - cudf::table_view const& right, - std::shared_ptr const& preprocessed_right, - cudf::detail::hash_table_t& hash_table, - bool has_nested_nulls, - null_equality nulls_equal, - [[maybe_unused]] bitmask_type const* bitmask, - cuda::stream_ref stream) +std::uint32_t hash_csr_capacity(size_type rows, double load_factor) { - CUDF_EXPECTS(0 != right.num_columns(), "Selected right dataset is empty", std::invalid_argument); - CUDF_EXPECTS(0 != right.num_rows(), "Right side table has no rows", std::invalid_argument); - - auto insert_rows = [&](auto const& right, auto const& d_hasher) { - auto const iter = cudf::detail::make_counting_transform_iterator(0, pair_fn{d_hasher}); - - if (nulls_equal == cudf::null_equality::EQUAL or not nullable(right)) { - hash_table.insert(iter, iter + right.num_rows(), stream.get()); - } else { - auto const stencil = cuda::counting_iterator{0}; - auto const pred = row_is_valid{bitmask}; - - hash_table.insert_if(iter, iter + right.num_rows(), stencil, pred, stream.get()); - } - }; - - auto const nulls = nullate::DYNAMIC{has_nested_nulls}; - - if (cudf::detail::is_primitive_row_op_compatible(right)) { - auto const d_hasher = cudf::detail::row::primitive::row_hasher{nulls, preprocessed_right}; - - insert_rows(right, d_hasher); - } else { - auto const row_hash = detail::row::hash::row_hasher{preprocessed_right}; - auto const d_hasher = row_hash.device_hasher(nulls); - - insert_rows(right, d_hasher); - } + auto const checked = checked_load_factor(load_factor); + auto const requested = std::max(static_cast(rows) + 1, + std::ceil(static_cast(rows) / checked)); + CUDF_EXPECTS(requested <= std::numeric_limits::max(), + "HashCSR table capacity is not representable", + std::overflow_error); + auto const capacity = cuda::std::bit_ceil(static_cast(requested)); + CUDF_EXPECTS(capacity <= std::numeric_limits::max(), + "HashCSR table capacity is not representable", + std::overflow_error); + return static_cast(capacity); } } // namespace @@ -119,32 +99,60 @@ hash_join::hash_join(cudf::table_view const& right, : _has_nulls(has_nulls), _is_empty{right.num_rows() == 0}, _nulls_equal{compare_nulls}, - _impl{std::make_unique(impl{typename impl::hash_table_t{ - cuco::extent{static_cast(right.num_rows())}, - checked_load_factor(load_factor), - cuco::empty_key{cuco::pair{std::numeric_limits::max(), cudf::JoinNoMatch}}, - {}, - {}, - {}, - {}, - rmm::mr::polymorphic_allocator{std::move(mr)}, - stream.get()}})}, _right{right}, - _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(_right, stream)} + _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(_right, stream)}, + _impl{std::make_unique( + hash_csr_capacity(right.num_rows(), load_factor), right.num_rows(), stream, std::move(mr))} { CUDF_FUNC_RANGE(); CUDF_EXPECTS(0 != right.num_columns(), "Hash join right table is empty", std::invalid_argument); if (_is_empty) { return; } - auto const row_bitmask = - cudf::detail::bitmask_and(right, stream, cudf::get_current_device_resource_ref()).first; - cudf::detail::build_hash_join(_right, - _preprocessed_right, - _impl->_hash_table, - _has_nulls, - _nulls_equal, - reinterpret_cast(row_bitmask.data()), + CUDF_CUDA_TRY(cudaMemsetAsync( + _impl->entries.data(), 0xff, _impl->entries.size() * sizeof(hash_csr_key_type), stream.get())); + CUDF_CUDA_TRY(cudaMemsetAsync(_impl->cumulative_ends.data(), + 0, + _impl->cumulative_ends.size() * sizeof(size_type), + stream.get())); + + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const row_bitmask = cudf::detail::bitmask_and(right, stream, temp_mr).first; + auto const valid_rows = _nulls_equal == null_equality::UNEQUAL + ? static_cast(row_bitmask.data()) + : nullptr; + rmm::device_uvector build_positions( + right.num_rows(), stream, temp_mr); + auto build = [&](auto equality, auto hasher) { + launch_hash_csr_build_count(right.num_rows(), + valid_rows, + build_positions.data(), + _impl->cumulative_ends.data(), + _impl->map_view(), + equality, + hasher, stream); + }; + dispatch_join_comparator( + right, right, _preprocessed_right, _preprocessed_right, _has_nulls, _nulls_equal, build); + std::size_t temp_storage_bytes{}; + CUDF_CUDA_TRY(cub::DeviceScan::InclusiveSum(nullptr, + temp_storage_bytes, + _impl->cumulative_ends.data(), + _impl->cumulative_ends.data(), + _impl->capacity, + stream.get())); + rmm::device_buffer temp_storage(temp_storage_bytes, stream, temp_mr); + CUDF_CUDA_TRY(cub::DeviceScan::InclusiveSum(temp_storage.data(), + temp_storage_bytes, + _impl->cumulative_ends.data(), + _impl->cumulative_ends.data(), + _impl->capacity, + stream.get())); + launch_hash_csr_build_fill(right.num_rows(), + build_positions.data(), + _impl->cumulative_ends.data(), + _impl->values.data(), + stream); } template hash_join::hash_join( diff --git a/cpp/src/join/hash_join/hash_join_impl.cuh b/cpp/src/join/hash_join/hash_join_impl.cuh index 18a65be10f76..4d9c2f9487d6 100644 --- a/cpp/src/join/hash_join/hash_join_impl.cuh +++ b/cpp/src/join/hash_join/hash_join_impl.cuh @@ -1,62 +1,47 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once +#include "hash_csr.cuh" + #include #include -#include +#include -#include -#include +#include +#include namespace cudf::detail { template struct hash_join::impl { - struct always_not_equal { - __device__ constexpr bool operator()( - cuco::pair const&, - cuco::pair const&) const noexcept - { - // multiset always insert - return false; - } - }; - - struct hasher1 { - __device__ constexpr hash_value_type operator()( - cuco::pair const& key) const noexcept - { - return key.first; - } - }; - - struct hasher2 { - hasher2(hash_value_type seed) : _hash{seed} {} - - __device__ constexpr hash_value_type operator()( - cuco::pair const& key) const noexcept - { - return _hash(key.first); - } - - private: - Hasher _hash; - }; - - using hash_table_t = - cuco::static_multiset, - cuco::extent, - cuda::thread_scope_device, - always_not_equal, - cuco::double_hashing, - rmm::mr::polymorphic_allocator, - cuco::storage<2>>; - - hash_table_t _hash_table; + impl(std::uint32_t capacity, + size_type rows, + cuda::stream_ref stream, + cuda::mr::any_resource mr) + : _mr{std::move(mr)}, + entries(capacity, stream, _mr), + cumulative_ends(capacity, stream, _mr), + values(rows, stream, _mr), + capacity{capacity} + { + } + + hash_csr_map_view map_view() const + { + return {const_cast(entries.data()), capacity, capacity - 1}; + } + + hash_csr_view csr_view() const { return {cumulative_ends.data(), values.data()}; } + + cuda::mr::any_resource _mr; + rmm::device_uvector entries; + rmm::device_uvector cumulative_ends; + rmm::device_uvector values; + std::uint32_t capacity; }; } // namespace cudf::detail diff --git a/cpp/src/join/hash_join/kernels_common.cuh b/cpp/src/join/hash_join/kernels_common.cuh deleted file mode 100644 index cd74a5bde4e1..000000000000 --- a/cpp/src/join/hash_join/kernels_common.cuh +++ /dev/null @@ -1,24 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -// Custom hash-join probe kernels that give cudf direct control over kernel launches. -// Uses the cuco ref type for hash-table access (storage, probing scheme, predicate). - -#pragma once - -#include "join/join_common_utils.hpp" - -#include -#include -#include - -#include - -namespace cudf::detail { - -/// The probe key type stored in the hash table: {hash_value, row_index}. -using probe_key_type = cuco::pair; - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/match_context.cu b/cpp/src/join/hash_join/match_context.cu index e6812f55b3be..e9a4429bf79e 100644 --- a/cpp/src/join/hash_join/match_context.cu +++ b/cpp/src/join/hash_join/match_context.cu @@ -5,35 +5,28 @@ #include "common.cuh" #include "dispatch.cuh" -#include "join/join_common_utils.cuh" -#include "partitioned_count_kernels.hpp" +#include "hash_csr_kernels.cuh" -#include +#include +#include #include #include -#include #include -#include namespace cudf::detail { -std::unique_ptr> make_join_match_counts( - table_view const& right, - std::shared_ptr const& preprocessed_right, - cudf::detail::hash_table_t const& hash_table, - bool is_empty, - bool has_nulls, - null_equality compare_nulls, +template +std::unique_ptr> hash_join::make_match_counts( join_kind join, table_view const& left, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + rmm::device_async_resource_ref mr) const { auto match_counts = std::make_unique>(left.num_rows(), stream, mr); - if (is_empty) { + if (_is_empty) { thrust::fill(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), match_counts->begin(), match_counts->end(), @@ -41,60 +34,52 @@ std::unique_ptr> make_join_match_counts( return match_counts; } - CUDF_EXPECTS(has_nulls || !cudf::has_nested_nulls(left), + CUDF_EXPECTS(_has_nulls || !cudf::has_nested_nulls(left), "Left table has nulls while right table was not hashed with null check.", std::invalid_argument); auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(left, stream); - auto const left_table_num_rows = left.num_rows(); + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const row_bitmask = cudf::detail::bitmask_and(left, stream, temp_mr).first; + auto const valid_rows = _nulls_equal == null_equality::UNEQUAL + ? static_cast(row_bitmask.data()) + : nullptr; - auto count_matches = [&](auto equality, auto d_hasher) { - // Precompute left keys: {hash(row_idx), row_idx} for each left row. - auto const n = static_cast(left_table_num_rows); - rmm::device_uvector left_keys(n, stream); - thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator(0), - cuda::counting_iterator(left_table_num_rows), - left_keys.begin(), - pair_fn{d_hasher}); - - auto const ref = hash_table.ref(cuco::op::count) - .rebind_key_eq(equality) - .rebind_hash_function(hash_table.hash_function()); + auto count_matches = [&](auto equality, auto hasher) { if (join == join_kind::INNER_JOIN) { - launch_partitioned_count(left_keys.data(), n, match_counts->begin(), ref, stream); + launch_hash_csr_probe_count(left.num_rows(), + valid_rows, + nullptr, + match_counts->data(), + nullptr, + nullptr, + _impl->map_view(), + _impl->csr_view(), + equality, + hasher, + stream); } else { - // IsOuter=true handles the clamp (zero → 1) for LEFT/FULL joins internally. - launch_partitioned_count(left_keys.data(), n, match_counts->begin(), ref, stream); + launch_hash_csr_probe_count(left.num_rows(), + valid_rows, + nullptr, + match_counts->data(), + nullptr, + nullptr, + _impl->map_view(), + _impl->csr_view(), + equality, + hasher, + stream); } }; dispatch_join_comparator( - right, left, preprocessed_right, preprocessed_left, has_nulls, compare_nulls, count_matches); + _right, left, _preprocessed_right, preprocessed_left, _has_nulls, _nulls_equal, count_matches); return match_counts; } -template -std::unique_ptr> hash_join::make_match_counts( - join_kind join, - cudf::table_view const& left, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) const -{ - return make_join_match_counts(_right, - _preprocessed_right, - _impl->_hash_table, - _is_empty, - _has_nulls, - _nulls_equal, - join, - left, - stream, - mr); -} - template std::unique_ptr> hash_join::make_match_counts(join_kind, cudf::table_view const&, diff --git a/cpp/src/join/hash_join/partitioned_count.cu b/cpp/src/join/hash_join/partitioned_count.cu deleted file mode 100644 index 4ddc8a9b5473..000000000000 --- a/cpp/src/join/hash_join/partitioned_count.cu +++ /dev/null @@ -1,20 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "partitioned_count_kernels.cuh" -#include "ref_types.cuh" - -namespace cudf::detail { - -template void launch_partitioned_count( - probe_key_type const*, thread_index_type, size_type*, primitive_count_ref_t, cuda::stream_ref); - -template void launch_partitioned_count( - probe_key_type const*, thread_index_type, size_type*, nested_count_ref_t, cuda::stream_ref); - -template void launch_partitioned_count( - probe_key_type const*, thread_index_type, size_type*, flat_count_ref_t, cuda::stream_ref); - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/partitioned_count_kernels.cuh b/cpp/src/join/hash_join/partitioned_count_kernels.cuh deleted file mode 100644 index e07965ab919a..000000000000 --- a/cpp/src/join/hash_join/partitioned_count_kernels.cuh +++ /dev/null @@ -1,94 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "kernels_common.cuh" - -#include - -#include -#include -#include - -namespace cudf::detail { - -/** - * @brief Count matching build-side rows for each probe key. - * - * Each probing tile (@p cg_size threads) calls `ref.count()` for one probe key - * and reduces the per-lane counts across the tile with a warp reduce. The result - * is written to @p output by a single elected thread via `invoke_one`. If - * @p IsOuter is true, keys with zero matches are recorded as 1 so every probe - * row contributes at least one output row in the subsequent retrieve pass. - * - * This is the first phase of the two-phase partitioned join: count then retrieve. - * The output array is consumed by `launch_partitioned_retrieve` to pre-allocate - * the output index buffers. - * - * @tparam IsOuter If true, zero-match keys produce a count of 1 - * @tparam Ref cuco open-addressing reference type (carries hash, equality, storage) - * @param keys Packed probe keys: `.first` = hash, `.second` = probe row index - * @param n Number of probe keys - * @param output Per-key match count output (one entry per probe key) - * @param ref cuco hash-table reference for counting - */ -template -CUDF_KERNEL void __launch_bounds__(DEFAULT_JOIN_BLOCK_SIZE) - partitioned_count_kernel(probe_key_type const* __restrict__ keys, - thread_index_type n, - size_type* __restrict__ output, - Ref ref) -{ - auto constexpr cg_size = DEFAULT_JOIN_CG_SIZE; - - auto idx = grid_1d::global_thread_id() / cg_size; - auto const stride = grid_1d::grid_stride() / cg_size; - - while (idx < n) { - auto const key = keys[idx]; - if constexpr (cg_size == 1) { - auto const match_count = ref.count(key); - if constexpr (IsOuter) { - output[idx] = (match_count == 0) ? size_type{1} : match_count; - } else { - output[idx] = match_count; - } - } else { - auto const tile = - cooperative_groups::tiled_partition(cooperative_groups::this_thread_block()); - auto const temp_count = static_cast(ref.count(tile, key)); - auto const match_count = - cooperative_groups::reduce(tile, temp_count, cooperative_groups::plus()); - cooperative_groups::invoke_one(tile, [&]() { - if constexpr (IsOuter) { - output[idx] = (match_count == 0) ? size_type{1} : match_count; - } else { - output[idx] = match_count; - } - }); - } - idx += stride; - } -} - -template -void launch_partitioned_count(probe_key_type const* keys, - thread_index_type n, - size_type* output, - Ref ref, - cuda::stream_ref stream) -{ - if (n == 0) { return; } - - auto const config = - grid_1d{static_cast(n * DEFAULT_JOIN_CG_SIZE), DEFAULT_JOIN_BLOCK_SIZE}; - - partitioned_count_kernel - <<>>(keys, n, output, ref); - CUDF_CUDA_TRY(cudaGetLastError()); -} - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/partitioned_count_kernels.hpp b/cpp/src/join/hash_join/partitioned_count_kernels.hpp deleted file mode 100644 index 14e6db4bc294..000000000000 --- a/cpp/src/join/hash_join/partitioned_count_kernels.hpp +++ /dev/null @@ -1,22 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "kernels_common.cuh" - -#include - -namespace cudf::detail { - -/// Launch the partitioned_count kernel. -template -void launch_partitioned_count(probe_key_type const* keys, - thread_index_type n, - size_type* output, - Ref ref, - cuda::stream_ref stream); - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/partitioned_count_outer.cu b/cpp/src/join/hash_join/partitioned_count_outer.cu deleted file mode 100644 index eb0a57632cbd..000000000000 --- a/cpp/src/join/hash_join/partitioned_count_outer.cu +++ /dev/null @@ -1,20 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "partitioned_count_kernels.cuh" -#include "ref_types.cuh" - -namespace cudf::detail { - -template void launch_partitioned_count( - probe_key_type const*, thread_index_type, size_type*, primitive_count_ref_t, cuda::stream_ref); - -template void launch_partitioned_count( - probe_key_type const*, thread_index_type, size_type*, nested_count_ref_t, cuda::stream_ref); - -template void launch_partitioned_count( - probe_key_type const*, thread_index_type, size_type*, flat_count_ref_t, cuda::stream_ref); - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/partitioned_join_retrieve.cu b/cpp/src/join/hash_join/partitioned_join_retrieve.cu index 7c64b9464c55..b8549a675833 100644 --- a/cpp/src/join/hash_join/partitioned_join_retrieve.cu +++ b/cpp/src/join/hash_join/partitioned_join_retrieve.cu @@ -5,53 +5,27 @@ #include "common.cuh" #include "dispatch.cuh" -#include "join/join_common_utils.cuh" +#include "hash_csr_kernels.cuh" #include "join/join_common_utils.hpp" -#include "partitioned_retrieve_kernels.hpp" #include +#include #include +#include +#include +#include #include #include #include +#include #include #include -#include -#include -#include -#include +#include +#include namespace cudf::detail { -namespace { - -/** - * @brief Returns trivial left/right index pairs for an outer join when the build side is empty. - */ -std::pair>, - std::unique_ptr>> -make_trivial_outer_indices(size_type left_start_idx, - size_type partition_size, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - auto left_indices = std::make_unique>(partition_size, stream, mr); - auto right_indices = std::make_unique>(partition_size, stream, mr); - auto out = cuda::zip_iterator(left_indices->begin(), right_indices->begin()); - thrust::tabulate(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - out, - out + partition_size, - cuda::proclaim_return_type>( - [left_start_idx] __device__(auto i) { - return cuda::std::tuple{static_cast(left_start_idx + i), - JoinNoMatch}; - })); - return std::pair(std::move(left_indices), std::move(right_indices)); -} - -} // namespace - template std::pair>, std::unique_ptr>> @@ -95,9 +69,15 @@ hash_join::partitioned_join_retrieve(join_kind join, if (join == join_kind::INNER_JOIN) { return std::pair(std::make_unique>(0, stream, mr), std::make_unique>(0, stream, mr)); - } else { - return make_trivial_outer_indices(left_start_idx, partition_size, stream, mr); } + auto left_indices = + std::make_unique>(partition_size, stream, mr); + auto right_indices = + std::make_unique>(partition_size, stream, mr); + auto const exec = rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()); + thrust::sequence(exec, left_indices->begin(), left_indices->end(), left_start_idx); + thrust::fill(exec, right_indices->begin(), right_indices->end(), JoinNoMatch); + return std::pair(std::move(left_indices), std::move(right_indices)); } // Slice the left table to the partition range @@ -109,49 +89,90 @@ hash_join::partitioned_join_retrieve(join_kind join, auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(left_partition_view, stream); - // For FULL_JOIN, probe with LEFT_JOIN semantics (no complement here) - bool const is_outer = (join != join_kind::INNER_JOIN); - - // launch_partitioned_retrieve reduces match counts to compute output size - // (total = last_offset + last_count), allocates output buffers, and launches the kernel. - auto const* partition_counts = match_ctx._match_counts->data() + left_start_idx; - auto const n = static_cast(partition_size); - - std::pair>, - std::unique_ptr>> - join_indices; - - auto retrieve_partition = [&](auto equality, auto d_hasher) { - // Precompute left keys for this partition slice. - rmm::device_uvector left_keys(n, stream); - thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator(0), - cuda::counting_iterator(partition_size), - left_keys.begin(), - pair_fn{d_hasher}); - - auto const ref = _impl->_hash_table.ref(cuco::op::count) - .rebind_key_eq(equality) - .rebind_hash_function(_impl->_hash_table.hash_function()); - - if (is_outer) { - join_indices = launch_partitioned_retrieve( - left_keys.data(), n, partition_counts, ref, left_start_idx, stream, mr); + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto counts = cudf::detail::make_zeroed_device_uvector_async( + static_cast(partition_size) + 1, stream, temp_mr); + CUDF_CUDA_TRY( + cudf::detail::memcpy_async(counts.data(), + match_ctx._match_counts->data() + left_start_idx, + static_cast(partition_size) * sizeof(size_type), + stream)); + auto offsets = cudf::detail::make_zeroed_device_uvector_async( + static_cast(partition_size) + 1, stream, temp_mr); + auto const output_size = + cudf::detail::sizes_to_offsets(counts.begin(), counts.end(), offsets.begin(), 0, stream); + CUDF_EXPECTS(output_size >= 0, "Join output size overflowed", std::overflow_error); + + rmm::device_uvector probe_slots(partition_size, stream, temp_mr); + auto const row_bitmask = cudf::detail::bitmask_and(left_partition_view, stream, temp_mr).first; + auto const valid_rows = _nulls_equal == null_equality::UNEQUAL + ? static_cast(row_bitmask.data()) + : nullptr; + auto save_slots = [&](auto equality, auto hasher) { + if (join == join_kind::INNER_JOIN) { + launch_hash_csr_probe_count(partition_size, + valid_rows, + probe_slots.data(), + nullptr, + nullptr, + nullptr, + _impl->map_view(), + _impl->csr_view(), + equality, + hasher, + stream); } else { - join_indices = launch_partitioned_retrieve( - left_keys.data(), n, partition_counts, ref, left_start_idx, stream, mr); + launch_hash_csr_probe_count(partition_size, + valid_rows, + probe_slots.data(), + nullptr, + nullptr, + nullptr, + _impl->map_view(), + _impl->csr_view(), + equality, + hasher, + stream); } }; - dispatch_join_comparator(_right, left_partition_view, _preprocessed_right, preprocessed_left, _has_nulls, _nulls_equal, - retrieve_partition); + save_slots); + + auto left_indices = std::make_unique>( + static_cast(output_size), stream, mr); + auto right_indices = std::make_unique>( + static_cast(output_size), stream, mr); + cudf::prefetch::detail::prefetch(*left_indices, stream); + cudf::prefetch::detail::prefetch(*right_indices, stream); + + if (join == join_kind::INNER_JOIN) { + launch_hash_csr_retrieve(output_size, + partition_size, + offsets.data(), + probe_slots.data(), + _impl->csr_view(), + left_start_idx, + left_indices->data(), + right_indices->data(), + stream); + } else { + launch_hash_csr_retrieve(output_size, + partition_size, + offsets.data(), + probe_slots.data(), + _impl->csr_view(), + left_start_idx, + left_indices->data(), + right_indices->data(), + stream); + } - return join_indices; + return {std::move(left_indices), std::move(right_indices)}; } template std::pair>, diff --git a/cpp/src/join/hash_join/partitioned_retrieve.cu b/cpp/src/join/hash_join/partitioned_retrieve.cu deleted file mode 100644 index 4780ce9b9250..000000000000 --- a/cpp/src/join/hash_join/partitioned_retrieve.cu +++ /dev/null @@ -1,41 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "partitioned_retrieve_kernels.cuh" -#include "ref_types.cuh" - -namespace cudf::detail { - -template std::pair>, - std::unique_ptr>> -launch_partitioned_retrieve(probe_key_type const*, - thread_index_type, - size_type const*, - primitive_count_ref_t, - size_type, - cuda::stream_ref, - rmm::device_async_resource_ref); - -template std::pair>, - std::unique_ptr>> -launch_partitioned_retrieve(probe_key_type const*, - thread_index_type, - size_type const*, - nested_count_ref_t, - size_type, - cuda::stream_ref, - rmm::device_async_resource_ref); - -template std::pair>, - std::unique_ptr>> -launch_partitioned_retrieve(probe_key_type const*, - thread_index_type, - size_type const*, - flat_count_ref_t, - size_type, - cuda::stream_ref, - rmm::device_async_resource_ref); - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/partitioned_retrieve_kernels.cuh b/cpp/src/join/hash_join/partitioned_retrieve_kernels.cuh deleted file mode 100644 index 72518018fb2f..000000000000 --- a/cpp/src/join/hash_join/partitioned_retrieve_kernels.cuh +++ /dev/null @@ -1,258 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "kernels_common.cuh" - -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace cudf::detail { - -/** - * @brief Count the number of set bits below a given position in a bitmask. - */ -__device__ __forceinline__ int count_lower_set_bits(unsigned int mask, int pos) -{ - return cuda::std::popcount(mask & ((1u << pos) - 1)); -} - -/** - * @brief Retrieve matching build-side rows for each probe key. - * - * Each probing tile (@p cg_size threads) walks the hash table for one probe key, - * collecting matches via warp ballot. Matches are staged in a per-flushing-tile (warp) - * shared-memory buffer instead of being written directly to global memory. When the buffer - * nears capacity, the flushing tile claims a contiguous range in the global output arrays - * via a single atomic and flushes with coalesced writes, amortising atomic overhead across - * many matches. If @p IsOuter is true, probe rows with no matches emit a - * `(left_index, JoinNoMatch)` pair. - * - * @tparam IsOuter If true, unmatched probe rows emit a null-padded output row - * @tparam Ref cuco open-addressing reference type (carries hash, equality, storage) - * @param keys Packed left keys: `.first` = hash, `.second` = left row index - * @param n Number of probe keys - * @param left_offset Added to each probe row index to produce an absolute left index - * @param left_output Output buffer for left (probe-side) row indices - * @param right_output Output buffer for right (build-side) row indices - * @param output_counter Global atomic counter tracking total pairs written so far - * @param ref cuco hash-table reference for probing - */ -template -CUDF_KERNEL void __launch_bounds__(DEFAULT_JOIN_BLOCK_SIZE) - partitioned_retrieve_kernel(probe_key_type const* __restrict__ keys, - thread_index_type n, - size_type left_offset, - size_type* __restrict__ left_output, - size_type* __restrict__ right_output, - size_type* __restrict__ output_counter, - Ref ref) -{ - namespace cg = cooperative_groups; - - auto constexpr cg_size = Ref::cg_size; - auto constexpr bucket_size = Ref::bucket_size; - auto constexpr flushing_tile_size = 32; // full warp for coalesced flushes - static_assert(flushing_tile_size >= cg_size); - static_assert(flushing_tile_size % cg_size == 0, - "Every probing tile must sit inside a single flushing tile"); - static_assert(DEFAULT_JOIN_BLOCK_SIZE % flushing_tile_size == 0); - - auto constexpr num_flushing_tiles = DEFAULT_JOIN_BLOCK_SIZE / flushing_tile_size; - auto constexpr tiles_in_block = DEFAULT_JOIN_BLOCK_SIZE / cg_size; - auto constexpr max_matches_per_step = flushing_tile_size * bucket_size; - // buffer_size leaves headroom so one full probing step can't overflow. - auto constexpr buffer_size = max_matches_per_step + flushing_tile_size; - - using index_pair = cuco::pair; - __shared__ index_pair buffers[num_flushing_tiles][buffer_size]; - __shared__ cuda::std::int32_t counters[num_flushing_tiles]; - - auto const block = cg::this_thread_block(); - auto const flushing_tile = cg::tiled_partition(block); - auto const probing_tile = cg::tiled_partition(block); - auto const flushing_tile_id = flushing_tile.meta_group_rank(); - auto const empty_sentinel = ref.empty_key_sentinel(); - auto const key_equal = ref.key_eq(); - - if (flushing_tile.thread_rank() == 0) { counters[flushing_tile_id] = 0; } - flushing_tile.sync(); - - auto atomic_counter = cuda::atomic_ref{*output_counter}; - - auto flush_buffers = [&](auto const& tile) { - auto const count = counters[flushing_tile_id]; - auto const offset = cg::invoke_one_broadcast(tile, [&]() { - return atomic_counter.fetch_add(static_cast(count), cuda::memory_order_relaxed); - }); - auto const rank = tile.thread_rank(); - for (int i = rank; i < count; i += tile.size()) { - left_output[offset + i] = buffers[flushing_tile_id][i].first; - right_output[offset + i] = buffers[flushing_tile_id][i].second; - } - }; - - auto const grid_stride_tiles = static_cast(gridDim.x) * tiles_in_block; - auto idx = - static_cast(blockIdx.x) * tiles_in_block + probing_tile.meta_group_rank(); - - while (flushing_tile.any(idx < n)) { - bool const active = idx < n; - auto const active_flushing_tile = - cg::binary_partition(flushing_tile, active); - - if (active) { - auto const probe_key = keys[idx]; - auto const left_index = probe_key.second + left_offset; - - auto probing_iter = ref.probing_scheme().template make_iterator( - probing_tile, probe_key, ref.storage_ref().extent()); - auto const init_probing_idx = *probing_iter; - - bool running = true; - [[maybe_unused]] bool found_match = false; - - while (active_flushing_tile.any(running)) { - if (running) { - auto const bucket_slots = ref.storage_ref()[*probing_iter]; - - bool equals[bucket_size]; - for (int i = 0; i < bucket_size; ++i) { - equals[i] = false; - if (running) { - if (bucket_slots[i] == empty_sentinel) { - running = false; - } else if (key_equal(probe_key, bucket_slots[i])) { - equals[i] = true; - } - } - } - - probing_tile.sync(); - running = probing_tile.all(running); - - cuda::std::int32_t exists[bucket_size]; - cuda::std::int32_t num_matches[bucket_size]; - cuda::std::int32_t total_matches = 0; - for (int i = 0; i < bucket_size; ++i) { - exists[i] = probing_tile.ballot(equals[i]); - num_matches[i] = cuda::std::popcount(static_cast(exists[i])); - total_matches += num_matches[i]; - } - - auto const lane_id = probing_tile.thread_rank(); - - if (total_matches > 0) { - if constexpr (IsOuter) { found_match = true; } - - cuda::std::int32_t output_idx = 0; - if (lane_id == 0) { - auto shared_ref = cuda::atomic_ref{ - counters[flushing_tile_id]}; - output_idx = shared_ref.fetch_add(total_matches, cuda::memory_order_relaxed); - } - output_idx = probing_tile.shfl(output_idx, 0); - - cuda::std::int32_t matches_offset = 0; - for (int i = 0; i < bucket_size; ++i) { - if (equals[i]) { - auto const lane_offset = count_lower_set_bits(exists[i], lane_id); - buffers[flushing_tile_id][output_idx + matches_offset + lane_offset] = { - left_index, bucket_slots[i].second}; - } - matches_offset += num_matches[i]; - } - } - - if constexpr (IsOuter) { - if (!running && !found_match && lane_id == 0) { - auto shared_ref = cuda::atomic_ref{ - counters[flushing_tile_id]}; - auto const output_idx = shared_ref.fetch_add(1, cuda::memory_order_relaxed); - buffers[flushing_tile_id][output_idx] = {left_index, cudf::JoinNoMatch}; - } - } - } // if running - - active_flushing_tile.sync(); - if (counters[flushing_tile_id] > (buffer_size - max_matches_per_step)) { - flush_buffers(active_flushing_tile); - active_flushing_tile.sync(); - if (active_flushing_tile.thread_rank() == 0) { counters[flushing_tile_id] = 0; } - active_flushing_tile.sync(); - } - - ++probing_iter; - if (*probing_iter == init_probing_idx) { running = false; } - } // while running - } // if active - - idx += grid_stride_tiles; - } // while idx < n - - flushing_tile.sync(); - if (counters[flushing_tile_id] > 0) { flush_buffers(flushing_tile); } -} - -template -std::pair>, - std::unique_ptr>> -launch_partitioned_retrieve(probe_key_type const* keys, - thread_index_type n, - size_type const* match_counts, - Ref ref, - size_type left_offset, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - if (n == 0) { - return std::pair(std::make_unique>(0, stream, mr), - std::make_unique>(0, stream, mr)); - } - - // Shared-memory buffered retrieve only needs the total output size, not - // per-row offsets. A reduce is cheaper than an exclusive_scan. - auto const total_output = - thrust::reduce(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - match_counts, - match_counts + n, - size_type{0}); - - if (total_output == 0) { - return std::pair(std::make_unique>(0, stream, mr), - std::make_unique>(0, stream, mr)); - } - - auto left_indices = std::make_unique>(total_output, stream, mr); - auto right_indices = std::make_unique>(total_output, stream, mr); - - // Global atomic counter claimed in bulk by each flushing-tile buffer flush. - cudf::detail::device_scalar output_counter(size_type{0}, stream); - - auto constexpr tiles_in_block = DEFAULT_JOIN_BLOCK_SIZE / Ref::cg_size; - auto const num_blocks = static_cast((n + tiles_in_block - 1) / tiles_in_block); - - partitioned_retrieve_kernel<<>>( - keys, n, left_offset, left_indices->data(), right_indices->data(), output_counter.data(), ref); - CUDF_CUDA_TRY(cudaGetLastError()); - - return std::pair(std::move(left_indices), std::move(right_indices)); -} - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/partitioned_retrieve_kernels.hpp b/cpp/src/join/hash_join/partitioned_retrieve_kernels.hpp deleted file mode 100644 index aa7f89fc066f..000000000000 --- a/cpp/src/join/hash_join/partitioned_retrieve_kernels.hpp +++ /dev/null @@ -1,42 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "kernels_common.cuh" - -#include - -#include - -#include - -#include -#include - -namespace cudf::detail { - -/** - * @brief Probes the hash table for each key and writes matching index pairs. - * - * Reduces match_counts to derive the total output size, allocates output buffers, - * and launches the retrieve kernel. `left_offset` is added to each stored probe-row index when - * writing to `left_indices`, so callers can produce indices in the full probe - * table's coordinate space directly from a slice-local `keys` array. - * - * @return A pair of device vectors [left_indices, right_indices]. - */ -template -std::pair>, - std::unique_ptr>> -launch_partitioned_retrieve(probe_key_type const* keys, - thread_index_type n, - size_type const* match_counts, - Ref ref, - size_type left_offset, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/partitioned_retrieve_outer.cu b/cpp/src/join/hash_join/partitioned_retrieve_outer.cu deleted file mode 100644 index afb2acb16bca..000000000000 --- a/cpp/src/join/hash_join/partitioned_retrieve_outer.cu +++ /dev/null @@ -1,41 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "partitioned_retrieve_kernels.cuh" -#include "ref_types.cuh" - -namespace cudf::detail { - -template std::pair>, - std::unique_ptr>> -launch_partitioned_retrieve(probe_key_type const*, - thread_index_type, - size_type const*, - primitive_count_ref_t, - size_type, - cuda::stream_ref, - rmm::device_async_resource_ref); - -template std::pair>, - std::unique_ptr>> -launch_partitioned_retrieve(probe_key_type const*, - thread_index_type, - size_type const*, - nested_count_ref_t, - size_type, - cuda::stream_ref, - rmm::device_async_resource_ref); - -template std::pair>, - std::unique_ptr>> -launch_partitioned_retrieve(probe_key_type const*, - thread_index_type, - size_type const*, - flat_count_ref_t, - size_type, - cuda::stream_ref, - rmm::device_async_resource_ref); - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/ref_types.cuh b/cpp/src/join/hash_join/ref_types.cuh deleted file mode 100644 index a1511c56dcb1..000000000000 --- a/cpp/src/join/hash_join/ref_types.cuh +++ /dev/null @@ -1,44 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -// Type aliases for the cuco hash table ref types and equality comparators -// used across hash join probe kernels. There are 3 dispatch paths: -// primitive, nested, non-nested. - -#pragma once - -#include "dispatch.cuh" -#include "hash_join_impl.cuh" - -#include - -#include - -namespace cudf::detail { - -// --- Equality types from the 3 dispatch paths --- - -using primitive_equality_t = primitive_pair_equal; - -using nested_equality_t = pair_equal>>; - -using flat_equality_t = pair_equal>>; - -// --- Count ref types (used by partitioned_count kernel) --- - -template -using count_ref_t = - decltype(std::declval() - .ref(cuco::op::count) - .rebind_key_eq(std::declval()) - .rebind_hash_function(std::declval().hash_function())); - -using primitive_count_ref_t = count_ref_t; -using nested_count_ref_t = count_ref_t; -using flat_count_ref_t = count_ref_t; - -} // namespace cudf::detail diff --git a/cpp/src/join/hash_join/retrieve_impl.cuh b/cpp/src/join/hash_join/retrieve_impl.cuh index e19b8e8e0c40..2440f31509dc 100644 --- a/cpp/src/join/hash_join/retrieve_impl.cuh +++ b/cpp/src/join/hash_join/retrieve_impl.cuh @@ -6,146 +6,22 @@ #include "common.cuh" #include "dispatch.cuh" -#include "join/join_common_utils.cuh" +#include "hash_csr_kernels.cuh" #include "join/join_common_utils.hpp" -#include "size_impl.cuh" -#include +#include #include +#include +#include #include #include #include #include #include -#include - -#include -#include namespace cudf::detail { -template -std::pair>, - std::unique_ptr>> -probe_join_hash_table( - cudf::table_view const& right_table, - cudf::table_view const& left_table, - std::shared_ptr const& preprocessed_right, - std::shared_ptr const& preprocessed_left, - cudf::detail::hash_table_t const& hash_table, - bool has_nulls, - null_equality compare_nulls, - std::optional output_size, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - static_assert(Join == join_kind::INNER_JOIN || Join == join_kind::LEFT_JOIN || - Join == join_kind::FULL_JOIN); - - constexpr auto size_join = Join == join_kind::FULL_JOIN ? join_kind::LEFT_JOIN : Join; - - std::size_t const join_size = output_size - ? *output_size - : compute_join_output_size(right_table, - left_table, - preprocessed_right, - preprocessed_left, - hash_table, - has_nulls, - compare_nulls, - stream); - - if (join_size == 0) { - return std::pair(std::make_unique>(0, stream, mr), - std::make_unique>(0, stream, mr)); - } - - auto left_indices = std::make_unique>(join_size, stream, mr); - auto right_indices = std::make_unique>(join_size, stream, mr); - cudf::prefetch::detail::prefetch(*left_indices, stream); - cudf::prefetch::detail::prefetch(*right_indices, stream); - - auto const left_table_num_rows = left_table.num_rows(); - auto const out_probe_begin = - cuda::make_transform_output_iterator(left_indices->begin(), output_fn{}); - auto const out_build_begin = - cuda::make_transform_output_iterator(right_indices->begin(), output_fn{}); - - auto retrieve_results = [&](auto equality, auto d_hasher) { - auto const iter = cudf::detail::make_counting_transform_iterator(0, pair_fn{d_hasher}); - if constexpr (Join == join_kind::INNER_JOIN) { - hash_table.retrieve(iter, - iter + left_table_num_rows, - equality, - hash_table.hash_function(), - out_probe_begin, - out_build_begin, - stream.get()); - } else { - [[maybe_unused]] auto out_probe_end = hash_table - .retrieve_outer(iter, - iter + left_table_num_rows, - equality, - hash_table.hash_function(), - out_probe_begin, - out_build_begin, - stream.get()) - .first; - - if constexpr (Join == join_kind::FULL_JOIN) { - auto const actual_size = cuda::std::distance(out_probe_begin, out_probe_end); - left_indices->resize(actual_size, stream); - right_indices->resize(actual_size, stream); - } - } - }; - - dispatch_join_comparator(right_table, - left_table, - preprocessed_right, - preprocessed_left, - has_nulls, - compare_nulls, - retrieve_results); - - return std::pair(std::move(left_indices), std::move(right_indices)); -} - -template -void retrieve_left_join_build_indices( - cudf::table_view const& right_table, - cudf::table_view const& left_table, - std::shared_ptr const& preprocessed_right, - std::shared_ptr const& preprocessed_left, - cudf::detail::hash_table_t const& hash_table, - bool has_nulls, - null_equality compare_nulls, - RightOutputIterator out_build_begin, - cuda::stream_ref stream) -{ - auto const left_table_num_rows = left_table.num_rows(); - - auto retrieve_results = [&](auto equality, auto d_hasher) { - auto const iter = cudf::detail::make_counting_transform_iterator(0, pair_fn{d_hasher}); - hash_table.retrieve_outer(iter, - iter + left_table_num_rows, - equality, - hash_table.hash_function(), - cuda::make_discard_iterator(), - out_build_begin, - stream.get()); - }; - - dispatch_join_comparator(right_table, - left_table, - preprocessed_right, - preprocessed_left, - has_nulls, - compare_nulls, - retrieve_results); -} - template template std::pair>, @@ -176,16 +52,70 @@ hash_join::join_retrieve(cudf::table_view const& left, auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(left, stream); - auto join_indices = cudf::detail::probe_join_hash_table(_right, - left, - _preprocessed_right, - preprocessed_left, - _impl->_hash_table, - _has_nulls, - _nulls_equal, - output_size, - stream, - mr); + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto match_counts = cudf::detail::make_zeroed_device_uvector_async( + static_cast(left.num_rows()) + 1, stream, temp_mr); + rmm::device_uvector probe_slots(left.num_rows(), stream, temp_mr); + auto const row_bitmask = cudf::detail::bitmask_and(left, stream, temp_mr).first; + auto const valid_rows = _nulls_equal == null_equality::UNEQUAL + ? static_cast(row_bitmask.data()) + : nullptr; + + auto count_matches = [&](auto equality, auto hasher) { + launch_hash_csr_probe_count(left.num_rows(), + valid_rows, + probe_slots.data(), + match_counts.data(), + nullptr, + nullptr, + _impl->map_view(), + _impl->csr_view(), + equality, + hasher, + stream); + }; + dispatch_join_comparator( + _right, left, _preprocessed_right, preprocessed_left, _has_nulls, _nulls_equal, count_matches); + + auto offsets = cudf::detail::make_zeroed_device_uvector_async( + static_cast(left.num_rows()) + 1, stream, temp_mr); + auto const actual_size = cudf::detail::sizes_to_offsets( + match_counts.begin(), match_counts.end(), offsets.begin(), 0, stream); + CUDF_EXPECTS(actual_size >= 0, "Join output size overflowed", std::overflow_error); + auto const join_size = Join != join_kind::FULL_JOIN && output_size.has_value() + ? *output_size + : static_cast(actual_size); + CUDF_EXPECTS(join_size == static_cast(actual_size), + "The provided join output size is incorrect"); + + auto left_indices = std::make_unique>(join_size, stream, mr); + auto right_indices = std::make_unique>(join_size, stream, mr); + cudf::prefetch::detail::prefetch(*left_indices, stream); + cudf::prefetch::detail::prefetch(*right_indices, stream); + + if constexpr (Join == join_kind::INNER_JOIN) { + launch_hash_csr_retrieve(actual_size, + left.num_rows(), + offsets.data(), + probe_slots.data(), + _impl->csr_view(), + 0, + left_indices->data(), + right_indices->data(), + stream); + } else { + launch_hash_csr_retrieve(actual_size, + left.num_rows(), + offsets.data(), + probe_slots.data(), + _impl->csr_view(), + 0, + left_indices->data(), + right_indices->data(), + stream); + } + + auto join_indices = std::pair(std::move(left_indices), std::move(right_indices)); if constexpr (Join == join_kind::FULL_JOIN) { return detail::finalize_full_join( diff --git a/cpp/src/join/hash_join/size_impl.cuh b/cpp/src/join/hash_join/size_impl.cuh index 3b1c44672fcc..88d2dd1a1e95 100644 --- a/cpp/src/join/hash_join/size_impl.cuh +++ b/cpp/src/join/hash_join/size_impl.cuh @@ -6,60 +6,16 @@ #include "common.cuh" #include "dispatch.cuh" -#include "join/join_common_utils.cuh" +#include "hash_csr_kernels.cuh" +#include +#include #include +#include -namespace cudf::detail { - -std::size_t get_full_join_size( - cudf::table_view const& right_table, - cudf::table_view const& left_table, - std::shared_ptr const& preprocessed_right, - std::shared_ptr const& preprocessed_left, - cudf::detail::hash_table_t const& hash_table, - bool has_nulls, - null_equality compare_nulls, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -template -std::size_t compute_join_output_size( - table_view const& right_table, - table_view const& left_table, - std::shared_ptr const& preprocessed_right, - std::shared_ptr const& preprocessed_left, - cudf::detail::hash_table_t const& hash_table, - bool has_nulls, - cudf::null_equality nulls_equal, - cuda::stream_ref stream) -{ - static_assert(Join == join_kind::INNER_JOIN || Join == join_kind::LEFT_JOIN); +#include - if (right_table.num_rows() == 0) { - return Join == join_kind::INNER_JOIN ? 0 : left_table.num_rows(); - } - - auto const left_table_num_rows = left_table.num_rows(); - - return dispatch_join_comparator( - right_table, - left_table, - preprocessed_right, - preprocessed_left, - has_nulls, - nulls_equal, - [&](auto equality, auto d_hasher) { - auto const iter = cudf::detail::make_counting_transform_iterator(0, pair_fn{d_hasher}); - if constexpr (Join == join_kind::LEFT_JOIN) { - return hash_table.count_outer( - iter, iter + left_table_num_rows, equality, hash_table.hash_function(), stream.get()); - } else { - return hash_table.count( - iter, iter + left_table_num_rows, equality, hash_table.hash_function(), stream.get()); - } - }); -} +namespace cudf::detail { template template @@ -83,14 +39,33 @@ std::size_t hash_join::join_size(cudf::table_view const& left, auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(left, stream); - return cudf::detail::compute_join_output_size(_right, - left, - _preprocessed_right, - preprocessed_left, - _impl->_hash_table, - _has_nulls, - _nulls_equal, - stream); + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto match_counts = + cudf::detail::make_zeroed_device_uvector_async(left.num_rows(), stream, temp_mr); + auto const row_bitmask = cudf::detail::bitmask_and(left, stream, temp_mr).first; + auto const valid_rows = _nulls_equal == null_equality::UNEQUAL + ? static_cast(row_bitmask.data()) + : nullptr; + + auto count_matches = [&](auto equality, auto hasher) { + launch_hash_csr_probe_count(left.num_rows(), + valid_rows, + nullptr, + match_counts.data(), + nullptr, + nullptr, + _impl->map_view(), + _impl->csr_view(), + equality, + hasher, + stream); + }; + dispatch_join_comparator( + _right, left, _preprocessed_right, preprocessed_left, _has_nulls, _nulls_equal, count_matches); + auto const output_size = cudf::detail::reduce( + match_counts.begin(), match_counts.end(), std::int64_t{0}, cuda::std::plus<>{}, stream); + CUDF_EXPECTS(output_size >= 0, "Join output size overflowed", std::overflow_error); + return static_cast(output_size); } template @@ -102,7 +77,6 @@ std::size_t hash_join::join_size(cudf::table_view const& left, static_assert(Join == join_kind::FULL_JOIN); CUDF_FUNC_RANGE(); - if (_is_empty) { return left.num_rows(); } CUDF_EXPECTS(_has_nulls || !cudf::has_nested_nulls(left), @@ -111,16 +85,42 @@ std::size_t hash_join::join_size(cudf::table_view const& left, auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(left, stream); - - return cudf::detail::get_full_join_size(_right, - left, - _preprocessed_right, - preprocessed_left, - _impl->_hash_table, - _has_nulls, - _nulls_equal, - stream, - mr); + auto match_counts = + cudf::detail::make_zeroed_device_uvector_async(left.num_rows(), stream, mr); + auto matched_slots = + cudf::detail::make_zeroed_device_uvector_async(_impl->capacity, stream, mr); + auto matched_build_rows = cudf::detail::device_scalar(0, stream, mr); + auto const row_bitmask = cudf::detail::bitmask_and(left, stream, mr).first; + auto const valid_rows = _nulls_equal == null_equality::UNEQUAL + ? static_cast(row_bitmask.data()) + : nullptr; + + auto count_matches = [&](auto equality, auto hasher) { + launch_hash_csr_probe_count(left.num_rows(), + valid_rows, + nullptr, + match_counts.data(), + matched_slots.data(), + matched_build_rows.data(), + _impl->map_view(), + _impl->csr_view(), + equality, + hasher, + stream); + }; + dispatch_join_comparator( + _right, left, _preprocessed_right, preprocessed_left, _has_nulls, _nulls_equal, count_matches); + + auto const left_output_size = cudf::detail::reduce( + match_counts.begin(), match_counts.end(), std::int64_t{0}, cuda::std::plus<>{}, stream); + auto const matched_right_rows = matched_build_rows.value(stream); + CUDF_EXPECTS(left_output_size >= 0, "Join output size overflowed", std::overflow_error); + auto const output_size = static_cast(left_output_size) + + static_cast(_right.num_rows()) - matched_right_rows; + CUDF_EXPECTS(output_size <= std::numeric_limits::max(), + "Join output size overflowed", + std::overflow_error); + return static_cast(output_size); } } // namespace cudf::detail