diff --git a/include/cucascade/cudf/gpu_data_representation.hpp b/include/cucascade/cudf/gpu_data_representation.hpp index d7e8154..9d2fc8c 100644 --- a/include/cucascade/cudf/gpu_data_representation.hpp +++ b/include/cucascade/cudf/gpu_data_representation.hpp @@ -53,9 +53,11 @@ class gpu_table_representation : public idata_representation { * 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). + * writer_stream) automatically, including when it is the CUDA default stream. + * + * @pre The calling thread's current CUDA device, @p writer_stream, @p memory_space, and every + * allocation in @p table refer to the same device. + * @pre @p writer_stream is the stream on which writes to @p table were last enqueued. * * @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 @@ -74,6 +76,10 @@ class gpu_table_representation : public idata_representation { * underlying data of @p table_view was last written. See the simple-table ctor * docstring for the writer_stream contract. * + * @pre The calling thread's current CUDA device, @p writer_stream, @p memory_space, and every + * allocation viewed by @p table_view refer to the same device. + * @pre @p writer_stream is the stream on which writes to @p table_view were last enqueued. + * * @param table_view View of the cuDF table (data ownership lives in @p owner) * @tparam Owner The type of the owner of the cuDF table (e.g., a specific operator or component) * @param owner Owner of the underlying data (transferred via std::any storage) @@ -121,6 +127,13 @@ class gpu_table_representation : public idata_representation { * * The cloned representation will have its own copy of the underlying cuDF table, * residing in the same memory space as the original. + * The copy waits for this representation's writer event and records its own writer event on + * @p stream. This method returns without synchronizing @p stream. + * + * @pre The calling thread's current CUDA device, @p stream, this representation's memory space, + * and its allocations refer to the same device. + * @pre This representation, its allocations, and their contents remain alive and unmodified + * until the work enqueued on @p stream completes. * * @param stream CUDA stream for memory operations * @return std::unique_ptr A new gpu_table_representation with copied data @@ -138,13 +151,22 @@ class gpu_table_representation : public idata_representation { * @brief Release ownership of the underlying cuDF table * * After calling this method, this representation no longer owns the table. + * The returned table is ordered after this representation's writer event on @p stream. For a + * view-backed representation, materialization is also enqueued on @p stream. This method returns + * without synchronizing @p stream. * - * @pre No stream other than @p stream may have in-flight work touching the table's device + * @pre Except for work represented by the writer event, 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 The calling thread's current CUDA device must match get_device_id(). + * * @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. * + * @pre For a view-backed representation, the external owner must keep the viewed allocations + * alive and unmodified until the work enqueued on @p stream completes. + * * @param stream Stream that will own deallocation ordering of the returned table's buffers * (also used to materialize the table from a view path before release) * @return std::unique_ptr The cuDF table @@ -189,8 +211,12 @@ class gpu_table_representation : public idata_representation { * 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. + * representation owns a single writer event handle that is reused). The CUDA default stream is + * a valid writer stream and records an event on the calling thread's current device. + * + * @pre The calling thread's current CUDA device, @p writer_stream, this representation's memory + * space, and its allocations refer to the same device. + * @pre @p writer_stream is the stream on which writes to this representation were last enqueued. * * @param writer_stream The stream on which the most recent writes to this * representation's memory were enqueued. @@ -198,16 +224,13 @@ class gpu_table_representation : public idata_representation { void record_writer_event(::cuda::stream_ref writer_stream) override; /** - * @brief Get the writer event recorded by record_writer_event(), or nullptr if none. + * @brief Get the writer event recorded by record_writer_event(). * * Readers that cross stream / device boundaries must call cudaStreamWaitEvent on - * this event (when non-null) before reading the underlying memory. When this - * returns nullptr, callers should fall back to a coarser sync (e.g. - * cudaDeviceSynchronize on the source device) — this is the legacy behavior - * preserved for representations constructed by code paths that have not yet - * been migrated to record writer events. + * this event before reading the underlying memory. Every successfully constructed + * gpu_table_representation has a non-null writer event. * - * @return cudaEvent_t The writer event, or nullptr if none has been recorded. + * @return cudaEvent_t The writer event */ [[nodiscard]] cudaEvent_t get_writer_event() const override; @@ -221,9 +244,9 @@ 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 the completion of the most recent writer-stream work that produced this + /// representation. Created and recorded during construction, then reused by + /// record_writer_event(). cudaEvent_t _writer_event{nullptr}; }; @@ -239,7 +262,7 @@ gpu_table_representation::gpu_table_representation(cudf::table_view 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/src/cudf/gpu_data_representation.cpp b/src/cudf/gpu_data_representation.cpp index 8adab21..6fba920 100644 --- a/src/cudf/gpu_data_representation.cpp +++ b/src/cudf/gpu_data_representation.cpp @@ -58,13 +58,8 @@ gpu_table_representation::gpu_table_representation(std::unique_ptr : 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); } + // representation is born with a recorded event, including on the CUDA default stream. + record_writer_event(writer_stream); } gpu_table_representation::~gpu_table_representation() @@ -105,11 +100,13 @@ std::unique_ptr gpu_table_representation::release_table(::cuda::str if (std::holds_alternative(_table)) { // The deep copy below is enqueued on `stream`, and its buffers are bound to it. validate_stream_device(stream, get_device_id()); + cucascade::cuda::cuda_event_view{_writer_event}.wait(stream); _table = std::make_unique(std::get(_table).view, stream); } 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. gpu_table_representation::rebind_stream(stream); + cucascade::cuda::cuda_event_view{_writer_event}.wait(stream); } return std::move(std::get>(_table)); } @@ -138,6 +135,7 @@ void gpu_table_representation::rebind_stream(::cuda::stream_ref stream) std::unique_ptr gpu_table_representation::clone(::cuda::stream_ref stream) { + cucascade::cuda::cuda_event_view{_writer_event}.wait(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 diff --git a/test/cudf/test_release_table_stream.cpp b/test/cudf/test_release_table_stream.cpp index f54ae19..a3ec360 100644 --- a/test/cudf/test_release_table_stream.cpp +++ b/test/cudf/test_release_table_stream.cpp @@ -18,6 +18,7 @@ #include "utils/cudf_test_utils.hpp" #include "utils/mock_test_utils.hpp" +#include #include #include #include @@ -43,9 +44,12 @@ #include +#include +#include #include #include #include +#include #include #include #include @@ -205,6 +209,88 @@ void CUDART_CB stall_stream_callback(void* /*user_data*/) std::this_thread::sleep_for(std::chrono::milliseconds(40)); } +struct writer_event_gate { + std::atomic entered{false}; + std::atomic released{false}; + + void release() noexcept + { + released.store(true, std::memory_order_release); + released.notify_all(); + } +}; + +void CUDART_CB wait_on_writer_event_gate(void* user_data) +{ + auto* gate = static_cast(user_data); + gate->entered.store(true, std::memory_order_release); + gate->entered.notify_all(); + gate->released.wait(false, std::memory_order_acquire); +} + +bool wait_until_set(std::atomic const& value, std::chrono::steady_clock::duration timeout) +{ + auto const deadline = std::chrono::steady_clock::now() + timeout; + while (!value.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + return value.load(std::memory_order_acquire); +} + +cudaError_t wait_until_stream_pending(::cuda::stream_ref stream, + std::chrono::steady_clock::duration timeout) +{ + auto const deadline = std::chrono::steady_clock::now() + timeout; + cudaError_t status = cudaSuccess; + do { + status = cudaStreamQuery(stream.get()); + if (status != cudaSuccess) { return status; } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } while (std::chrono::steady_clock::now() < deadline); + return status; +} + +cudaError_t observe_event_pending_for(cudaEvent_t event, + std::chrono::steady_clock::duration duration) +{ + auto const deadline = std::chrono::steady_clock::now() + duration; + do { + auto const status = cudaEventQuery(event); + if (status != cudaErrorNotReady) { return status; } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } while (std::chrono::steady_clock::now() < deadline); + return cudaErrorNotReady; +} + +class writer_event_gate_cleanup { + public: + writer_event_gate_cleanup(writer_event_gate& gate, ::cuda::stream_ref stream) + : _gate(gate), _stream(stream) + { + } + + ~writer_event_gate_cleanup() noexcept { drain(); } + + writer_event_gate_cleanup(writer_event_gate_cleanup const&) = delete; + writer_event_gate_cleanup& operator=(writer_event_gate_cleanup const&) = delete; + + cudaError_t drain() noexcept + { + if (!_drained) { + _gate.release(); + _status = cudaStreamSynchronize(_stream.get()); + _drained = true; + } + return _status; + } + + private: + writer_event_gate& _gate; + ::cuda::stream_ref _stream; + cudaError_t _status{cudaSuccess}; + bool _drained{false}; +}; + /// Pinned so the D2H readback enqueues asynchronously instead of staging synchronously. struct pinned_buffer { void* ptr{nullptr}; @@ -219,6 +305,111 @@ struct pinned_buffer { } // namespace +TEST_CASE("gpu_table_representation clone waits for pending writer work without host blocking", + "[gpu_data_representation][clone][writer_event]") +{ + using namespace std::chrono_literals; + + CUCASCADE_CUDA_TRY(cudaSetDevice(0)); + auto& gpu_space = shared_gpu_space(); + rmm::cuda_stream writer_stream; + rmm::cuda_stream clone_stream; + + constexpr cudf::size_type num_rows = 1 << 18; + constexpr std::size_t num_bytes = static_cast(num_rows) * sizeof(int32_t); + constexpr unsigned char stale_pattern = 0x11; + constexpr unsigned char written_pattern = 0x5A; + + auto column = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, + num_rows, + cudf::mask_state::UNALLOCATED, + writer_stream.view(), + gpu_space->get_default_allocator()); + CUCASCADE_CUDA_TRY(cudaMemsetAsync( + column->mutable_view().head(), stale_pattern, num_bytes, writer_stream.value())); + writer_stream.synchronize(); + + std::vector> columns; + columns.push_back(std::move(column)); + auto source = std::make_shared( + std::make_unique(std::move(columns)), *gpu_space, writer_stream.view()); + + auto const consumer_initial_status = cudaStreamQuery(clone_stream.value()); + writer_event_gate gate; + std::future> clone_future; + CUCASCADE_CUDA_TRY(cudaLaunchHostFunc(writer_stream.value(), wait_on_writer_event_gate, &gate)); + writer_event_gate_cleanup gate_cleanup{gate, writer_stream.view()}; + + bool const gate_entered = wait_until_set(gate.entered, 5s); + if (!gate_entered) { + auto const cleanup_status = gate_cleanup.drain(); + REQUIRE(cleanup_status == cudaSuccess); + REQUIRE(gate_entered); + return; + } + + CUCASCADE_CUDA_TRY(cudaMemsetAsync(const_cast(source->get_table_view().column(0).head()), + written_pattern, + num_bytes, + writer_stream.value())); + source->record_writer_event(writer_stream.view()); + + clone_future = std::async(std::launch::async, [&] { + CUCASCADE_CUDA_TRY(cudaSetDevice(0)); + return source->clone(clone_stream.view()); + }); + + auto const consumer_pending_status = wait_until_stream_pending(clone_stream.view(), 5s); + bool const returned_while_gated = clone_future.wait_for(5s) == std::future_status::ready; + + std::unique_ptr cloned_base; + std::exception_ptr clone_error; + auto consume_clone_result = [&] { + try { + cloned_base = clone_future.get(); + } catch (...) { + clone_error = std::current_exception(); + } + }; + + if (returned_while_gated) { consume_clone_result(); } + + auto* cloned = dynamic_cast(cloned_base.get()); + auto const clone_writer_status = cloned != nullptr && cloned->get_writer_event() != nullptr + ? observe_event_pending_for(cloned->get_writer_event(), 100ms) + : cudaErrorInvalidValue; + + auto const writer_cleanup_status = gate_cleanup.drain(); + if (!returned_while_gated) { + clone_future.wait(); + consume_clone_result(); + } + auto const clone_sync_status = cudaStreamSynchronize(clone_stream.value()); + + std::vector actual(num_bytes); + auto const readback_status = cloned != nullptr + ? cudaMemcpy(actual.data(), + cloned->get_table_view().column(0).head(), + actual.size(), + cudaMemcpyDeviceToHost) + : cudaErrorInvalidValue; + bool const copied_post_gate_bytes = + readback_status == cudaSuccess && + std::all_of( + actual.cbegin(), actual.cend(), [](unsigned char value) { return value == written_pattern; }); + + REQUIRE(consumer_initial_status == cudaSuccess); + REQUIRE(consumer_pending_status == cudaErrorNotReady); + REQUIRE(returned_while_gated); + REQUIRE(clone_error == nullptr); + REQUIRE(clone_writer_status == cudaErrorNotReady); + REQUIRE(writer_cleanup_status == cudaSuccess); + REQUIRE(clone_sync_status == cudaSuccess); + REQUIRE(cloned != nullptr); + REQUIRE(readback_status == cudaSuccess); + REQUIRE(copied_post_gate_bytes); +} + TEST_CASE("release_table rebinds owned-table buffers to the release stream", "[release_table][stream]") { @@ -375,6 +566,143 @@ TEST_CASE("view-branch release_table deep-copies on the release stream and leave REQUIRE(src_checked >= 5); } +TEST_CASE("release_table waits for pending writer work without host blocking", + "[release_table][writer_event]") +{ + using namespace std::chrono_literals; + + CUCASCADE_CUDA_TRY(cudaSetDevice(0)); + auto& gpu_space = shared_gpu_space(); + rmm::cuda_stream writer_stream; + rmm::cuda_stream release_stream; + + constexpr cudf::size_type num_rows = 1 << 18; + constexpr std::size_t num_bytes = static_cast(num_rows) * sizeof(int32_t); + constexpr unsigned char stale_pattern = 0x22; + constexpr unsigned char written_pattern = 0x6B; + + auto column = cudf::make_numeric_column(cudf::data_type{cudf::type_id::INT32}, + num_rows, + cudf::mask_state::UNALLOCATED, + writer_stream.view(), + gpu_space->get_default_allocator()); + CUCASCADE_CUDA_TRY(cudaMemsetAsync( + column->mutable_view().head(), stale_pattern, num_bytes, writer_stream.value())); + writer_stream.synchronize(); + + std::vector> columns; + columns.push_back(std::move(column)); + auto table = std::make_unique(std::move(columns)); + + bool view_backed = false; + std::shared_ptr owner; + std::weak_ptr owner_lifetime; + std::unique_ptr rep; + SECTION("owned table") + { + rep = std::make_unique( + std::move(table), *gpu_space, writer_stream.view()); + } + SECTION("view-backed table") + { + view_backed = true; + owner = std::shared_ptr{std::move(table)}; + owner_lifetime = owner; + rep = std::make_unique(owner->view(), + std::shared_ptr{owner}, + owner->alloc_size(), + *gpu_space, + writer_stream.view()); + } + + auto const consumer_initial_status = cudaStreamQuery(release_stream.value()); + writer_event_gate gate; + std::future> release_future; + CUCASCADE_CUDA_TRY(cudaLaunchHostFunc(writer_stream.value(), wait_on_writer_event_gate, &gate)); + writer_event_gate_cleanup gate_cleanup{gate, writer_stream.view()}; + + bool const gate_entered = wait_until_set(gate.entered, 5s); + if (!gate_entered) { + auto const cleanup_status = gate_cleanup.drain(); + REQUIRE(cleanup_status == cudaSuccess); + REQUIRE(gate_entered); + return; + } + + CUCASCADE_CUDA_TRY(cudaMemsetAsync(const_cast(rep->get_table_view().column(0).head()), + written_pattern, + num_bytes, + writer_stream.value())); + rep->record_writer_event(writer_stream.view()); + + release_future = std::async(std::launch::async, [&] { + CUCASCADE_CUDA_TRY(cudaSetDevice(0)); + return rep->release_table(release_stream.view()); + }); + + auto const consumer_pending_status = wait_until_stream_pending(release_stream.view(), 5s); + bool const returned_while_gated = release_future.wait_for(5s) == std::future_status::ready; + + std::unique_ptr released; + std::exception_ptr release_error; + auto consume_release_result = [&] { + try { + released = release_future.get(); + } catch (...) { + release_error = std::current_exception(); + } + }; + + if (returned_while_gated) { consume_release_result(); } + + auto const post_return_stream_status = returned_while_gated && released != nullptr + ? cudaStreamQuery(release_stream.value()) + : cudaErrorInvalidValue; + std::unique_ptr release_tail; + cudaError_t release_tail_status = cudaSuccess; + if (view_backed && returned_while_gated && released != nullptr) { + release_tail = std::make_unique(); + release_tail->record(release_stream.view()); + release_tail_status = observe_event_pending_for(release_tail->get(), 100ms); + } + bool const owner_retained_while_copy_pending = + !view_backed || (!owner_lifetime.expired() && owner.use_count() == 1); + + auto const writer_cleanup_status = gate_cleanup.drain(); + if (!returned_while_gated) { + release_future.wait(); + consume_release_result(); + } + auto const release_sync_status = cudaStreamSynchronize(release_stream.value()); + + std::vector actual(num_bytes); + auto const readback_status = + released != nullptr + ? cudaMemcpy( + actual.data(), released->view().column(0).head(), actual.size(), cudaMemcpyDeviceToHost) + : cudaErrorInvalidValue; + bool const copied_post_gate_bytes = + readback_status == cudaSuccess && + std::all_of( + actual.cbegin(), actual.cend(), [](unsigned char value) { return value == written_pattern; }); + if (view_backed) { owner.reset(); } + bool const owner_released_after_copy = !view_backed || owner_lifetime.expired(); + + REQUIRE(consumer_initial_status == cudaSuccess); + REQUIRE(consumer_pending_status == cudaErrorNotReady); + REQUIRE(returned_while_gated); + REQUIRE(release_error == nullptr); + REQUIRE(released != nullptr); + REQUIRE(post_return_stream_status == cudaErrorNotReady); + if (view_backed) { REQUIRE(release_tail_status == cudaErrorNotReady); } + REQUIRE(owner_retained_while_copy_pending); + REQUIRE(writer_cleanup_status == cudaSuccess); + REQUIRE(release_sync_status == cudaSuccess); + REQUIRE(readback_status == cudaSuccess); + REQUIRE(copied_post_gate_bytes); + REQUIRE(owner_released_after_copy); +} + TEST_CASE("release_table then cudf::rebind_stream to the same stream composes", "[release_table][stream]") { diff --git a/test/data/test_data_representation.cpp b/test/data/test_data_representation.cpp index 12f6208..c7e45de 100644 --- a/test/data/test_data_representation.cpp +++ b/test/data/test_data_representation.cpp @@ -39,6 +39,7 @@ #include #include +#include #include #include @@ -168,6 +169,29 @@ TEST_CASE("gpu_table_representation Construction", "[gpu_data_representation]") REQUIRE(repr.get_size_in_bytes() > 0); } +TEST_CASE("gpu_table_representation constructors record the CUDA default stream", + "[gpu_data_representation][writer_event]") +{ + CUCASCADE_CUDA_TRY(cudaSetDevice(0)); + auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 0); + auto stream = ::cuda::stream_ref{cudaStream_t{cudaStreamDefault}}; + + gpu_table_representation owned( + std::make_unique(std::vector>{}), + *gpu_space, + stream); + + auto view_owner = std::make_shared(std::vector>{}); + gpu_table_representation view_backed(view_owner->view(), + std::shared_ptr{view_owner}, + view_owner->alloc_size(), + *gpu_space, + stream); + + REQUIRE(owned.get_writer_event() != nullptr); + REQUIRE(view_backed.get_writer_event() != nullptr); +} + TEST_CASE("gpu_table_representation get_size_in_bytes", "[gpu_data_representation]") { auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 0); @@ -259,6 +283,7 @@ TEST_CASE("gpu_table_representation device_id", "[gpu_data_representation]") return; } + rmm::cuda_set_device_raii const pin_device{rmm::cuda_device_id{1}}; auto gpu_space = make_mock_memory_space(memory::Tier::GPU, 1); auto table = create_simple_cudf_table(100, gpu_space->get_default_allocator()); gpu_table_representation repr(std::make_unique(std::move(table)),