From f0573ef5c09673dcd40d7c357c4e7c6eae40a375 Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Fri, 28 Aug 2026 22:00:20 +0000 Subject: [PATCH] Support concurrent aggregate calls on one streaming_groupby --- cpp/include/cudf/groupby.hpp | 11 +++ .../groupby/streaming_groupby/aggregate.cu | 29 ++++--- cpp/src/groupby/streaming_groupby/common.cuh | 45 ++++++++++- cpp/src/groupby/streaming_groupby/impl.cu | 9 ++- cpp/src/groupby/streaming_groupby/insert.cuh | 9 ++- cpp/src/groupby/streaming_groupby/merge.cu | 13 +++- cpp/tests/groupby/streaming_groupby_test.cpp | 77 +++++++++++++++++++ 7 files changed, 173 insertions(+), 20 deletions(-) diff --git a/cpp/include/cudf/groupby.hpp b/cpp/include/cudf/groupby.hpp index a3e21b58093f..da98e4e97142 100644 --- a/cpp/include/cudf/groupby.hpp +++ b/cpp/include/cudf/groupby.hpp @@ -519,6 +519,13 @@ class streaming_groupby { * are updated atomically. The input `data` table is not referenced after this * call returns. * + * This function may be called concurrently from multiple host threads on the same object, + * and each call may supply a different stream. Callers do not need to serialize the calls or + * synchronize between them. Key insertion is serialized internally, on the host and across + * streams, because a batch's newly discovered keys are held in a transient encoding that is + * only valid while that one insertion is in flight. The aggregation that follows each + * insertion updates every group atomically, so those phases overlap freely across streams. + * * @param data Table containing both key and value columns * @param stream CUDA stream used for device memory operations and kernel launches * @@ -532,6 +539,10 @@ class streaming_groupby { * * Extracts the other object's accumulated intermediate state and merges it into this * object's persistent hash table. The other object is not modified. + * + * This function shares the insertion path with `aggregate()` and is serialized against it, so + * it is safe to call while other host threads are calling `aggregate()` on this object. The + * source object must not be mutated concurrently. * Both objects must have been constructed with compatible aggregation requests, * and this object must have had at least one `aggregate()` call. * diff --git a/cpp/src/groupby/streaming_groupby/aggregate.cu b/cpp/src/groupby/streaming_groupby/aggregate.cu index c024565720d1..3ae891f729db 100644 --- a/cpp/src/groupby/streaming_groupby/aggregate.cu +++ b/cpp/src/groupby/streaming_groupby/aggregate.cu @@ -17,16 +17,13 @@ #include #include +#include #include namespace cudf::groupby { void streaming_groupby::impl::do_aggregate(table_view const& data, cuda::stream_ref stream) { - CUDF_EXPECTS(!_invalidated, - "streaming_groupby is in an invalidated state from a prior failure; " - "no further aggregate()/merge() is allowed. finalize() may still be called."); - auto const batch_size = data.num_rows(); if (batch_size == 0) { return; } @@ -40,15 +37,29 @@ void streaming_groupby::impl::do_aggregate(table_view const& data, cuda::stream_ "Transient key encoding (max_distinct_keys + batch_size) would overflow size_type.", std::invalid_argument); - if (!_initialized) { initialize(data, stream); } + // The transient key encoding is only valid while a single insertion is in flight, so + // insertion is serialized across concurrent callers on the host and, via the event, on the + // device. The aggregation below is per-group atomic and runs unserialized. + auto const result = [&] { + std::lock_guard const lock{_insert_mutex}; + + CUDF_EXPECTS(!_invalidated, + "streaming_groupby is in an invalidated state from a prior failure; " + "no further aggregate()/merge() is allowed. finalize() may still be called."); + + if (!_initialized) { initialize(data, stream); } - auto const batch_keys = data.select(_key_indices); + auto const batch_keys = data.select(_key_indices); - update_nullable_state(batch_keys); + update_nullable_state(batch_keys); - if (!_key_set) { create_key_set(stream); } + if (!_key_set) { create_key_set(stream); } - auto result = probe_and_insert(batch_keys, stream); + _insert_done.wait(stream); + auto inserted = probe_and_insert(batch_keys, stream); + _insert_done.record(stream); + return inserted; + }(); auto const values_view = data.select(_value_col_indices); auto const d_values = table_device_view::create(values_view, stream); diff --git a/cpp/src/groupby/streaming_groupby/common.cuh b/cpp/src/groupby/streaming_groupby/common.cuh index 100837863792..255cb41c81d5 100644 --- a/cpp/src/groupby/streaming_groupby/common.cuh +++ b/cpp/src/groupby/streaming_groupby/common.cuh @@ -26,11 +26,40 @@ #include #include +#include #include +#include #include namespace cudf::groupby { +/* + * Minimal owning wrapper around a CUDA event, used to order the insertion phase of + * `aggregate()` / `merge()` calls that overlap on different streams. Only the insertion + * phase needs this ordering; the aggregation phase updates each group with atomics and is + * safe to overlap. + */ +class insert_order_event { + public: + insert_order_event() { CUDF_CUDA_TRY(cudaEventCreateWithFlags(&_event, cudaEventDisableTiming)); } + ~insert_order_event() { cudaEventDestroy(_event); } + insert_order_event(insert_order_event const&) = delete; + insert_order_event& operator=(insert_order_event const&) = delete; + + /// Makes `stream` wait for the most recently recorded insertion. No-op before the first + /// `record()`, which is exactly the behavior the first call needs. + void wait(cuda::stream_ref stream) const + { + CUDF_CUDA_TRY(cudaStreamWaitEvent(stream.get(), _event)); + } + + /// Records completion of the insertion just enqueued on `stream`. + void record(cuda::stream_ref stream) { CUDF_CUDA_TRY(cudaEventRecord(_event, stream.get())); } + + private: + cudaEvent_t _event{}; +}; + /* * Companion location for a stored dense ID: which compacted batch table the key * lives in (`first`) and the row index within that table (`second`). Packed into @@ -250,6 +279,15 @@ struct streaming_groupby::impl { null_policy _null_handling; cuda::mr::any_resource _mr; + /* + * Serializes the insertion phase of `aggregate()` and `merge()`. Callers may invoke those + * from multiple host threads; everything they mutate on the host, and the transient key + * encoding they place in the hash set, is guarded here. + */ + std::mutex _insert_mutex; + /// Orders the insertion phase across calls that supply different streams. + insert_order_event _insert_done; + bool _initialized{false}; /// Set true once an `aggregate()` / `merge()` call has thrown after touching the /// hash set. Subsequent `aggregate()` / `merge()` calls fail fast; only @@ -260,7 +298,7 @@ struct streaming_groupby::impl { * mark of dense IDs in the persistent hash set: stored slot values are in * [0, _distinct_keys). */ - size_type _distinct_keys{0}; + std::atomic _distinct_keys{0}; bool _has_nullable_keys{false}; bool _has_nested_keys{false}; @@ -301,7 +339,10 @@ struct streaming_groupby::impl { std::unique_ptr _key_set; [[nodiscard]] size_type num_keys() const { return static_cast(_key_indices.size()); } - [[nodiscard]] bool has_state() const { return _initialized && _distinct_keys > 0; } + [[nodiscard]] bool has_state() const + { + return _initialized && _distinct_keys.load(std::memory_order_relaxed) > 0; + } impl(host_span key_indices, host_span requests, diff --git a/cpp/src/groupby/streaming_groupby/impl.cu b/cpp/src/groupby/streaming_groupby/impl.cu index de8ddcbb3254..19b94b7ce132 100644 --- a/cpp/src/groupby/streaming_groupby/impl.cu +++ b/cpp/src/groupby/streaming_groupby/impl.cu @@ -243,7 +243,9 @@ std::unique_ptr streaming_groupby::impl::gather_agg_results( // The results we care about are dense in `[0, _distinct_keys)` and can be extracted by // slice+copy. auto const sliced = - cudf::detail::slice(_agg_results->view(), {0, _distinct_keys}, stream).front(); + cudf::detail::slice( + _agg_results->view(), {0, _distinct_keys.load(std::memory_order_relaxed)}, stream) + .front(); return std::make_unique
(sliced, stream, mr); } @@ -378,7 +380,10 @@ std::pair, std::vector> streaming_gro return _impl->do_finalize(stream, mr); } -size_type streaming_groupby::distinct_keys() const noexcept { return _impl->_distinct_keys; } +size_type streaming_groupby::distinct_keys() const noexcept +{ + return _impl->_distinct_keys.load(std::memory_order_relaxed); +} bool is_streaming_groupby_supported(data_type values_type, aggregation::Kind kind) { diff --git a/cpp/src/groupby/streaming_groupby/insert.cuh b/cpp/src/groupby/streaming_groupby/insert.cuh index 9df90ecfcb93..d2c51acd5667 100644 --- a/cpp/src/groupby/streaming_groupby/insert.cuh +++ b/cpp/src/groupby/streaming_groupby/insert.cuh @@ -96,9 +96,10 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_ // Bound check: the hash set has already been written above (transient slot values), // so on failure the object is left invalidated; further aggregate()/merge() calls // will throw immediately while finalize() can still recover partial results. - if (_distinct_keys + new_distinct_keys > _max_distinct_keys) { + auto const distinct_so_far = _distinct_keys.load(std::memory_order_relaxed); + if (distinct_so_far + new_distinct_keys > _max_distinct_keys) { _invalidated = true; - CUDF_FAIL("Distinct key count (" + std::to_string(_distinct_keys + new_distinct_keys) + + CUDF_FAIL("Distinct key count (" + std::to_string(distinct_so_far + new_distinct_keys) + ") would exceed max_distinct_keys (" + std::to_string(_max_distinct_keys) + ")."); } @@ -115,7 +116,7 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_ // Store the compacted batch. auto const new_batch_id = static_cast(_compacted_batches.size()); - auto const dense_id_offset = _distinct_keys; + auto const dense_id_offset = distinct_so_far; _compacted_batches.push_back(std::move(compacted)); _preprocessed_batches.push_back(preprocessed_compacted); @@ -138,7 +139,7 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_ update_transient_target_indices_fn{ base, slot_offsets.data(), _max_distinct_keys, target_indices.data()}); - _distinct_keys += new_distinct_keys; + _distinct_keys.fetch_add(new_distinct_keys, std::memory_order_relaxed); } // If new_distinct_keys == 0, target_indices is already final from Pass 1 — every // slot held a dense ID at probe time, so *iter was already the correct dense ID. diff --git a/cpp/src/groupby/streaming_groupby/merge.cu b/cpp/src/groupby/streaming_groupby/merge.cu index 6d79fdd7c18b..90efe95ba26d 100644 --- a/cpp/src/groupby/streaming_groupby/merge.cu +++ b/cpp/src/groupby/streaming_groupby/merge.cu @@ -18,6 +18,7 @@ #include #include +#include #include namespace cudf::groupby { @@ -90,6 +91,9 @@ struct merge_single_pass_aggs_fn { void streaming_groupby::impl::do_merge(impl const& other, cuda::stream_ref stream) { + // `other` is only read from, so a single lock on this object's insertion state is enough. + std::lock_guard const lock{_insert_mutex}; + CUDF_EXPECTS(!_invalidated, "streaming_groupby is in an invalidated state from a prior failure; " "no further aggregate()/merge() is allowed. finalize() may still be called."); @@ -99,8 +103,9 @@ void streaming_groupby::impl::do_merge(impl const& other, cuda::stream_ref strea CUDF_EXPECTS(_initialized, "Cannot merge into an uninitialized streaming_groupby. " "Call aggregate() at least once before merge()."); - CUDF_EXPECTS(other._distinct_keys <= _max_distinct_keys, - "Merge source distinct keys (" + std::to_string(other._distinct_keys) + + auto const other_keys_count = other._distinct_keys.load(std::memory_order_relaxed); + CUDF_EXPECTS(other_keys_count <= _max_distinct_keys, + "Merge source distinct keys (" + std::to_string(other_keys_count) + ") exceeds max_distinct_keys (" + std::to_string(_max_distinct_keys) + ").", std::invalid_argument); CUDF_EXPECTS(other._agg_kinds == _agg_kinds, @@ -117,14 +122,16 @@ void streaming_groupby::impl::do_merge(impl const& other, cuda::stream_ref strea auto other_keys = other.gather_distinct_keys(stream, mr); auto const other_key_view = other_keys->view(); - auto const other_distinct_keys = other._distinct_keys; + auto const other_distinct_keys = other._distinct_keys.load(std::memory_order_relaxed); if (other_distinct_keys == 0) { return; } update_nullable_state(other_key_view); if (!_key_set) { create_key_set(stream); } + _insert_done.wait(stream); auto result = probe_and_insert(other_key_view, stream); + _insert_done.record(stream); // Merge aggregation values using dense target indices. We only read from // `other._agg_results`; no need to deep-copy the source rows like keys. diff --git a/cpp/tests/groupby/streaming_groupby_test.cpp b/cpp/tests/groupby/streaming_groupby_test.cpp index 885e704ad94b..98add79386f7 100644 --- a/cpp/tests/groupby/streaming_groupby_test.cpp +++ b/cpp/tests/groupby/streaming_groupby_test.cpp @@ -19,8 +19,12 @@ #include #include +#include +#include #include +#include +#include #include static std::vector const KEY_COL{0}; @@ -325,6 +329,79 @@ TEST_F(StreamingGroupbyTest, MergeTwoObjects) check(keys, results, cudf::table_view{{ek}}, {ev}); } +TEST_F(StreamingGroupbyTest, ConcurrentAggregate) +{ + using K = int32_t; + using V = int32_t; + + constexpr int num_batches = 8; + + // Every batch re-hits keys 0 and 1 and introduces one key of its own, so concurrent calls + // both collide on existing groups and discover new keys at the same time. + std::vector> keys; + std::vector> vals; + keys.reserve(num_batches); + vals.reserve(num_batches); + for (int i = 0; i < num_batches; ++i) { + keys.emplace_back(std::initializer_list{0, 1, static_cast(i + 2)}); + vals.emplace_back(std::initializer_list{1, 10, 100}); + } + + std::vector batches; + batches.reserve(num_batches); + for (int i = 0; i < num_batches; ++i) { + batches.push_back(cudf::table_view{{keys[i], vals[i]}}); + } + + std::vector> streams; + streams.reserve(num_batches); + for (int i = 0; i < num_batches; ++i) { + streams.push_back(std::make_unique()); + } + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + + auto const device = rmm::get_current_cuda_device(); + std::vector threads; + std::vector errors(num_batches); + // `ready` lets the main thread wait until every worker is spinning, and `start` then releases + // them together, so the aggregate() calls actually overlap. + std::atomic ready{0}; + std::atomic start{false}; + threads.reserve(num_batches); + for (int i = 0; i < num_batches; ++i) { + threads.emplace_back([&, i] { + rmm::cuda_set_device_raii const device_guard{device}; + ready.fetch_add(1, std::memory_order_relaxed); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + try { + streaming_agg.aggregate(batches[i], streams[i]->view()); + } catch (...) { + errors[i] = std::current_exception(); + } + }); + } + while (ready.load(std::memory_order_relaxed) != num_batches) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto& thread : threads) { + thread.join(); + } + for (auto const& error : errors) { + EXPECT_FALSE(error); + } + for (auto const& stream : streams) { + stream->synchronize(); + } + + auto [out_keys, results] = streaming_agg.finalize(); + verify_against_groupby(out_keys, results, batches, KEY_COL, reqs); +} + TEST_F(StreamingGroupbyTest, EmptyBatch) { using K = int32_t;