Skip to content

fix(data): order GPU representation reads after writers #182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kevkrist wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix(data): order GPU representation reads after writers #182

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
605fe3f
fix(cudf): order GPU reads after writers
kevkrist Aug 12, 2026
c5bac4c
test(cudf): cover cross-stream clone ordering
kevkrist Aug 13, 2026
160ca5f
docs(cudf): streamline ordering comments
kevkrist Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 15 additions & 9 deletions include/cucascade/cudf/gpu_data_representation.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<idata_representation> 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<idata_representation> clone(rmm::cuda_stream_view stream) override;

Expand All @@ -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<cudf::table> 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<cudf::table> release_table(rmm::cuda_stream_view stream);

Expand Down
44 changes: 33 additions & 11 deletions src/cudf/gpu_data_representation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
#include <cudf/copying.hpp>
#include <cudf/utilities/traits.hpp>

#include <rmm/cuda_device.hpp>

namespace cucascade {

gpu_table_representation::gpu_table_representation(std::unique_ptr<cudf::table> table,
Expand Down Expand Up @@ -73,11 +75,24 @@ cudf::table_view gpu_table_representation::get_table_view() const
}
}

std::unique_ptr<cudf::table> gpu_table_representation::release_table(
[[maybe_unused]] rmm::cuda_stream_view stream)
std::unique_ptr<cudf::table> gpu_table_representation::release_table(rmm::cuda_stream_view stream)
{
if (std::holds_alternative<owning_table_view>(_table)) {
_table = std::make_unique<cudf::table>(std::get<owning_table_view>(_table).view, stream);
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);
} else {
CUCASCADE_CUDA_TRY(cudaDeviceSynchronize());
}

auto materialized = std::make_unique<cudf::table>(
std::get<owning_table_view>(_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<std::unique_ptr<cudf::table>>(_table));
}
Expand All @@ -102,14 +117,21 @@ void gpu_table_representation::rebind_stream(rmm::cuda_stream_view stream)

std::unique_ptr<idata_representation> 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<gpu_table_representation>(
std::make_unique<cudf::table>(view, stream), get_memory_space(), stream);
return cloned;
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);
} else {
CUCASCADE_CUDA_TRY(cudaDeviceSynchronize());
}

auto cloned_table = std::make_unique<cudf::table>(
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<gpu_table_representation>(
std::move(cloned_table), get_memory_space(), stream);
}

void gpu_table_representation::record_writer_event(rmm::cuda_stream_view writer_stream)
Expand Down
73 changes: 27 additions & 46 deletions src/cudf/representation_converter_builtins.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
#include <cudf/utilities/traits.hpp>
#include <cudf/utilities/type_dispatcher.hpp>

#include <rmm/cuda_device.hpp>
#include <rmm/cuda_stream.hpp>
#include <rmm/device_buffer.hpp>
#include <rmm/device_uvector.hpp>
Expand Down Expand Up @@ -81,6 +82,22 @@ inline cudf::type_id as_cudf_type_id(int32_t type_id)
return static_cast<cudf::type_id>(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
Expand All @@ -101,11 +118,9 @@ std::unique_ptr<idata_representation> 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<gpu_table_representation>();
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<memory::fixed_size_host_memory_resource>();
Expand Down Expand Up @@ -497,7 +512,9 @@ std::unique_ptr<idata_representation> convert_gpu_to_host_fast(
rmm::cuda_stream_view stream,
memory::reservation* reservation)
{
auto& gpu_source = source.cast<gpu_table_representation>();
auto& gpu_source = source.cast<gpu_table_representation>();
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 ---
Expand Down Expand Up @@ -866,11 +883,6 @@ std::unique_ptr<idata_representation> 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<gpu_table_representation>();

// Same-device case: clone via source's own clone() method.
Expand All @@ -881,47 +893,14 @@ std::unique_ptr<idata_representation> 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
// and target-side allocations are issued on this stream so they observe in-order
// 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();

Expand Down Expand Up @@ -1614,8 +1593,10 @@ static std::unique_ptr<idata_representation> 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<gpu_table_representation>();
auto& backend = target_memory_space->get_io_backend();
auto& gpu_source = source.cast<gpu_table_representation>();
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
Expand Down
Loading