diff --git a/cpp/src/join/hash_join/hash_csr_kernels.cuh b/cpp/src/join/hash_join/hash_csr_kernels.cuh index ae577f547c32..1052605a3000 100644 --- a/cpp/src/join/hash_join/hash_csr_kernels.cuh +++ b/cpp/src/join/hash_join/hash_csr_kernels.cuh @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -127,7 +128,7 @@ CUDF_KERNEL void hash_csr_build_fill_kernel(size_type num_rows, } } -template +template CUDF_KERNEL void hash_csr_probe_count_kernel(size_type num_rows, bitmask_type const* valid_rows, size_type* probe_groups, @@ -156,7 +157,7 @@ CUDF_KERNEL void hash_csr_probe_count_kernel(size_type num_rows, } auto const match_count = IsOuter ? cuda::std::max(count, size_type{1}) : count; if (match_counts != nullptr) { match_counts[index] = match_count; } - thread_matches += static_cast(match_count); + if constexpr (ComputeSum) { thread_matches += static_cast(match_count); } // Only right and full joins consume the matched-row tally, and `matched_groups` is null for // every other kind, so this whole block compiles away outside outer joins rather than costing @@ -175,7 +176,7 @@ CUDF_KERNEL void hash_csr_probe_count_kernel(size_type num_rows, } } - if (total_matches != nullptr) { + if constexpr (ComputeSum) { using block_reduce = cub::BlockReduce; __shared__ typename block_reduce::TempStorage temp_storage; auto const block_matches = block_reduce(temp_storage).Sum(thread_matches); @@ -251,18 +252,22 @@ void launch_hash_csr_probe_count_kernel(size_type num_rows, { 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_groups, - match_counts, - total_matches, - matched_groups, - matched_build_rows, - map, - csr, - equal, - hasher); + // Keep reduction registers and shared memory out of the kernels used by match contexts and + // retrieval. A device-side null check still reserves those resources for every probe. + dispatch_bool(total_matches != nullptr, [&](auto compute_sum) { + hash_csr_probe_count_kernel + <<>>(num_rows, + valid_rows, + probe_groups, + match_counts, + total_matches, + matched_groups, + matched_build_rows, + map, + csr, + equal, + hasher); + }); CUDF_CUDA_TRY(cudaGetLastError()); }