diff --git a/java/src/main/native/include/maps_column_view.hpp b/java/src/main/native/include/maps_column_view.hpp index 379f5a1e19b3..2cc0f61b087e 100644 --- a/java/src/main/native/include/maps_column_view.hpp +++ b/java/src/main/native/include/maps_column_view.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -10,7 +10,7 @@ #include #include -#include +#include namespace cudf { @@ -29,7 +29,7 @@ namespace jni { class maps_column_view { public: maps_column_view(lists_column_view const& lists_of_structs, - rmm::cuda_stream_view stream = cudf::get_default_stream()); + cuda::stream_ref stream = cudf::get_default_stream()); // Rule of 5. maps_column_view(maps_column_view const& maps_view) = default; @@ -74,7 +74,7 @@ class maps_column_view { */ std::unique_ptr get_values_for( column_view const& keys, - 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()) const; /** @@ -93,7 +93,7 @@ class maps_column_view { */ std::unique_ptr get_values_for( scalar const& key, - 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()) const; /** @@ -114,7 +114,7 @@ class maps_column_view { */ std::unique_ptr contains( scalar const& key, - 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()) const; /** @@ -136,7 +136,7 @@ class maps_column_view { std::unique_ptr contains( column_view const& key, - 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()) const; private: diff --git a/java/src/main/native/src/ColumnViewJni.cu b/java/src/main/native/src/ColumnViewJni.cu index 028e3f5a548d..8491de9da8ac 100644 --- a/java/src/main/native/src/ColumnViewJni.cu +++ b/java/src/main/native/src/ColumnViewJni.cu @@ -33,7 +33,7 @@ namespace cudf::jni { std::unique_ptr generate_list_offsets(cudf::column_view const& list_length, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { CUDF_EXPECTS(list_length.type().id() == cudf::type_id::INT32, "Input column does not have type INT32."); @@ -47,7 +47,7 @@ std::unique_ptr generate_list_offsets(cudf::column_view const& lis auto d_offsets = offsets_view.template begin(); thrust::inclusive_scan(rmm::exec_policy_nosync(stream), begin_iter, end_iter, d_offsets + 1); - CUDF_CUDA_TRY(cudaMemsetAsync(d_offsets, 0, sizeof(int32_t), stream)); + CUDF_CUDA_TRY(cudaMemsetAsync(d_offsets, 0, sizeof(int32_t), stream.get())); return offsets_column; } @@ -73,7 +73,7 @@ __device__ bool list_has_nulls(list_device_view list) void post_process_list_overlap(cudf::column_view const& lhs, cudf::column_view const& rhs, std::unique_ptr const& overlap_result, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { // If both of the input columns do not have nulls, we don't need to do anything here. if (!lists_column_view{lhs}.child().has_nulls() && !lists_column_view{rhs}.child().has_nulls()) { @@ -143,7 +143,7 @@ void post_process_list_overlap(cudf::column_view const& lhs, } std::unique_ptr lists_distinct_by_key(cudf::lists_column_view const& input, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { if (input.is_empty()) { return empty_like(input.parent()); } diff --git a/java/src/main/native/src/ColumnViewJni.hpp b/java/src/main/native/src/ColumnViewJni.hpp index 3201e786a98e..0cb29aade905 100644 --- a/java/src/main/native/src/ColumnViewJni.hpp +++ b/java/src/main/native/src/ColumnViewJni.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,7 +9,7 @@ #include #include -#include +#include namespace cudf::jni { @@ -25,7 +25,7 @@ namespace cudf::jni { * @return The column represents list offsets. */ std::unique_ptr generate_list_offsets( - cudf::column_view const& list_length, rmm::cuda_stream_view stream = cudf::get_default_stream()); + cudf::column_view const& list_length, cuda::stream_ref stream = cudf::get_default_stream()); /** * @brief Perform a special treatment for the results of `cudf::lists::have_overlap` to produce the @@ -48,7 +48,7 @@ std::unique_ptr generate_list_offsets( void post_process_list_overlap(cudf::column_view const& lhs, cudf::column_view const& rhs, std::unique_ptr const& overlap_result, - rmm::cuda_stream_view stream = cudf::get_default_stream()); + cuda::stream_ref stream = cudf::get_default_stream()); /** * @brief Generates lists column by copying elements that are distinct by key from each input list @@ -64,6 +64,6 @@ void post_process_list_overlap(cudf::column_view const& lhs, * @return A new list columns in which the elements in each list are distinct by key. */ std::unique_ptr lists_distinct_by_key(cudf::lists_column_view const& input, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); } // namespace cudf::jni diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index be5dab644425..3ea47e68c483 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -12,7 +12,6 @@ #include #include -#include #include #include #include @@ -28,6 +27,7 @@ #include #include +#include #include #include @@ -937,8 +937,8 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocInternal(JNIEnv* env, { cudf::jni::auto_set_device(env); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); - auto c_stream = rmm::cuda_stream_view(reinterpret_cast(stream)); - void* ret = mr.allocate(c_stream, size); + auto c_stream = cuda::stream_ref(reinterpret_cast(stream)); + void* ret = mr.allocate(c_stream, size); return reinterpret_cast(ret); } JNI_CATCH(env, 0); @@ -952,7 +952,7 @@ Java_ai_rapids_cudf_Rmm_free(JNIEnv* env, jclass clazz, jlong ptr, jlong size, j cudf::jni::auto_set_device(env); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); void* cptr = reinterpret_cast(ptr); - auto c_stream = rmm::cuda_stream_view(reinterpret_cast(stream)); + auto c_stream = cuda::stream_ref(reinterpret_cast(stream)); mr.deallocate(c_stream, cptr, size); } JNI_CATCH(env, ); diff --git a/java/src/main/native/src/aggregation128_utils.cu b/java/src/main/native/src/aggregation128_utils.cu index 2daa0622f2be..37b81d4c16d6 100644 --- a/java/src/main/native/src/aggregation128_utils.cu +++ b/java/src/main/native/src/aggregation128_utils.cu @@ -70,7 +70,7 @@ namespace cudf::jni { std::unique_ptr extract_chunk32(cudf::column_view const& in_col, cudf::data_type type, int chunk_idx, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { CUDF_EXPECTS(in_col.type().id() == cudf::type_id::DECIMAL128, "not a 128-bit type"); CUDF_EXPECTS(chunk_idx >= 0 && chunk_idx < 4, "invalid chunk index"); @@ -95,7 +95,7 @@ std::unique_ptr extract_chunk32(cudf::column_view const& in_col, // Reassemble a column of 128-bit values from four 64-bit integer columns with overflow detection. std::unique_ptr assemble128_from_sum(cudf::table_view const& chunks_table, cudf::data_type output_type, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { CUDF_EXPECTS(output_type.id() == cudf::type_id::DECIMAL128, "not a 128-bit type"); CUDF_EXPECTS(chunks_table.num_columns() == 4, "must be 4 column table"); diff --git a/java/src/main/native/src/aggregation128_utils.hpp b/java/src/main/native/src/aggregation128_utils.hpp index e4fb15f52916..fd6be4da55ef 100644 --- a/java/src/main/native/src/aggregation128_utils.hpp +++ b/java/src/main/native/src/aggregation128_utils.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,7 +9,7 @@ #include #include -#include +#include #include @@ -31,11 +31,10 @@ namespace cudf::jni { * @param stream CUDA stream to use * @return A column containing the extracted 32-bit integer values */ -std::unique_ptr extract_chunk32( - cudf::column_view const& col, - cudf::data_type dtype, - int chunk_idx, - rmm::cuda_stream_view stream = cudf::get_default_stream()); +std::unique_ptr extract_chunk32(cudf::column_view const& col, + cudf::data_type dtype, + int chunk_idx, + cuda::stream_ref stream = cudf::get_default_stream()); /** * @brief Reassemble a 128-bit column from four 64-bit integer columns with overflow detection. @@ -60,6 +59,6 @@ std::unique_ptr extract_chunk32( std::unique_ptr assemble128_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()); } // namespace cudf::jni diff --git a/java/src/main/native/src/check_nvcomp_output_sizes.cu b/java/src/main/native/src/check_nvcomp_output_sizes.cu index 9d4205247500..1a069a264866 100644 --- a/java/src/main/native/src/check_nvcomp_output_sizes.cu +++ b/java/src/main/native/src/check_nvcomp_output_sizes.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "check_nvcomp_output_sizes.hpp" @@ -28,7 +28,7 @@ namespace java { bool check_nvcomp_output_sizes(std::size_t const* dev_uncompressed_sizes, std::size_t const* dev_actual_uncompressed_sizes, std::size_t num_chunks, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { NVTX3_FUNC_RANGE_IN(java_domain); return thrust::equal(rmm::exec_policy_nosync(stream), diff --git a/java/src/main/native/src/check_nvcomp_output_sizes.hpp b/java/src/main/native/src/check_nvcomp_output_sizes.hpp index 5ecfaa566d91..f7fedb028bdd 100644 --- a/java/src/main/native/src/check_nvcomp_output_sizes.hpp +++ b/java/src/main/native/src/check_nvcomp_output_sizes.hpp @@ -1,13 +1,14 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include #include +#include + namespace cudf { namespace java { @@ -18,6 +19,6 @@ namespace java { bool check_nvcomp_output_sizes(std::size_t const* dev_uncompressed_sizes, std::size_t const* dev_actual_uncompressed_sizes, std::size_t num_chunks, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); } // namespace java } // namespace cudf diff --git a/java/src/main/native/src/maps_column_view.cu b/java/src/main/native/src/maps_column_view.cu index e50369aef421..f9935cfb459a 100644 --- a/java/src/main/native/src/maps_column_view.cu +++ b/java/src/main/native/src/maps_column_view.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -29,7 +29,7 @@ column_view make_lists(column_view const& lists_child, lists_column_view const& } // namespace maps_column_view::maps_column_view(lists_column_view const& lists_of_structs, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) : keys_{make_lists(lists_of_structs.child().child(0), lists_of_structs)}, values_{make_lists(lists_of_structs.child().child(1), lists_of_structs)} { @@ -43,7 +43,7 @@ maps_column_view::maps_column_view(lists_column_view const& lists_of_structs, template std::unique_ptr get_values_for_impl(maps_column_view const& maps_view, KeyT const& lookup_keys, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { auto const keys_ = maps_view.keys(); @@ -66,7 +66,7 @@ std::unique_ptr get_values_for_impl(maps_column_view const& maps_view, } std::unique_ptr maps_column_view::get_values_for(column_view const& lookup_keys, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) const { CUDF_EXPECTS(lookup_keys.size() == size(), @@ -76,7 +76,7 @@ std::unique_ptr maps_column_view::get_values_for(column_view const& look } std::unique_ptr maps_column_view::get_values_for(cudf::scalar const& lookup_key, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) const { return get_values_for_impl(*this, lookup_key, stream, mr); @@ -85,7 +85,7 @@ std::unique_ptr maps_column_view::get_values_for(cudf::scalar const& loo template std::unique_ptr contains_impl(maps_column_view const& maps_view, KeyT const& lookup_keys, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { auto const keys = maps_view.keys(); @@ -99,7 +99,7 @@ std::unique_ptr contains_impl(maps_column_view const& maps_view, } std::unique_ptr maps_column_view::contains(column_view const& lookup_keys, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) const { CUDF_EXPECTS(lookup_keys.size() == size(), @@ -109,7 +109,7 @@ std::unique_ptr maps_column_view::contains(column_view const& lookup_key } std::unique_ptr maps_column_view::contains(cudf::scalar const& lookup_key, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) const { return contains_impl(*this, lookup_key, stream, mr);