Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 5 additions & 3 deletions src/main/cpp/benchmarks/bloom_filter.cu
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@

#include <cudf_test/column_utilities.hpp>

#include <cuda/stream>

#include <bloom_filter.hpp>
#include <hash/hash.hpp>
#include <nvbench/nvbench.cuh>
Expand All @@ -39,13 +41,13 @@ void bloom_filter_put_impl(nvbench::state& state, int version)
auto const src = create_random_table({{cudf::type_id::INT64}}, row_count{num_rows}, builder);
auto const input = spark_rapids_jni::xxhash64(*src);

auto const stream = cudf::get_default_stream();
state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value()));
cuda::stream_ref const stream = cudf::get_default_stream();
state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.get()));
state.exec(nvbench::exec_tag::timer | nvbench::exec_tag::sync,
[&](nvbench::launch&, auto& timer) {
timer.start();
spark_rapids_jni::bloom_filter_put(*bloom_filter, *input);
stream.synchronize();
stream.sync();
timer.stop();
});

Expand Down
10 changes: 6 additions & 4 deletions src/main/cpp/benchmarks/cast_long_to_binary_string.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,19 @@

#include <cudf/io/types.hpp>

#include <cuda/stream>

#include <cast_string.hpp>
#include <nvbench/nvbench.cuh>

static void long_to_binary_string(nvbench::state& state)
{
auto const num_rows = static_cast<cudf::size_type>(state.get_int64("num_rows"));

auto const input_table = create_random_table({cudf::type_id::INT64}, row_count{num_rows});
auto const long_col = input_table->get_column(0);
auto const stream = cudf::get_default_stream();
state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value()));
auto const input_table = create_random_table({cudf::type_id::INT64}, row_count{num_rows});
auto const long_col = input_table->get_column(0);
cuda::stream_ref const stream = cudf::get_default_stream();
state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.get()));
state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) {
spark_rapids_jni::long_to_binary_string(long_col, stream);
});
Expand Down
4 changes: 2 additions & 2 deletions src/main/cpp/benchmarks/get_json_object.cu
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ struct strings_to_host_fn {
void operator()(std::vector<std::string>& host_data,
char const* chars,
cudf::column_view const& offsets,
rmm::cuda_stream_view stream)
cuda::stream_ref stream)
{
auto const h_offsets = cudf::detail::make_std_vector_sync(
cudf::device_span<OffsetType const>(offsets.data<OffsetType>(), offsets.size()), stream);
Expand All @@ -62,7 +62,7 @@ struct strings_to_host_fn {
void operator()(std::vector<std::string>&,
char const*,
cudf::column_view const&,
rmm::cuda_stream_view)
cuda::stream_ref)
{
CUDF_FAIL("invalid offsets type");
}
Expand Down
4 changes: 2 additions & 2 deletions src/main/cpp/src/aggregation64_utils.cu
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ namespace spark_rapids_jni {
std::unique_ptr<cudf::column> extract_chunk32_from_64bit(cudf::column_view const& in_col,
cudf::data_type type,
int chunk_idx,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
CUDF_EXPECTS(
Expand Down Expand Up @@ -112,7 +112,7 @@ std::unique_ptr<cudf::column> extract_chunk32_from_64bit(cudf::column_view const
// Reassemble a column of 64-bit values from two 64-bit integer columns with overflow detection.
std::unique_ptr<cudf::table> assemble64_from_sum(cudf::table_view const& chunks_table,
cudf::data_type output_type,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
CUDF_EXPECTS(
Expand Down
8 changes: 4 additions & 4 deletions src/main/cpp/src/aggregation64_utils.hpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2025, NVIDIA CORPORATION.
* Copyright (c) 2025-2026, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand All @@ -20,7 +20,7 @@
#include <cudf/table/table.hpp>
#include <cudf/utilities/default_stream.hpp>

#include <rmm/cuda_stream_view.hpp>
#include <cuda/stream>

#include <memory>

Expand All @@ -46,7 +46,7 @@ std::unique_ptr<cudf::column> extract_chunk32_from_64bit(
cudf::column_view const& col,
cudf::data_type dtype,
int chunk_idx,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
cuda::stream_ref stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());

/**
Expand All @@ -70,7 +70,7 @@ std::unique_ptr<cudf::column> extract_chunk32_from_64bit(
std::unique_ptr<cudf::table> assemble64_from_sum(
cudf::table_view const& chunks_table,
cudf::data_type output_type,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
cuda::stream_ref stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());

} // namespace spark_rapids_jni
32 changes: 16 additions & 16 deletions src/main/cpp/src/bloom_filter.cu
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,12 @@
#include <cudf/utilities/memory_resource.hpp>
#include <cudf/utilities/span.hpp>

#include <rmm/cuda_stream_view.hpp>
#include <rmm/exec_policy.hpp>

#include <cuda/atomic>
#include <cuda/functional>
#include <cuda/std/utility>
#include <cuda/stream>
#include <thrust/logical.h>

#include <byteswap.h>
Expand Down Expand Up @@ -154,22 +154,22 @@ struct bloom_probe_functor {

void pack_bloom_filter_header(cudf::device_span<uint8_t> buf,
bloom_filter_header const& header,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
int32_t seed)
{
if (header.version == bloom_filter_version_1) {
bloom_filter_header_v1 raw = {byte_swap_int32(header.version),
byte_swap_int32(header.num_hashes),
byte_swap_int32(header.num_longs)};
CUDF_CUDA_TRY(cudaMemcpyAsync(
buf.data(), &raw, bloom_filter_header_v1_size_bytes, cudaMemcpyDefault, stream));
buf.data(), &raw, bloom_filter_header_v1_size_bytes, cudaMemcpyDefault, stream.get()));
} else {
bloom_filter_header_v2 raw = {byte_swap_int32(header.version),
byte_swap_int32(header.num_hashes),
byte_swap_int32(seed),
byte_swap_int32(header.num_longs)};
CUDF_CUDA_TRY(cudaMemcpyAsync(
buf.data(), &raw, bloom_filter_header_v2_size_bytes, cudaMemcpyDefault, stream));
buf.data(), &raw, bloom_filter_header_v2_size_bytes, cudaMemcpyDefault, stream.get()));
}
}

Expand All @@ -188,7 +188,7 @@ void pack_bloom_filter_header(cudf::device_span<uint8_t> buf,
for V2, the value stored in the serialized header.
*/
std::tuple<bloom_filter_header, cudf::device_span<cudf::bitmask_type const>, int64_t, int32_t>
unpack_bloom_filter(cudf::device_span<uint8_t const> bloom_filter, rmm::cuda_stream_view stream)
unpack_bloom_filter(cudf::device_span<uint8_t const> bloom_filter, cuda::stream_ref stream)
{
CUDF_EXPECTS(bloom_filter.size() >= static_cast<size_t>(bloom_filter_header_v1_size_bytes),
"Encountered truncated bloom filter");
Expand All @@ -200,8 +200,8 @@ unpack_bloom_filter(cudf::device_span<uint8_t const> bloom_filter, rmm::cuda_str
// TODO (future): Consider using pinned host memory for cudaMemcpyAsync.
// Refer to https://github.com/NVIDIA/spark-rapids-jni/issues/4407.
CUDF_CUDA_TRY(
cudaMemcpyAsync(raw_ints, bloom_filter.data(), read_size, cudaMemcpyDefault, stream));
stream.synchronize();
cudaMemcpyAsync(raw_ints, bloom_filter.data(), read_size, cudaMemcpyDefault, stream.get()));
stream.sync();

int const version = byte_swap_int32(raw_ints[0]);
CUDF_EXPECTS(version == bloom_filter_version_1 || version == bloom_filter_version_2,
Expand Down Expand Up @@ -237,7 +237,7 @@ unpack_bloom_filter(cudf::device_span<uint8_t const> bloom_filter, rmm::cuda_str
}

std::tuple<bloom_filter_header, cudf::device_span<cudf::bitmask_type const>, int64_t, int32_t>
unpack_bloom_filter(cudf::column_view const& bloom_filter, rmm::cuda_stream_view stream)
unpack_bloom_filter(cudf::column_view const& bloom_filter, cuda::stream_ref stream)
{
return unpack_bloom_filter(
cudf::device_span<uint8_t const>{bloom_filter.data<uint8_t>(),
Expand Down Expand Up @@ -300,7 +300,7 @@ std::unique_ptr<cudf::list_scalar> bloom_filter_create(int version,
int num_hashes,
int bloom_filter_longs,
int seed,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
SRJ_FUNC_RANGE();
Expand All @@ -320,8 +320,8 @@ std::unique_ptr<cudf::list_scalar> bloom_filter_create(int version,
stream,
(version == bloom_filter_version_1 ? 0 : seed));

CUDF_CUDA_TRY(
cudaMemsetAsync(static_cast<uint8_t*>(buf.data()) + hdr_size, 0, bloom_filter_size, stream));
CUDF_CUDA_TRY(cudaMemsetAsync(
static_cast<uint8_t*>(buf.data()) + hdr_size, 0, bloom_filter_size, stream.get()));

return std::make_unique<cudf::list_scalar>(
cudf::column(
Expand All @@ -333,7 +333,7 @@ std::unique_ptr<cudf::list_scalar> bloom_filter_create(int version,

void bloom_filter_put(cudf::list_scalar& bloom_filter,
cudf::column_view const& input,
rmm::cuda_stream_view stream)
cuda::stream_ref stream)
{
SRJ_FUNC_RANGE();
auto [header, buffer, bloom_filter_bits, seed] = unpack_bloom_filter(bloom_filter.view(), stream);
Expand All @@ -351,7 +351,7 @@ void bloom_filter_put(cudf::list_scalar& bloom_filter,

auto launch = [&](auto version_tag, auto nullable_tag) {
gpu_bloom_filter_put<decltype(version_tag)::value, decltype(nullable_tag)::value>
<<<grid.num_blocks, block_size, 0, stream.value()>>>(
<<<grid.num_blocks, block_size, 0, stream.get()>>>(
mutable_buffer, bloom_filter_bits, *d_input, header.num_hashes, seed);
};

Expand All @@ -373,7 +373,7 @@ void bloom_filter_put(cudf::list_scalar& bloom_filter,
}

std::unique_ptr<cudf::list_scalar> bloom_filter_merge(cudf::column_view const& bloom_filters,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
SRJ_FUNC_RANGE();
Expand Down Expand Up @@ -450,7 +450,7 @@ std::unique_ptr<cudf::list_scalar> bloom_filter_merge(cudf::column_view const& b

std::unique_ptr<cudf::column> bloom_filter_probe(cudf::column_view const& input,
cudf::device_span<uint8_t const> bloom_filter,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
SRJ_FUNC_RANGE();
Expand Down Expand Up @@ -489,7 +489,7 @@ std::unique_ptr<cudf::column> bloom_filter_probe(cudf::column_view const& input,

std::unique_ptr<cudf::column> bloom_filter_probe(cudf::column_view const& input,
cudf::list_scalar& bloom_filter,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
SRJ_FUNC_RANGE();
Expand Down
13 changes: 7 additions & 6 deletions src/main/cpp/src/bloom_filter.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,10 @@
#include <cudf/utilities/default_stream.hpp>
#include <cudf/utilities/span.hpp>

#include <rmm/cuda_stream_view.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda/stream>

namespace spark_rapids_jni {

constexpr int bloom_filter_version_1 = 1;
Expand Down Expand Up @@ -90,7 +91,7 @@ std::unique_ptr<cudf::list_scalar> bloom_filter_create(
int num_hashes,
int bloom_filter_longs,
int seed = 0,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
cuda::stream_ref stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref());

/**
Expand All @@ -105,7 +106,7 @@ std::unique_ptr<cudf::list_scalar> bloom_filter_create(
*/
void bloom_filter_put(cudf::list_scalar& bloom_filter,
cudf::column_view const& input,
rmm::cuda_stream_view stream = cudf::get_default_stream());
cuda::stream_ref stream = cudf::get_default_stream());

/**
* @brief Probe a bloom filter with an input column of int64_t values.
Expand All @@ -121,7 +122,7 @@ void bloom_filter_put(cudf::list_scalar& bloom_filter,
std::unique_ptr<cudf::column> bloom_filter_probe(
cudf::column_view const& input,
cudf::device_span<uint8_t const> bloom_filter,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
cuda::stream_ref stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref());

/**
Expand All @@ -138,7 +139,7 @@ std::unique_ptr<cudf::column> bloom_filter_probe(
std::unique_ptr<cudf::column> bloom_filter_probe(
cudf::column_view const& input,
cudf::list_scalar& bloom_filter,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
cuda::stream_ref stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref());

/**
Expand All @@ -156,7 +157,7 @@ std::unique_ptr<cudf::column> bloom_filter_probe(
*/
std::unique_ptr<cudf::list_scalar> bloom_filter_merge(
cudf::column_view const& bloom_filters,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
cuda::stream_ref stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref());

} // namespace spark_rapids_jni
4 changes: 2 additions & 2 deletions src/main/cpp/src/case_when.cu
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ struct select_first_true_fn {
} // anonymous namespace

std::unique_ptr<cudf::column> select_first_true_index(cudf::table_view const& when_bool_columns,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
// checks
Expand Down Expand Up @@ -95,7 +95,7 @@ std::unique_ptr<cudf::column> select_first_true_index(cudf::table_view const& wh
} // namespace detail

std::unique_ptr<cudf::column> select_first_true_index(cudf::table_view const& when_bool_columns,
rmm::cuda_stream_view stream,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
return detail::select_first_true_index(when_bool_columns, stream, mr);
Expand Down
6 changes: 4 additions & 2 deletions src/main/cpp/src/case_when.hpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2024-2025, NVIDIA CORPORATION.
* Copyright (c) 2024-2026, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand All @@ -22,6 +22,8 @@
#include <rmm/mr/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda/stream>

#include <memory>

namespace spark_rapids_jni {
Expand All @@ -47,7 +49,7 @@ namespace spark_rapids_jni {
*/
std::unique_ptr<cudf::column> select_first_true_index(
cudf::table_view const& when_bool_columns,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
cuda::stream_ref stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref());

} // namespace spark_rapids_jni
Loading