From 605fe3f05d76ddc36323c0a272dd8a905b8a0268 Mon Sep 17 00:00:00 2001 From: Kevin Kristensen Date: Wed, 12 Aug 2026 21:00:33 +0000 Subject: [PATCH 1/4] fix(cudf): order GPU reads after writers --- .../cudf/gpu_data_representation.hpp | 24 +++--- src/cudf/gpu_data_representation.cpp | 46 +++++++++--- .../representation_converter_builtins.cpp | 73 +++++++------------ 3 files changed, 77 insertions(+), 66 deletions(-) diff --git a/include/cucascade/cudf/gpu_data_representation.hpp b/include/cucascade/cudf/gpu_data_representation.hpp index b6d463ad..33b83fd2 100644 --- a/include/cucascade/cudf/gpu_data_representation.hpp +++ b/include/cucascade/cudf/gpu_data_representation.hpp @@ -118,13 +118,15 @@ class gpu_table_representation : public idata_representation { std::size_t get_uncompressed_data_size_in_bytes() const override; /** - * @brief Create a deep copy of this GPU table representation. + * @brief Create an independently owned copy of this GPU table * - * The cloned representation will have its own copy of the underlying cuDF table, - * residing in the same memory space as the original. + * Orders the copy after the recorded writer event, or synchronizes the source device if no event + * is available. The copy uses this memory space's default allocator on @p stream, and the method + * synchronizes @p stream before returning. A stream with a non-null handle is recorded as the + * result's writer stream. * - * @param stream CUDA stream for memory operations - * @return std::unique_ptr A new gpu_table_representation with copied data + * @param stream Stream on this representation's device used for the copy + * @return Independently owned copy in the same memory space */ std::unique_ptr clone(rmm::cuda_stream_view stream) override; @@ -136,12 +138,16 @@ class gpu_table_representation : public idata_representation { cudf::table_view get_table_view() const; /** - * @brief Release ownership of the underlying cuDF table + * @brief Move out an owned cuDF table or materialize a table view * - * After calling this method, this representation no longer owns the table. + * An owned table is moved out without synchronization; the caller must order subsequent access + * after any outstanding writer work. A view-backed table is copied on @p stream using this memory + * space's default allocator after the recorded writer event, or after synchronizing the source + * device when no event exists. The method synchronizes @p stream before releasing the external + * owner. In either case, this representation is left without a table. * - * @param stream CUDA stream (used to materialize the table from a view path before release) - * @return std::unique_ptr The cuDF table + * @param stream Stream on this representation's device used only for view materialization + * @return Moved table or independently owned materialization of the table view */ std::unique_ptr release_table(rmm::cuda_stream_view stream); diff --git a/src/cudf/gpu_data_representation.cpp b/src/cudf/gpu_data_representation.cpp index d56985a3..9ae41dd2 100644 --- a/src/cudf/gpu_data_representation.cpp +++ b/src/cudf/gpu_data_representation.cpp @@ -23,6 +23,8 @@ #include #include +#include + namespace cucascade { gpu_table_representation::gpu_table_representation(std::unique_ptr table, @@ -73,11 +75,25 @@ cudf::table_view gpu_table_representation::get_table_view() const } } -std::unique_ptr gpu_table_representation::release_table( - [[maybe_unused]] rmm::cuda_stream_view stream) +std::unique_ptr gpu_table_representation::release_table(rmm::cuda_stream_view stream) { if (std::holds_alternative(_table)) { - _table = std::make_unique(std::get(_table).view, stream); + rmm::cuda_set_device_raii device_guard{rmm::cuda_device_id{get_device_id()}}; + // Wait for the latest writer before materializing the view. Eventless representations require + // a source-device synchronization. + if (_writer_event != nullptr) { + cucascade::cuda::cuda_event_view{_writer_event}.wait(stream); + } else { + CUCASCADE_CUDA_TRY(cudaDeviceSynchronize()); + } + + auto materialized = std::make_unique( + std::get(_table).view, stream, get_memory_space().get_default_allocator()); + // cuDF enqueues the deep copy asynchronously. Replacing the variant destroys the external + // owner, so the materialization stream must finish reading the view before that owner can + // release its source buffers. + stream.synchronize(); + _table = std::move(materialized); } return std::move(std::get>(_table)); } @@ -102,14 +118,22 @@ void gpu_table_representation::rebind_stream(rmm::cuda_stream_view stream) std::unique_ptr gpu_table_representation::clone(rmm::cuda_stream_view stream) { - // Create a deep copy of the cuDF table using the provided stream. - // STREAM-LINEAGE: the clone has been written by `stream`; record an event on - // it so any cross-stream/cross-device reader of the clone honors the - // producer-consumer ordering established by record_writer_event(). - cudf::table_view view = get_table_view(); - auto cloned = std::make_unique( - std::make_unique(view, stream), get_memory_space(), stream); - return cloned; + rmm::cuda_set_device_raii device_guard{rmm::cuda_device_id{get_device_id()}}; + // Wait for the latest writer before copying the source. Eventless representations require a + // source-device synchronization. + if (_writer_event != nullptr) { + cucascade::cuda::cuda_event_view{_writer_event}.wait(stream); + } else { + CUCASCADE_CUDA_TRY(cudaDeviceSynchronize()); + } + + auto cloned_table = std::make_unique( + get_table_view(), stream, get_memory_space().get_default_allocator()); + // The source may be destroyed as soon as clone() returns, so finish all asynchronous reads from + // it before publishing the independently owned result. + stream.synchronize(); + return std::make_unique( + std::move(cloned_table), get_memory_space(), stream); } void gpu_table_representation::record_writer_event(rmm::cuda_stream_view writer_stream) diff --git a/src/cudf/representation_converter_builtins.cpp b/src/cudf/representation_converter_builtins.cpp index aab773bf..e7ce39c5 100644 --- a/src/cudf/representation_converter_builtins.cpp +++ b/src/cudf/representation_converter_builtins.cpp @@ -44,6 +44,7 @@ #include #include +#include #include #include #include @@ -81,6 +82,22 @@ inline cudf::type_id as_cudf_type_id(int32_t type_id) return static_cast(type_id); } +// Orders `source_read_stream` after the source's latest recorded writer. An event-backed wait is +// asynchronous and does not extend source lifetime; callers must retain the source until their +// reads complete. If no event is recorded, this function synchronizes the source device before +// returning. The device associated with `source_read_stream` must be current on entry. +void wait_for_gpu_source(gpu_table_representation const& source, + rmm::cuda_stream_view source_read_stream) +{ + if (auto const writer_event = source.get_writer_event(); writer_event != nullptr) { + cuda::cuda_event_view{writer_event}.wait(source_read_stream); + return; + } + + rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; + CUCASCADE_CUDA_TRY(cudaDeviceSynchronize()); +} + // Forward declaration. convert_gpu_to_gpu is defined below convert_gpu_to_host_fast // so it can reuse BatchCopyAccumulator and the column-tree reconstruction helpers, // peer-copying each column buffer directly and avoiding cudf::pack (whose internal @@ -101,11 +118,9 @@ std::unique_ptr convert_gpu_to_host( rmm::cuda_stream_view stream, memory::reservation* reservation) { - // Synchronize the stream to ensure any prior operations (like table creation) - // are complete before we read from the source table - stream.synchronize(); - auto& gpu_source = source.cast(); + rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; + wait_for_gpu_source(gpu_source, stream); auto packed_data = cudf::pack(gpu_source.get_table_view(), stream); auto mr = target_memory_space->get_memory_resource_as(); @@ -497,7 +512,9 @@ std::unique_ptr convert_gpu_to_host_fast( rmm::cuda_stream_view stream, memory::reservation* reservation) { - auto& gpu_source = source.cast(); + auto& gpu_source = source.cast(); + rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; + wait_for_gpu_source(gpu_source, stream); const cudf::table_view view = gpu_source.get_table_view(); // --- Pass 1: plan the allocation layout --- @@ -866,11 +883,6 @@ std::unique_ptr convert_gpu_to_gpu( rmm::cuda_stream_view stream, [[maybe_unused]] memory::reservation* reservation) { - // Sync the caller's stream so the source table's buffers are stable on the source - // device before we issue peer copies. The caller's stream is the one that produced - // (or last touched) the source representation. - stream.synchronize(); - auto& gpu_source = source.cast(); // Same-device case: clone via source's own clone() method. @@ -881,25 +893,6 @@ std::unique_ptr convert_gpu_to_gpu( auto const src_device_id = gpu_source.get_device_id(); auto const dst_device_id = target_memory_space->get_device_id(); - // STREAM-LINEAGE INVARIANT: cross-device peer copies of cudaMallocAsync - // allocations require explicit event-ordered synchronization with the - // writer stream. A source-device-wide cudaDeviceSynchronize() does NOT - // establish the cross-mempool visibility the driver needs — under - // compute-sanitizer this site emits hundreds of stream-ordered-race errors - // even with a brute-force device sync. Producer-consumer pairing: - // producer = the stream that wrote gpu_source (recorded via - // gpu_table_representation::record_writer_event) - // consumer = target_stream (acquired from target memory space below) - // We resolve this in two passes: - // 1) Wait on the writer event (if recorded) on the *target* stream so the - // reader sees the writer's allocation/copy ordering. This is the precise - // primitive the sanitizer recognizes as closing the race. - // 2) Keep the source-device cudaDeviceSynchronize() as defense-in-depth for - // callers that have not yet been migrated to record writer events - // (get_writer_event() == nullptr). When the writer event is set the - // cudaDeviceSynchronize is technically redundant but harmless. - cudaEvent_t const writer_event = gpu_source.get_writer_event(); - rmm::cuda_set_device_raii target_guard{rmm::cuda_device_id{dst_device_id}}; // Target-bound stream from the target memory_space's stream pool. All peer copies @@ -907,21 +900,7 @@ std::unique_ptr convert_gpu_to_gpu( // completion without explicit cross-stream events. auto target_stream = target_memory_space->acquire_stream(); auto mr = target_memory_space->get_default_allocator(); - - if (writer_event != nullptr) { - // STREAM-LINEAGE pass 1: tie the reader stream's timeline to the writer's - // recorded event. After this point the target_stream observes all - // writer-side cudaMallocAsync allocations and writes in proper order. - cucascade::cuda::cuda_event_view{writer_event}.wait(target_stream); - } else { - // STREAM-LINEAGE pass 2 (fallback): no writer event recorded — fall back to - // a coarser source-device sync. This path is documented as insufficient for - // cross-mempool cudaMallocAsync allocations but is preserved for - // representations produced by code paths that have not yet been migrated to - // record_writer_event(). - rmm::cuda_set_device_raii src_sync_guard{rmm::cuda_device_id{src_device_id}}; - CUCASCADE_CUDA_TRY(cudaDeviceSynchronize()); - } + wait_for_gpu_source(gpu_source, target_stream); cudf::table_view const src_view = gpu_source.get_table_view(); @@ -1614,8 +1593,10 @@ static std::unique_ptr convert_gpu_to_disk( rmm::cuda_stream_view stream, [[maybe_unused]] memory::reservation* reservation) { - auto& backend = target_memory_space->get_io_backend(); - auto& gpu_source = source.cast(); + auto& backend = target_memory_space->get_io_backend(); + auto& gpu_source = source.cast(); + rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; + wait_for_gpu_source(gpu_source, stream); cudf::table_view tv = gpu_source.get_table_view(); // Generate unique file path under the disk memory space's mount directory From c5bac4c3c7b4c976c4a6a855f0d40452187601da Mon Sep 17 00:00:00 2001 From: Kevin Kristensen Date: Thu, 13 Aug 2026 14:43:17 +0000 Subject: [PATCH 2/4] test(cudf): cover cross-stream clone ordering --- test/data/test_data_representation.cpp | 156 ++++++++++++++++++ .../test_reservation_manager_configurator.cpp | 2 +- 2 files changed, 157 insertions(+), 1 deletion(-) diff --git a/test/data/test_data_representation.cpp b/test/data/test_data_representation.cpp index d2639539..ac5122e9 100644 --- a/test/data/test_data_representation.cpp +++ b/test/data/test_data_representation.cpp @@ -46,8 +46,13 @@ #include +#include #include +#include +#include +#include #include +#include #include #include #include @@ -639,6 +644,70 @@ TEST_CASE("Representations polymorphism", // Clone Tests // ============================================================================= +namespace { + +/** + * @brief Deterministically hold a CUDA stream inside a host callback until released. + * + * The callback only uses C++ atomics; CUDA APIs are forbidden from CUDA host callbacks. + */ +class cuda_stream_gate { + public: + cuda_stream_gate() = default; + cuda_stream_gate(cuda_stream_gate const&) = delete; + cuda_stream_gate& operator=(cuda_stream_gate const&) = delete; + + ~cuda_stream_gate() + { + release(); + if (_enqueued) { _exited.wait(false, std::memory_order_acquire); } + } + + void enqueue(rmm::cuda_stream_view stream) + { + CUCASCADE_CUDA_TRY(cudaLaunchHostFunc(stream.value(), &cuda_stream_gate::wait, this)); + _enqueued = true; + } + + void wait_until_entered() const { _entered.wait(false, std::memory_order_acquire); } + + void release() noexcept + { + _released.store(true, std::memory_order_release); + _released.notify_all(); + } + + private: + static void CUDART_CB wait(void* data) + { + auto& gate = *static_cast(data); + gate._entered.store(true, std::memory_order_release); + gate._entered.notify_all(); + gate._released.wait(false, std::memory_order_acquire); + gate._exited.store(true, std::memory_order_release); + gate._exited.notify_all(); + } + + bool _enqueued{false}; + mutable std::atomic _entered{false}; + std::atomic _released{false}; + std::atomic _exited{false}; +}; + +class scoped_stream_gate_release { + public: + explicit scoped_stream_gate_release(cuda_stream_gate& gate) : _gate(gate) {} + ~scoped_stream_gate_release() { _gate.release(); } + + scoped_stream_gate_release(scoped_stream_gate_release const&) = delete; + scoped_stream_gate_release& operator=(scoped_stream_gate_release const&) = delete; + + private: + cuda_stream_gate& _gate; +}; + +} // namespace + TEST_CASE("gpu_table_representation clone creates independent copy", "[gpu_data_representation]") { auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 0); @@ -674,6 +743,93 @@ TEST_CASE("gpu_table_representation clone creates independent copy", "[gpu_data_ } } +TEST_CASE("gpu_table_representation clone waits for a distinct writer stream", + "[gpu_data_representation][stream_ordering]") +{ + using namespace std::chrono_literals; + + memory::gpu_memory_space_config config; + config.device_id = 0; + config.memory_capacity = 64ULL << 20; + config.mr_factory_fn = test::make_shared_current_device_resource; + auto gpu_space = std::make_shared(config); + CUCASCADE_CUDA_TRY(cudaSetDevice(config.device_id)); + rmm::cuda_stream producer_stream; + rmm::cuda_stream consumer_stream; + + constexpr cudf::size_type num_rows = 1024; + constexpr std::size_t data_size = + static_cast(num_rows) * sizeof(std::int32_t); + constexpr unsigned char expected_byte = 0x5a; + + auto column = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, + num_rows, + cudf::mask_state::UNALLOCATED, + consumer_stream.view(), + gpu_space->get_default_allocator()); + + // Establish a known stale value before deliberately blocking the real producer write. + CUCASCADE_CUDA_TRY(cudaMemsetAsync( + column->mutable_view().head(), 0, data_size, consumer_stream.value())); + consumer_stream.synchronize(); + + cuda_stream_gate producer_gate; + producer_gate.enqueue(producer_stream.view()); + CUCASCADE_CUDA_TRY(cudaMemsetAsync(column->mutable_view().head(), + expected_byte, + data_size, + producer_stream.value())); + + std::vector> columns; + columns.push_back(std::move(column)); + gpu_table_representation source(std::make_unique(std::move(columns)), + *gpu_space, + producer_stream.view()); + + // Make sure the producer cannot reach either the write or source's recorded writer event. + producer_gate.wait_until_entered(); + auto const writer_status_while_blocked = cudaEventQuery(source.get_writer_event()); + + std::atomic clone_started{false}; + std::future> clone_future; + scoped_stream_gate_release release_on_exit{producer_gate}; + clone_future = std::async(std::launch::async, [&] { + CUCASCADE_CUDA_TRY(cudaSetDevice(source.get_device_id())); + clone_started.store(true, std::memory_order_release); + clone_started.notify_all(); + return source.clone(consumer_stream.view()); + }); + clone_started.wait(false, std::memory_order_acquire); + + // The fixed implementation waits for source's writer event and cannot return while the + // producer is gated. Main queues the copy without that wait and returns immediately. + auto const status_while_writer_blocked = clone_future.wait_for(1s); + if (status_while_writer_blocked == std::future_status::ready) { + // On the buggy implementation, finish the premature copy while the source still contains the + // stale pattern. This turns the ordering failure into deterministic data corruption too. + consumer_stream.synchronize(); + } + + producer_gate.release(); + auto cloned_base = clone_future.get(); + producer_stream.synchronize(); + consumer_stream.synchronize(); + + REQUIRE(writer_status_while_blocked == cudaErrorNotReady); + CHECK(status_while_writer_blocked == std::future_status::timeout); + + auto* clone = dynamic_cast(cloned_base.get()); + REQUIRE(clone != nullptr); + + std::vector bytes(data_size); + CUCASCADE_CUDA_TRY(cudaMemcpy(bytes.data(), + clone->get_table_view().column(0).head(), + data_size, + cudaMemcpyDeviceToHost)); + REQUIRE(std::all_of( + bytes.cbegin(), bytes.cend(), [](uint8_t value) { return value == expected_byte; })); +} + TEST_CASE("gpu_table_representation clone empty table", "[gpu_data_representation]") { auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 0); diff --git a/test/memory/test_reservation_manager_configurator.cpp b/test/memory/test_reservation_manager_configurator.cpp index fcb5001a..f9bd1e2f 100644 --- a/test/memory/test_reservation_manager_configurator.cpp +++ b/test/memory/test_reservation_manager_configurator.cpp @@ -26,7 +26,7 @@ #include #include -#include +#include #include #include From 160ca5f0d2021804619053d3be7d41fb66edf1e3 Mon Sep 17 00:00:00 2001 From: Kevin Kristensen Date: Thu, 13 Aug 2026 15:02:29 +0000 Subject: [PATCH 3/4] docs(cudf): streamline ordering comments --- src/cudf/gpu_data_representation.cpp | 6 ++-- test/data/test_data_representation.cpp | 39 +++++++++----------------- 2 files changed, 15 insertions(+), 30 deletions(-) diff --git a/src/cudf/gpu_data_representation.cpp b/src/cudf/gpu_data_representation.cpp index 9ae41dd2..5c5f472e 100644 --- a/src/cudf/gpu_data_representation.cpp +++ b/src/cudf/gpu_data_representation.cpp @@ -79,8 +79,7 @@ std::unique_ptr gpu_table_representation::release_table(rmm::cuda_s { if (std::holds_alternative(_table)) { rmm::cuda_set_device_raii device_guard{rmm::cuda_device_id{get_device_id()}}; - // Wait for the latest writer before materializing the view. Eventless representations require - // a source-device synchronization. + // Without an event, the producing stream is unknown and requires a device-wide fallback. if (_writer_event != nullptr) { cucascade::cuda::cuda_event_view{_writer_event}.wait(stream); } else { @@ -119,8 +118,7 @@ void gpu_table_representation::rebind_stream(rmm::cuda_stream_view stream) std::unique_ptr gpu_table_representation::clone(rmm::cuda_stream_view stream) { rmm::cuda_set_device_raii device_guard{rmm::cuda_device_id{get_device_id()}}; - // Wait for the latest writer before copying the source. Eventless representations require a - // source-device synchronization. + // Without an event, the producing stream is unknown and requires a device-wide fallback. if (_writer_event != nullptr) { cucascade::cuda::cuda_event_view{_writer_event}.wait(stream); } else { diff --git a/test/data/test_data_representation.cpp b/test/data/test_data_representation.cpp index ac5122e9..1b2d72d5 100644 --- a/test/data/test_data_representation.cpp +++ b/test/data/test_data_representation.cpp @@ -646,11 +646,7 @@ TEST_CASE("Representations polymorphism", namespace { -/** - * @brief Deterministically hold a CUDA stream inside a host callback until released. - * - * The callback only uses C++ atomics; CUDA APIs are forbidden from CUDA host callbacks. - */ +// CUDA host callbacks cannot call CUDA APIs, so use C++ atomics to gate the stream. class cuda_stream_gate { public: cuda_stream_gate() = default; @@ -757,9 +753,8 @@ TEST_CASE("gpu_table_representation clone waits for a distinct writer stream", rmm::cuda_stream producer_stream; rmm::cuda_stream consumer_stream; - constexpr cudf::size_type num_rows = 1024; - constexpr std::size_t data_size = - static_cast(num_rows) * sizeof(std::int32_t); + constexpr cudf::size_type num_rows = 1024; + constexpr std::size_t data_size = static_cast(num_rows) * sizeof(std::int32_t); constexpr unsigned char expected_byte = 0x5a; auto column = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, @@ -769,24 +764,20 @@ TEST_CASE("gpu_table_representation clone waits for a distinct writer stream", gpu_space->get_default_allocator()); // Establish a known stale value before deliberately blocking the real producer write. - CUCASCADE_CUDA_TRY(cudaMemsetAsync( - column->mutable_view().head(), 0, data_size, consumer_stream.value())); + CUCASCADE_CUDA_TRY( + cudaMemsetAsync(column->mutable_view().head(), 0, data_size, consumer_stream.value())); consumer_stream.synchronize(); cuda_stream_gate producer_gate; producer_gate.enqueue(producer_stream.view()); - CUCASCADE_CUDA_TRY(cudaMemsetAsync(column->mutable_view().head(), - expected_byte, - data_size, - producer_stream.value())); + CUCASCADE_CUDA_TRY(cudaMemsetAsync( + column->mutable_view().head(), expected_byte, data_size, producer_stream.value())); std::vector> columns; columns.push_back(std::move(column)); - gpu_table_representation source(std::make_unique(std::move(columns)), - *gpu_space, - producer_stream.view()); + gpu_table_representation source( + std::make_unique(std::move(columns)), *gpu_space, producer_stream.view()); - // Make sure the producer cannot reach either the write or source's recorded writer event. producer_gate.wait_until_entered(); auto const writer_status_while_blocked = cudaEventQuery(source.get_writer_event()); @@ -801,12 +792,10 @@ TEST_CASE("gpu_table_representation clone waits for a distinct writer stream", }); clone_started.wait(false, std::memory_order_acquire); - // The fixed implementation waits for source's writer event and cannot return while the - // producer is gated. Main queues the copy without that wait and returns immediately. + // clone() must not return until the source writer event can complete. auto const status_while_writer_blocked = clone_future.wait_for(1s); if (status_while_writer_blocked == std::future_status::ready) { - // On the buggy implementation, finish the premature copy while the source still contains the - // stale pattern. This turns the ordering failure into deterministic data corruption too. + // Complete a premature copy before releasing the producer to make stale data deterministic. consumer_stream.synchronize(); } @@ -822,10 +811,8 @@ TEST_CASE("gpu_table_representation clone waits for a distinct writer stream", REQUIRE(clone != nullptr); std::vector bytes(data_size); - CUCASCADE_CUDA_TRY(cudaMemcpy(bytes.data(), - clone->get_table_view().column(0).head(), - data_size, - cudaMemcpyDeviceToHost)); + CUCASCADE_CUDA_TRY(cudaMemcpy( + bytes.data(), clone->get_table_view().column(0).head(), data_size, cudaMemcpyDeviceToHost)); REQUIRE(std::all_of( bytes.cbegin(), bytes.cend(), [](uint8_t value) { return value == expected_byte; })); } From e62614a216f14f9039e73c76c37214cfcd389730 Mon Sep 17 00:00:00 2001 From: Kevin Kristensen Date: Wed, 16 Sep 2026 22:28:42 +0000 Subject: [PATCH 4/4] fix(data): make GPU cloning asynchronously safe --- docs/data-management.md | 10 +- .../cudf/gpu_data_representation.hpp | 46 ++-- include/cucascade/data/common.hpp | 13 +- include/cucascade/data/data_batch.hpp | 40 ++-- src/cudf/gpu_data_representation.cpp | 64 +++--- .../representation_converter_builtins.cpp | 85 +++++--- src/data/data_batch.cpp | 42 +++- test/cudf/test_release_table_stream.cpp | 135 ++++++++++-- test/data/test_data_batch.cpp | 201 ++++++++++++++++++ test/data/test_data_representation.cpp | 146 +------------ 10 files changed, 500 insertions(+), 282 deletions(-) diff --git a/docs/data-management.md b/docs/data-management.md index e0f2fce0..3637915b 100644 --- a/docs/data-management.md +++ b/docs/data-management.md @@ -88,7 +88,8 @@ public: }; ``` -- `clone()` performs a deep copy using `cudf::table(table.view(), stream)` +- `clone()` enqueues a deep copy with independently owned backing storage on the supplied stream + and returns without synchronizing - Owns the `cudf::table` object but not the underlying GPU memory (managed by the allocator) ### Host Table Representation (Direct Copy) @@ -280,9 +281,10 @@ auto cloned = ro.clone(new_batch_id, stream); auto cloned = ro.clone_to(registry, new_batch_id, host_space, stream); ``` -Cloning produces a new `shared_ptr` in `idle` state with the given ID. The clone -contains a deep copy of the data representation, residing in the same memory space as the original -(or a different space when using `clone_to`). +Cloning produces a new `shared_ptr` in `idle` state with the given ID and a deep copy of +the representation's backing data. A read-only batch accessor keeps the source protected until an +asynchronous clone's queued reads complete. Before consuming the clone on another stream, acquire a +read-only accessor for it and call `wait_until_ready(consumer_stream)`. --- diff --git a/include/cucascade/cudf/gpu_data_representation.hpp b/include/cucascade/cudf/gpu_data_representation.hpp index 270be733..13480a31 100644 --- a/include/cucascade/cudf/gpu_data_representation.hpp +++ b/include/cucascade/cudf/gpu_data_representation.hpp @@ -49,13 +49,8 @@ class gpu_table_representation : public idata_representation { /** * @brief Construct a new gpu_table_representation object. * - * STREAM-LINEAGE: every gpu_table_representation must be born with a recorded - * writer event so cross-stream / cross-device readers (notably - * representation_converter.cpp's convert_gpu_to_gpu()) can establish ordering - * via cudaStreamWaitEvent. The constructor calls record_writer_event(@p - * writer_stream) automatically — passing a default-constructed - * stream_ref records no event (legacy, only acceptable for paths whose - * data was never produced on any stream). + * Records a writer event on @p writer_stream so later readers can establish ordering. Default + * stream handles are valid and resolve to the current CUDA device. * * @param table Unique pointer to the cuDF table with the data (ownership is transferred) * @param memory_space The memory space where the GPU table resides @@ -117,15 +112,14 @@ class gpu_table_representation : public idata_representation { std::size_t get_uncompressed_data_size_in_bytes() const override; /** - * @brief Create an independently owned copy of this GPU table + * @brief Enqueue a deep copy of this GPU table. * - * Orders the copy after the recorded writer event, or synchronizes the source device if no event - * is available. The copy uses this memory space's default allocator on @p stream, and the method - * synchronizes @p stream before returning. A stream with a non-null handle is recorded as the - * result's writer stream. + * Orders @p stream after this representation's writer event, enqueues the copy using this memory + * space's allocator, records the result's writer event, and returns without synchronizing. The + * clone owns backing storage independent from this representation. * * @param stream Stream on this representation's device used for the copy - * @return Independently owned copy in the same memory space + * @return Independently owned deep copy in the same memory space */ std::unique_ptr clone(::cuda::stream_ref stream) override; @@ -139,14 +133,13 @@ class gpu_table_representation : public idata_representation { /** * @brief Move out an owned cuDF table or materialize a table view * - * An owned table is moved out without synchronization; the caller must order subsequent access - * after any outstanding writer work. A view-backed table is copied on @p stream using this memory - * space's default allocator after the recorded writer event, or after synchronizing the source - * device when no event exists. The method synchronizes @p stream before releasing the external - * owner. In either case, this representation is left without a table. + * A view-backed table is copied after its writer event and @p stream is synchronized before the + * external owner is released. In either case, this representation is left without a table. * - * @pre No stream other than @p stream may have in-flight work touching the table's device - * memory: binding the buffers to @p stream does not insert cross-stream ordering. + * @pre Untracked readers of the table's device memory must have completed. View-backed + * materialization waits for the recorded writer automatically. Releasing an owned table only + * rebinds its buffers, so the caller must order any in-flight work on other streams before the + * call; rebinding itself does not insert cross-stream ordering. * * @pre @p stream must belong to get_device_id() — the device owning this representation's * memory. Default stream handles resolve to the caller's current device. @@ -194,9 +187,8 @@ class gpu_table_representation : public idata_representation { * cudaStreamWaitEvent(reader_stream, get_writer_event(), 0) before peer-copying * source buffers. * - * Calling this multiple times overwrites the previously recorded event (the - * representation owns a single writer event handle that is reused). Passing a - * default-constructed stream_ref records no event and clears any prior one. + * Calling this multiple times records the same owned event at a new point. Default stream handles + * are valid and resolve to the current CUDA device. * * @param writer_stream The stream on which the most recent writes to this * representation's memory were enqueued. @@ -227,9 +219,7 @@ class gpu_table_representation : public idata_representation { std::variant, owning_table_view> _table; ///< cudf::table is the underlying representation of the data - /// Lazily-created CUDA event recording the completion of the most recent - /// writer-stream work that produced this representation. Null until the first - /// call to record_writer_event(). + /// CUDA event recording completion of the most recent writer-stream work. cudaEvent_t _writer_event{nullptr}; }; @@ -243,9 +233,7 @@ gpu_table_representation::gpu_table_representation(cudf::table_view table_view, _table( owning_table_view{std::make_any(std::forward(owner)), alloc_size, table_view}) { - // STREAM-LINEAGE: record writer event so cross-stream/cross-device readers - // can establish ordering via cudaStreamWaitEvent. - if (writer_stream.get() != nullptr) { record_writer_event(writer_stream); } + record_writer_event(writer_stream); } } // namespace cucascade diff --git a/include/cucascade/data/common.hpp b/include/cucascade/data/common.hpp index 3c1e8365..29686b67 100644 --- a/include/cucascade/data/common.hpp +++ b/include/cucascade/data/common.hpp @@ -103,11 +103,14 @@ class idata_representation { /** * @brief Create a deep copy of this data representation. * - * The cloned representation will have its own copy of the underlying data, - * residing in the same memory space as the original. - * - * @param stream CUDA stream for memory operations - * @return std::unique_ptr A new data representation with copied data + * The cloned representation has its own copy of the underlying data in the same memory space. + * Implementations that enqueue asynchronous work must put every source read and destination + * write on @p stream and return without synchronizing it. The caller must keep this + * representation alive and unmodified until that work completes, including when this function + * throws after enqueueing work. + * + * @param stream CUDA stream for asynchronous memory operations + * @return A new data representation with independently owned backing data */ virtual std::unique_ptr clone(::cuda::stream_ref stream) = 0; diff --git a/include/cucascade/data/data_batch.hpp b/include/cucascade/data/data_batch.hpp index e6a67ceb..664c66b2 100644 --- a/include/cucascade/data/data_batch.hpp +++ b/include/cucascade/data/data_batch.hpp @@ -422,13 +422,27 @@ class read_only_data_batch { } /** - * @brief Create an independent deep copy of the batch data. + * @brief Order @p stream after the batch's latest asynchronous write * - * The clone has a new batch ID and its own copy of the data representation, - * residing in the same memory space as the original. + * For GPU data, waits on the representation's writer event without synchronizing the host. If + * the representation has no writer event, synchronizes its device as a safe fallback. This is a + * no-op for non-GPU data. + * + * @param stream Stream that will read the batch + * @throws cucascade::logic_error if @p stream belongs to a different CUDA device + * @throws cucascade::cuda_error if CUDA stream/event operations fail + */ + void wait_until_ready(::cuda::stream_ref stream) const; + + /** + * @brief Create an independent deep copy of the batch data + * + * Asynchronous copies are enqueued on @p stream and this function returns without synchronizing + * it. The source batch remains protected from mutation and destruction until the enqueued reads + * complete. Call wait_until_ready() before consuming the clone on another stream. * * @param new_batch_id Batch ID for the cloned batch. - * @param stream CUDA stream for memory operations. + * @param stream CUDA stream for asynchronous memory operations. * @return A new data_batch wrapped in shared_ptr. * @throws std::runtime_error if the data is null. */ @@ -514,7 +528,7 @@ class read_only_data_batch { * * Holds an exclusive lock on the parent data_batch's mutex, permitting a single * writer with no concurrent readers. Provides all read methods plus write methods - * (set_data, convert_to) and clone operations (clone, clone_to). + * (set_data, convert_to) and clone-to operations (clone_to). * * Move-only. The exclusive lock is released when this object is destroyed or moved-from. */ @@ -598,22 +612,6 @@ class mutable_data_batch { */ void rebind_stream(::cuda::stream_ref stream); - /** - * @brief Create an independent deep copy of the batch data. - * - * The clone has a new batch ID and its own copy of the data representation, - * residing in the same memory space as the original. - * - * @param new_batch_id Batch ID for the cloned batch. - * @param stream CUDA stream for memory operations. - * @return A new data_batch wrapped in shared_ptr. - * @throws std::runtime_error if the data is null. - */ - [[nodiscard]] std::shared_ptr clone( - uint64_t new_batch_id, - ::cuda::stream_ref stream, - std::unique_ptr probe = std::make_unique()) const; - /** * @brief Create an independent deep copy with representation conversion. * diff --git a/src/cudf/gpu_data_representation.cpp b/src/cudf/gpu_data_representation.cpp index 968a3895..e4f6d2ca 100644 --- a/src/cudf/gpu_data_representation.cpp +++ b/src/cudf/gpu_data_representation.cpp @@ -59,22 +59,18 @@ gpu_table_representation::gpu_table_representation(std::unique_ptr ::cuda::stream_ref writer_stream) : idata_representation(memory_space), _table(std::move(table)) { - // STREAM-LINEAGE: record the writer event in the constructor body so every - // representation is born with a recorded event. Skipping when the caller - // passes a default-constructed (per-thread default) stream view preserves - // legacy behavior for callers that genuinely have no writer stream — they - // will fall back to cudaDeviceSynchronize on the source device in - // convert_gpu_to_gpu(). All non-legacy callers MUST pass a real writer - // stream. - if (writer_stream.get() != nullptr) { record_writer_event(writer_stream); } + record_writer_event(writer_stream); } gpu_table_representation::~gpu_table_representation() { - // STREAM-LINEAGE: release the writer event if one was recorded. if (_writer_event != nullptr) { - CUCASCADE_ASSERT_CUDA_SUCCESS(cudaEventDestroy(_writer_event)); + int original_device = -1; + bool const restore_device = ::cudaGetDevice(&original_device) == cudaSuccess; + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaSetDevice(get_device_id())); + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaEventDestroy(_writer_event)); _writer_event = nullptr; + if (restore_device) { CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaSetDevice(original_device)); } } } @@ -107,20 +103,24 @@ std::unique_ptr gpu_table_representation::release_table(::cuda::str if (std::holds_alternative(_table)) { validate_stream_device(stream, get_device_id()); rmm::cuda_set_device_raii device_guard{rmm::cuda_device_id{get_device_id()}}; - // Without an event, the producing stream is unknown and requires a device-wide fallback. if (_writer_event != nullptr) { - cucascade::cuda::cuda_event_view{_writer_event}.wait(stream); + cuda::cuda_event_view{_writer_event}.wait(stream); } else { - CUCASCADE_CUDA_TRY(cudaDeviceSynchronize()); + CUCASCADE_CUDA_TRY(::cudaDeviceSynchronize()); } - auto materialized = std::make_unique( - std::get(_table).view, stream, get_memory_space().get_default_allocator()); - // cuDF enqueues the deep copy asynchronously. Replacing the variant destroys the external - // owner, so the materialization stream must finish reading the view before that owner can - // release its source buffers. - stream.sync(); - _table = std::move(materialized); + try { + auto materialized = std::make_unique(std::get(_table).view, + stream, + get_memory_space().get_default_allocator()); + stream.sync(); + _table = std::move(materialized); + } catch (...) { + // Construction may enqueue reads before throwing. Keep the view owner in _table and drain + // those reads without throwing while the source device is still current. + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaStreamSynchronize(stream.get())); + throw; + } } else { // Rebind so the returned table's frees stay stream-ordered behind the caller's reads. // rebind_stream() applies the same device guard, so this branch needs no separate check. @@ -153,29 +153,35 @@ void gpu_table_representation::rebind_stream(::cuda::stream_ref stream) std::unique_ptr gpu_table_representation::clone(::cuda::stream_ref stream) { + validate_stream_device(stream, get_device_id()); rmm::cuda_set_device_raii device_guard{rmm::cuda_device_id{get_device_id()}}; - // Without an event, the producing stream is unknown and requires a device-wide fallback. if (_writer_event != nullptr) { - cucascade::cuda::cuda_event_view{_writer_event}.wait(stream); + cuda::cuda_event_view{_writer_event}.wait(stream); } else { - CUCASCADE_CUDA_TRY(cudaDeviceSynchronize()); + CUCASCADE_CUDA_TRY(::cudaDeviceSynchronize()); } auto cloned_table = std::make_unique( get_table_view(), stream, get_memory_space().get_default_allocator()); - // The source may be destroyed as soon as clone() returns, so finish all asynchronous reads from - // it before publishing the independently owned result. - stream.synchronize(); return std::make_unique( std::move(cloned_table), get_memory_space(), stream); } void gpu_table_representation::record_writer_event(::cuda::stream_ref writer_stream) { - // STREAM-LINEAGE: lazily create the event on first call (cudaEventDisableTiming — - // used solely for cross-stream ordering, never for elapsed-time queries). + validate_stream_device(writer_stream, get_device_id()); + rmm::cuda_set_device_raii device_guard{rmm::cuda_device_id{get_device_id()}}; if (_writer_event == nullptr) { - CUCASCADE_CUDA_TRY(cudaEventCreateWithFlags(&_writer_event, cudaEventDisableTiming)); + cudaEvent_t new_event = nullptr; + CUCASCADE_CUDA_TRY(::cudaEventCreateWithFlags(&new_event, cudaEventDisableTiming)); + try { + cucascade::cuda::cuda_event_view{new_event}.record(writer_stream); + } catch (...) { + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaEventDestroy(new_event)); + throw; + } + _writer_event = new_event; + return; } cucascade::cuda::cuda_event_view{_writer_event}.record(writer_stream); } diff --git a/src/cudf/representation_converter_builtins.cpp b/src/cudf/representation_converter_builtins.cpp index 1a466d2b..250271f7 100644 --- a/src/cudf/representation_converter_builtins.cpp +++ b/src/cudf/representation_converter_builtins.cpp @@ -86,22 +86,6 @@ inline cudf::type_id as_cudf_type_id(int32_t type_id) return static_cast(type_id); } -// Orders `source_read_stream` after the source's latest recorded writer. An event-backed wait is -// asynchronous and does not extend source lifetime; callers must retain the source until their -// reads complete. If no event is recorded, this function synchronizes the source device before -// returning. The device associated with `source_read_stream` must be current on entry. -void wait_for_gpu_source(gpu_table_representation const& source, - rmm::cuda_stream_view source_read_stream) -{ - if (auto const writer_event = source.get_writer_event(); writer_event != nullptr) { - cuda::cuda_event_view{writer_event}.wait(source_read_stream); - return; - } - - rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; - CUCASCADE_CUDA_TRY(cudaDeviceSynchronize()); -} - // Forward declaration. convert_gpu_to_gpu is defined below convert_gpu_to_host_fast // so it can reuse BatchCopyAccumulator and the column-tree reconstruction helpers, // peer-copying each column buffer directly and avoiding cudf::pack (whose internal @@ -122,9 +106,11 @@ std::unique_ptr convert_gpu_to_host( ::cuda::stream_ref stream, memory::reservation* reservation) { + // Synchronize the stream to ensure any prior operations (like table creation) + // are complete before we read from the source table + stream.sync(); + auto& gpu_source = source.cast(); - rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; - wait_for_gpu_source(gpu_source, stream); auto packed_data = cudf::pack(gpu_source.get_table_view(), stream); auto mr = target_memory_space->get_memory_resource_as(); @@ -516,9 +502,7 @@ std::unique_ptr convert_gpu_to_host_fast( ::cuda::stream_ref stream, memory::reservation* reservation) { - auto& gpu_source = source.cast(); - rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; - wait_for_gpu_source(gpu_source, stream); + auto& gpu_source = source.cast(); const cudf::table_view view = gpu_source.get_table_view(); // --- Pass 1: plan the allocation layout --- @@ -891,12 +875,47 @@ std::unique_ptr convert_gpu_to_gpu( // Same-device case: clone via source's own clone() method. if (source.get_device_id() == target_memory_space->get_device_id()) { - return source.clone(stream); + rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; + stream.sync(); + try { + auto result = source.clone(stream); + stream.sync(); + return result; + } catch (...) { + // clone() may have enqueued source reads before throwing. Complete them before clone_to() + // releases its source lock or a mutable accessor reuses the source representation. + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaStreamSynchronize(stream.get())); + throw; + } } + // Sync the caller's stream so the source table's buffers are stable on the source + // device before we issue peer copies. The caller's stream is the one that produced + // (or last touched) the source representation. + stream.sync(); + auto const src_device_id = gpu_source.get_device_id(); auto const dst_device_id = target_memory_space->get_device_id(); + // STREAM-LINEAGE INVARIANT: cross-device peer copies of cudaMallocAsync + // allocations require explicit event-ordered synchronization with the + // writer stream. A source-device-wide cudaDeviceSynchronize() does NOT + // establish the cross-mempool visibility the driver needs — under + // compute-sanitizer this site emits hundreds of stream-ordered-race errors + // even with a brute-force device sync. Producer-consumer pairing: + // producer = the stream that wrote gpu_source (recorded via + // gpu_table_representation::record_writer_event) + // consumer = target_stream (acquired from target memory space below) + // We resolve this in two passes: + // 1) Wait on the writer event (if recorded) on the *target* stream so the + // reader sees the writer's allocation/copy ordering. This is the precise + // primitive the sanitizer recognizes as closing the race. + // 2) Keep the source-device cudaDeviceSynchronize() as defense-in-depth for + // callers that have not yet been migrated to record writer events + // (get_writer_event() == nullptr). When the writer event is set the + // cudaDeviceSynchronize is technically redundant but harmless. + cudaEvent_t const writer_event = gpu_source.get_writer_event(); + rmm::cuda_set_device_raii target_guard{rmm::cuda_device_id{dst_device_id}}; // Target-bound stream from the target memory_space's stream pool. All peer copies @@ -904,7 +923,21 @@ std::unique_ptr convert_gpu_to_gpu( // completion without explicit cross-stream events. auto target_stream = target_memory_space->acquire_stream(); auto mr = target_memory_space->get_default_allocator(); - wait_for_gpu_source(gpu_source, target_stream); + + if (writer_event != nullptr) { + // STREAM-LINEAGE pass 1: tie the reader stream's timeline to the writer's + // recorded event. After this point the target_stream observes all + // writer-side cudaMallocAsync allocations and writes in proper order. + cucascade::cuda::cuda_event_view{writer_event}.wait(target_stream); + } else { + // STREAM-LINEAGE pass 2 (fallback): no writer event recorded — fall back to + // a coarser source-device sync. This path is documented as insufficient for + // cross-mempool cudaMallocAsync allocations but is preserved for + // representations produced by code paths that have not yet been migrated to + // record_writer_event(). + rmm::cuda_set_device_raii src_sync_guard{rmm::cuda_device_id{src_device_id}}; + CUCASCADE_CUDA_TRY(cudaDeviceSynchronize()); + } cudf::table_view const src_view = gpu_source.get_table_view(); @@ -1626,10 +1659,8 @@ static std::unique_ptr convert_gpu_to_disk( ::cuda::stream_ref stream, [[maybe_unused]] memory::reservation* reservation) { - auto& backend = target_memory_space->get_io_backend(); - auto& gpu_source = source.cast(); - rmm::cuda_set_device_raii source_device_guard{rmm::cuda_device_id{source.get_device_id()}}; - wait_for_gpu_source(gpu_source, stream); + auto& backend = target_memory_space->get_io_backend(); + auto& gpu_source = source.cast(); cudf::table_view tv = gpu_source.get_table_view(); // Generate unique file path under the disk memory space's mount directory diff --git a/src/data/data_batch.cpp b/src/data/data_batch.cpp index e43ae1ee..8ce505f0 100644 --- a/src/data/data_batch.cpp +++ b/src/data/data_batch.cpp @@ -15,6 +15,7 @@ * limitations under the License. */ +#include #include #include @@ -24,6 +25,7 @@ #include #include #include +#include namespace cucascade { @@ -359,11 +361,41 @@ read_only_data_batch::~read_only_data_batch() } } +void read_only_data_batch::wait_until_ready(::cuda::stream_ref stream) const +{ + auto const* representation = get_data(); + if (representation == nullptr || representation->get_current_tier() != memory::Tier::GPU) { + return; + } + + int stream_device = -1; + CUCASCADE_CUDA_TRY(::cudaStreamGetDevice(stream.get(), &stream_device)); + if (stream_device != representation->get_device_id()) { + CUCASCADE_FAIL("stream belongs to CUDA device " + std::to_string(stream_device) + + " but this batch's data lives on device " + + std::to_string(representation->get_device_id())); + } + + rmm::cuda_set_device_raii device_guard{rmm::cuda_device_id{stream_device}}; + if (auto const writer_event = representation->get_writer_event(); writer_event != nullptr) { + cuda::cuda_event_view{writer_event}.wait(stream); + } else { + CUCASCADE_CUDA_TRY(::cudaDeviceSynchronize()); + } +} + std::shared_ptr read_only_data_batch::clone( uint64_t new_batch_id, ::cuda::stream_ref stream, std::unique_ptr probe) const { if (_batch->_data == nullptr) { throw std::runtime_error("Cannot clone: data is null"); } - auto cloned_data = _batch->_data->clone(stream); + std::unique_ptr cloned_data; + try { + cloned_data = _batch->_data->clone(stream); + } catch (...) { + record_reader_event(stream); + throw; + } + record_reader_event(stream); return data_batch::make(new_batch_id, std::move(cloned_data), std::move(probe)); } @@ -410,12 +442,4 @@ void mutable_data_batch::rebind_stream(::cuda::stream_ref stream) if (auto* repr = _batch->get_data()) { repr->rebind_stream(stream); } } -std::shared_ptr mutable_data_batch::clone( - uint64_t new_batch_id, ::cuda::stream_ref stream, std::unique_ptr probe) const -{ - if (_batch->_data == nullptr) { throw std::runtime_error("Cannot clone: data is null"); } - auto cloned_data = _batch->_data->clone(stream); - return data_batch::make(new_batch_id, std::move(cloned_data), std::move(probe)); -} - } // namespace cucascade diff --git a/test/cudf/test_release_table_stream.cpp b/test/cudf/test_release_table_stream.cpp index f54ae19e..4cf7b43e 100644 --- a/test/cudf/test_release_table_stream.cpp +++ b/test/cudf/test_release_table_stream.cpp @@ -43,9 +43,11 @@ #include +#include #include #include #include +#include #include #include #include @@ -205,6 +207,44 @@ void CUDART_CB stall_stream_callback(void* /*user_data*/) std::this_thread::sleep_for(std::chrono::milliseconds(40)); } +void CUDART_CB wait_for_release_callback(void* user_data) +{ + auto& released = *static_cast*>(user_data); + released.wait(false, std::memory_order_acquire); +} + +class release_gate_cleanup { + public: + release_gate_cleanup(std::atomic& gate, ::cuda::stream_ref stream, int device_id) + : _gate(gate), _stream(stream), _device_id(device_id) + { + } + + ~release_gate_cleanup() noexcept + { + release(); + int original_device = -1; + bool const restore_device = ::cudaGetDevice(&original_device) == cudaSuccess; + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaSetDevice(_device_id)); + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaStreamSynchronize(_stream.get())); + if (restore_device) { CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaSetDevice(original_device)); } + } + + release_gate_cleanup(release_gate_cleanup const&) = delete; + release_gate_cleanup& operator=(release_gate_cleanup const&) = delete; + + void release() noexcept + { + _gate.store(true, std::memory_order_release); + _gate.notify_all(); + } + + private: + std::atomic& _gate; + ::cuda::stream_ref _stream; + int _device_id; +}; + /// Pinned so the D2H readback enqueues asynchronously instead of staging synchronously. struct pinned_buffer { void* ptr{nullptr}; @@ -375,6 +415,80 @@ TEST_CASE("view-branch release_table deep-copies on the release stream and leave REQUIRE(src_checked >= 5); } +TEST_CASE("view-branch release_table keeps its sole owner alive until materialization completes", + "[release_table][view][lifetime]") +{ + using namespace std::chrono_literals; + + auto& gpu_space = async_gpu_space(); + rmm::cuda_stream producer_stream_storage; + rmm::cuda_stream release_stream_storage; + ::cuda::stream_ref producer_stream = producer_stream_storage; + ::cuda::stream_ref release_stream = release_stream_storage; + + constexpr cudf::size_type num_rows = 1 << 18; + std::vector expected(static_cast(num_rows)); + std::iota(expected.begin(), expected.end(), 17); + + auto column = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, + num_rows, + cudf::mask_state::UNALLOCATED, + producer_stream, + gpu_space->get_default_allocator()); + CUCASCADE_CUDA_TRY(cudaMemcpyAsync(column->mutable_view().data(), + expected.data(), + expected.size() * sizeof(int32_t), + cudaMemcpyHostToDevice, + producer_stream.get())); + producer_stream.sync(); + + std::vector> columns; + columns.push_back(std::move(column)); + auto owner = std::make_shared(std::move(columns)); + std::weak_ptr owner_lifetime = owner; + gpu_table_representation rep(owner->view(), + std::shared_ptr{owner}, + owner->alloc_size(), + *gpu_space, + producer_stream); + owner.reset(); + REQUIRE_FALSE(owner_lifetime.expired()); + + std::atomic release_gate{false}; + std::atomic worker_started{false}; + std::future> release_future; + CUCASCADE_CUDA_TRY( + cudaLaunchHostFunc(release_stream.get(), wait_for_release_callback, &release_gate)); + release_gate_cleanup release_gate_on_exit{release_gate, release_stream, 0}; + release_future = std::async(std::launch::async, [&] { + worker_started.store(true, std::memory_order_release); + worker_started.notify_all(); + CUCASCADE_CUDA_TRY(cudaSetDevice(0)); + return rep.release_table(release_stream); + }); + worker_started.wait(false, std::memory_order_acquire); + + auto const blocked_status = release_future.wait_for(100ms); + bool const owner_alive_while_copy_is_blocked = !owner_lifetime.expired(); + release_gate_on_exit.release(); + + auto released = release_future.get(); + release_stream.sync(); + + REQUIRE(blocked_status == std::future_status::timeout); + REQUIRE(owner_alive_while_copy_is_blocked); + REQUIRE(owner_lifetime.expired()); + + std::vector actual(expected.size()); + CUCASCADE_CUDA_TRY(cudaMemcpyAsync(actual.data(), + released->view().column(0).data(), + actual.size() * sizeof(int32_t), + cudaMemcpyDeviceToHost, + release_stream.get())); + release_stream.sync(); + REQUIRE(actual == expected); +} + TEST_CASE("release_table then cudf::rebind_stream to the same stream composes", "[release_table][stream]") { @@ -430,24 +544,17 @@ TEST_CASE("release_table accepts every same-device stream handle", } } -TEST_CASE("release_table rejects a stream whose device differs from the representation's", - "[release_table][stream][device]") +TEST_CASE("gpu_table_representation rejects a writer stream from another device", + "[gpu_data_representation][stream][device]") { - // Exercises the guard's comparison on any host, including single-GPU CI where the true - // multi-GPU cases below can only skip. The space claims device 1 while the table's memory and - // stream are really on device 0, so get_device_id() and the stream device disagree — which is - // exactly the mismatch the guard exists to catch. rmm::cuda_set_device_raii const pin_device{rmm::cuda_device_id{0}}; auto mismatched_space = test::make_mock_memory_space(memory::Tier::GPU, 1); - gpu_table_representation rep( - make_patterned_table(shared_stream()), *mismatched_space, shared_stream()); - - REQUIRE_THROWS_AS(rep.release_table(shared_stream()), cucascade::logic_error); - REQUIRE_THROWS_AS(rep.rebind_stream(shared_stream()), cucascade::logic_error); - - // Rejected before any mutation: the representation still owns its table. - REQUIRE(rep.get_table_view().num_columns() == 3); + auto construct_with_mismatched_stream = [&] { + gpu_table_representation rep( + make_patterned_table(shared_stream()), *mismatched_space, shared_stream()); + }; + REQUIRE_THROWS_AS(construct_with_mismatched_stream(), cucascade::logic_error); } TEST_CASE("release_table and rebind_stream reject a stream owned by another device", diff --git a/test/data/test_data_batch.cpp b/test/data/test_data_batch.cpp index e1acbc0a..4cb04d6a 100644 --- a/test/data/test_data_batch.cpp +++ b/test/data/test_data_batch.cpp @@ -19,6 +19,7 @@ #include "utils/mock_test_utils.hpp" #include +#include #include #include @@ -36,6 +37,7 @@ #include #include #include +#include #include #include #include @@ -1098,6 +1100,205 @@ static void CUDART_CB stream_gate_callback(void* user_data) gate->released.wait(false, std::memory_order_acquire); } +class throwing_async_clone_representation : private cucascade::test::mock_memory_space_holder, + public idata_representation { + public: + explicit throwing_async_clone_representation(stream_gate& gate) + : mock_memory_space_holder(memory::Tier::GPU, 0), idata_representation(*space), _gate(gate) + { + } + + std::size_t get_size_in_bytes() const override { return 0; } + std::size_t get_uncompressed_data_size_in_bytes() const override { return 0; } + + std::unique_ptr clone(::cuda::stream_ref stream) override + { + CUCASCADE_CUDA_TRY(::cudaLaunchHostFunc(stream.get(), stream_gate_callback, &_gate)); + throw std::runtime_error("clone failed after enqueue"); + } + + private: + stream_gate& _gate; +}; + +TEST_CASE("read-only GPU clone is asynchronous and carries both ordering dependencies", + "[data_batch][clone][reader_event][stream_ordering]") +{ + using namespace std::chrono_literals; + + constexpr cudf::size_type num_rows = 1 << 18; + constexpr std::size_t data_size = static_cast(num_rows) * sizeof(std::int32_t); + auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 0); + rmm::cuda_stream producer_stream_storage; + rmm::cuda_stream clone_stream_storage; + rmm::cuda_stream consumer_stream_storage; + ::cuda::stream_ref producer_stream = producer_stream_storage; + ::cuda::stream_ref clone_stream = clone_stream_storage; + ::cuda::stream_ref consumer_stream = consumer_stream_storage; + + auto column = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, + num_rows, + cudf::mask_state::UNALLOCATED, + producer_stream, + gpu_space->get_default_allocator()); + CUCASCADE_CUDA_TRY( + ::cudaMemsetAsync(column->mutable_view().head(), 0, data_size, producer_stream.get())); + producer_stream.sync(); + + stream_gate gate; + CUCASCADE_CUDA_TRY(::cudaLaunchHostFunc(producer_stream.get(), stream_gate_callback, &gate)); + stream_gate_release_guard release_gate_on_exit{gate, producer_stream}; + CUCASCADE_CUDA_TRY( + ::cudaMemsetAsync(column->mutable_view().head(), 0x5A, data_size, producer_stream.get())); + std::vector> columns; + columns.push_back(std::move(column)); + auto batch = data_batch::make( + 1, + std::make_unique( + std::make_unique(std::move(columns)), *gpu_space, producer_stream)); + + std::future> clone_future; + std::future consumer_future; + auto source = batch->to_read_only(); + clone_future = std::async(std::launch::async, [&] { + CUCASCADE_CUDA_TRY(::cudaSetDevice(0)); + return source.clone(2, clone_stream); + }); + + auto const clone_status = clone_future.wait_for(1s); + bool const returned_asynchronously = clone_status == std::future_status::ready; + if (!returned_asynchronously) { gate.release(); } + auto clone = clone_future.get(); + + auto clone_reader = clone->to_read_only(); + auto const clone_writer_event = clone_reader.get_writer_event(); + REQUIRE(clone_writer_event != nullptr); + bool const destination_pending = + returned_asynchronously && ::cudaEventQuery(clone_writer_event) == cudaErrorNotReady; + + batch = data_batch::to_idle(std::move(source)); + auto mutable_source = batch->try_to_mutable(); + bool const source_reclamation_blocked = !mutable_source.has_value(); + mutable_source.reset(); + + clone_reader.wait_until_ready(consumer_stream); + std::atomic consumer_started{false}; + consumer_future = std::async(std::launch::async, [&] { + consumer_started.store(true, std::memory_order_release); + consumer_started.notify_all(); + consumer_stream.sync(); + }); + consumer_started.wait(false, std::memory_order_acquire); + auto const consumer_status = consumer_future.wait_for(100ms); + + gate.release(); + consumer_future.get(); + clone_stream.sync(); + + auto const* clone_representation = + dynamic_cast(clone_reader.get_data()); + REQUIRE(clone_representation != nullptr); + std::vector copied_bytes(data_size); + CUCASCADE_CUDA_TRY(::cudaMemcpy(copied_bytes.data(), + clone_representation->get_table_view().column(0).head(), + data_size, + cudaMemcpyDeviceToHost)); + + REQUIRE(returned_asynchronously); + REQUIRE(destination_pending); + REQUIRE(source_reclamation_blocked); + REQUIRE(consumer_status == std::future_status::timeout); + REQUIRE(std::all_of( + copied_bytes.begin(), copied_bytes.end(), [](std::uint8_t value) { return value == 0x5A; })); + REQUIRE(batch->try_to_mutable().has_value()); +} + +TEST_CASE("read-only clone records a reader when the implementation throws after enqueue", + "[data_batch][clone][reader_event]") +{ + rmm::cuda_stream stream_storage; + ::cuda::stream_ref stream = stream_storage; + stream_gate gate; + stream_gate_release_guard release_gate_on_exit{gate, stream}; + auto batch = data_batch::make(1, std::make_unique(gate)); + auto source = batch->to_read_only(); + + REQUIRE_THROWS_WITH(source.clone(2, stream), "clone failed after enqueue"); + batch = data_batch::to_idle(std::move(source)); + REQUIRE_FALSE(batch->try_to_mutable().has_value()); + + gate.release(); + stream.sync(); + REQUIRE(batch->try_to_mutable().has_value()); +} + +TEST_CASE("same-device mutable clone_to returns a ready GPU deep copy", + "[data_batch][clone_to][gpu][stream_ordering]") +{ + using namespace std::chrono_literals; + + constexpr cudf::size_type num_rows = 1 << 18; + constexpr std::size_t data_size = static_cast(num_rows) * sizeof(std::int32_t); + auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 0); + rmm::cuda_stream producer_stream_storage; + rmm::cuda_stream clone_stream_storage; + ::cuda::stream_ref producer_stream = producer_stream_storage; + ::cuda::stream_ref clone_stream = clone_stream_storage; + + auto column = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, + num_rows, + cudf::mask_state::UNALLOCATED, + producer_stream, + gpu_space->get_default_allocator()); + stream_gate gate; + CUCASCADE_CUDA_TRY(::cudaLaunchHostFunc(producer_stream.get(), stream_gate_callback, &gate)); + stream_gate_release_guard release_gate_on_exit{gate, producer_stream}; + CUCASCADE_CUDA_TRY( + ::cudaMemsetAsync(column->mutable_view().head(), 0x5A, data_size, producer_stream.get())); + + std::vector> columns; + columns.push_back(std::move(column)); + auto batch = data_batch::make( + 1, + std::make_unique( + std::make_unique(std::move(columns)), *gpu_space, producer_stream)); + + representation_converter_registry registry; + register_builtin_converters(registry); + auto source = batch->to_mutable(); + std::atomic clone_started{false}; + auto clone_future = std::async( + std::launch::async, + [source = std::move(source), ®istry, clone_stream, &gpu_space, &clone_started]() mutable { + CUCASCADE_CUDA_TRY(::cudaSetDevice(0)); + clone_started.store(true, std::memory_order_release); + clone_started.notify_all(); + return source.clone_to(registry, 2, gpu_space.get(), clone_stream); + }); + clone_started.wait(false, std::memory_order_acquire); + + auto const clone_status = clone_future.wait_for(100ms); + gate.release(); + auto clone = clone_future.get(); + + auto clone_reader = clone->to_read_only(); + auto const clone_writer_event = clone_reader.get_writer_event(); + REQUIRE(clone_status == std::future_status::timeout); + REQUIRE(clone_writer_event != nullptr); + REQUIRE(::cudaEventQuery(clone_writer_event) == cudaSuccess); + + auto const* clone_representation = + dynamic_cast(clone_reader.get_data()); + REQUIRE(clone_representation != nullptr); + std::vector copied_bytes(data_size); + CUCASCADE_CUDA_TRY(::cudaMemcpy(copied_bytes.data(), + clone_representation->get_table_view().column(0).head(), + data_size, + cudaMemcpyDeviceToHost)); + REQUIRE(std::all_of( + copied_bytes.begin(), copied_bytes.end(), [](std::uint8_t value) { return value == 0x5A; })); +} + TEST_CASE("mutable_data_batch holds exclusive lock during convert_to stream sync", "[data_batch][convert_to]") { diff --git a/test/data/test_data_representation.cpp b/test/data/test_data_representation.cpp index eeabf234..ee4049a2 100644 --- a/test/data/test_data_representation.cpp +++ b/test/data/test_data_representation.cpp @@ -46,13 +46,8 @@ #include -#include #include -#include -#include -#include #include -#include #include #include #include @@ -658,66 +653,6 @@ TEST_CASE("Representations polymorphism", // Clone Tests // ============================================================================= -namespace { - -// CUDA host callbacks cannot call CUDA APIs, so use C++ atomics to gate the stream. -class cuda_stream_gate { - public: - cuda_stream_gate() = default; - cuda_stream_gate(cuda_stream_gate const&) = delete; - cuda_stream_gate& operator=(cuda_stream_gate const&) = delete; - - ~cuda_stream_gate() - { - release(); - if (_enqueued) { _exited.wait(false, std::memory_order_acquire); } - } - - void enqueue(rmm::cuda_stream_view stream) - { - CUCASCADE_CUDA_TRY(cudaLaunchHostFunc(stream.value(), &cuda_stream_gate::wait, this)); - _enqueued = true; - } - - void wait_until_entered() const { _entered.wait(false, std::memory_order_acquire); } - - void release() noexcept - { - _released.store(true, std::memory_order_release); - _released.notify_all(); - } - - private: - static void CUDART_CB wait(void* data) - { - auto& gate = *static_cast(data); - gate._entered.store(true, std::memory_order_release); - gate._entered.notify_all(); - gate._released.wait(false, std::memory_order_acquire); - gate._exited.store(true, std::memory_order_release); - gate._exited.notify_all(); - } - - bool _enqueued{false}; - mutable std::atomic _entered{false}; - std::atomic _released{false}; - std::atomic _exited{false}; -}; - -class scoped_stream_gate_release { - public: - explicit scoped_stream_gate_release(cuda_stream_gate& gate) : _gate(gate) {} - ~scoped_stream_gate_release() { _gate.release(); } - - scoped_stream_gate_release(scoped_stream_gate_release const&) = delete; - scoped_stream_gate_release& operator=(scoped_stream_gate_release const&) = delete; - - private: - cuda_stream_gate& _gate; -}; - -} // namespace - TEST_CASE("gpu_table_representation clone creates independent copy", "[gpu_data_representation]") { auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 0); @@ -727,9 +662,10 @@ TEST_CASE("gpu_table_representation clone creates independent copy", "[gpu_data_ *gpu_space, ::cuda::stream_ref{cudaStream_t{cudaStreamDefault}}); - // Clone the representation auto cloned_base = repr.clone(::cuda::stream_ref{cudaStream_t{cudaStreamDefault}}); REQUIRE(cloned_base != nullptr); + REQUIRE(repr.get_writer_event() != nullptr); + REQUIRE(cloned_base->get_writer_event() != nullptr); // Verify it's a gpu_table_representation auto* cloned = dynamic_cast(cloned_base.get()); @@ -756,84 +692,6 @@ TEST_CASE("gpu_table_representation clone creates independent copy", "[gpu_data_ } } -TEST_CASE("gpu_table_representation clone waits for a distinct writer stream", - "[gpu_data_representation][stream_ordering]") -{ - using namespace std::chrono_literals; - - memory::gpu_memory_space_config config; - config.device_id = 0; - config.memory_capacity = 64ULL << 20; - config.mr_factory_fn = test::make_shared_current_device_resource; - auto gpu_space = std::make_shared(config); - CUCASCADE_CUDA_TRY(cudaSetDevice(config.device_id)); - rmm::cuda_stream producer_stream; - rmm::cuda_stream consumer_stream; - - constexpr cudf::size_type num_rows = 1024; - constexpr std::size_t data_size = static_cast(num_rows) * sizeof(std::int32_t); - constexpr unsigned char expected_byte = 0x5a; - - auto column = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, - num_rows, - cudf::mask_state::UNALLOCATED, - consumer_stream.view(), - gpu_space->get_default_allocator()); - - // Establish a known stale value before deliberately blocking the real producer write. - CUCASCADE_CUDA_TRY( - cudaMemsetAsync(column->mutable_view().head(), 0, data_size, consumer_stream.value())); - consumer_stream.synchronize(); - - cuda_stream_gate producer_gate; - producer_gate.enqueue(producer_stream.view()); - CUCASCADE_CUDA_TRY(cudaMemsetAsync( - column->mutable_view().head(), expected_byte, data_size, producer_stream.value())); - - std::vector> columns; - columns.push_back(std::move(column)); - gpu_table_representation source( - std::make_unique(std::move(columns)), *gpu_space, producer_stream.view()); - - producer_gate.wait_until_entered(); - auto const writer_status_while_blocked = cudaEventQuery(source.get_writer_event()); - - std::atomic clone_started{false}; - std::future> clone_future; - scoped_stream_gate_release release_on_exit{producer_gate}; - clone_future = std::async(std::launch::async, [&] { - CUCASCADE_CUDA_TRY(cudaSetDevice(source.get_device_id())); - clone_started.store(true, std::memory_order_release); - clone_started.notify_all(); - return source.clone(consumer_stream.view()); - }); - clone_started.wait(false, std::memory_order_acquire); - - // clone() must not return until the source writer event can complete. - auto const status_while_writer_blocked = clone_future.wait_for(1s); - if (status_while_writer_blocked == std::future_status::ready) { - // Complete a premature copy before releasing the producer to make stale data deterministic. - consumer_stream.synchronize(); - } - - producer_gate.release(); - auto cloned_base = clone_future.get(); - producer_stream.synchronize(); - consumer_stream.synchronize(); - - REQUIRE(writer_status_while_blocked == cudaErrorNotReady); - CHECK(status_while_writer_blocked == std::future_status::timeout); - - auto* clone = dynamic_cast(cloned_base.get()); - REQUIRE(clone != nullptr); - - std::vector bytes(data_size); - CUCASCADE_CUDA_TRY(cudaMemcpy( - bytes.data(), clone->get_table_view().column(0).head(), data_size, cudaMemcpyDeviceToHost)); - REQUIRE(std::all_of( - bytes.cbegin(), bytes.cend(), [](uint8_t value) { return value == expected_byte; })); -} - TEST_CASE("gpu_table_representation clone empty table", "[gpu_data_representation]") { auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 0);