-
Notifications
You must be signed in to change notification settings - Fork 1.1k
refactor: replace rmm::device_scalar with cudf::detail::device_scalar #23618
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,5 @@ | ||
| /* | ||
| * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. | ||
| * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
|
|
@@ -10,15 +10,22 @@ | |
| #include <cudf/detail/utilities/vector_factories.hpp> | ||
|
|
||
| #include <rmm/cuda_stream_view.hpp> | ||
| #include <rmm/device_scalar.hpp> | ||
| #include <rmm/device_uvector.hpp> | ||
| #include <rmm/resource_ref.hpp> | ||
|
|
||
| #include <type_traits> | ||
| #include <utility> | ||
|
|
||
| namespace CUDF_EXPORT cudf { | ||
| namespace detail { | ||
|
|
||
| template <typename T> | ||
| class device_scalar : public rmm::device_scalar<T> { | ||
| class device_scalar { | ||
| public: | ||
| static_assert(std::is_trivially_copyable_v<T>, | ||
| "cudf::detail::device_scalar<T> requires T to be trivially copyable"); | ||
| using value_type = T; | ||
|
|
||
| #ifdef __CUDACC__ | ||
| #pragma nv_exec_check_disable | ||
| #endif | ||
|
|
@@ -35,48 +42,52 @@ class device_scalar : public rmm::device_scalar<T> { | |
| explicit device_scalar( | ||
| rmm::cuda_stream_view stream, | ||
| rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) | ||
| : rmm::device_scalar<T>(stream, mr), bounce_buffer{make_pinned_vector<T>(1, stream)} | ||
| : _storage{1, stream, std::move(mr)}, bounce_buffer{make_pinned_vector<T>(1, stream)} | ||
| { | ||
| } | ||
|
|
||
| explicit device_scalar( | ||
| T const& initial_value, | ||
| rmm::cuda_stream_view stream, | ||
| rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) | ||
| : rmm::device_scalar<T>(stream, mr), bounce_buffer{make_pinned_vector<T>(1, stream)} | ||
| : _storage{1, stream, std::move(mr)}, bounce_buffer{make_pinned_vector<T>(1, stream)} | ||
| { | ||
| bounce_buffer[0] = initial_value; | ||
| cuda_memcpy_async<T>(device_span<T>{this->data(), 1}, bounce_buffer, stream); | ||
| set_value_async(initial_value, stream); | ||
| } | ||
|
|
||
| device_scalar(device_scalar const& other, | ||
| rmm::cuda_stream_view stream, | ||
| rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) | ||
| : rmm::device_scalar<T>(other, stream, mr), bounce_buffer{make_pinned_vector<T>(1, stream)} | ||
| : _storage{other._storage, stream, mr}, bounce_buffer{make_pinned_vector<T>(1, stream)} | ||
| { | ||
| } | ||
|
|
||
| [[nodiscard]] T value(rmm::cuda_stream_view stream) const | ||
| { | ||
| cuda_memcpy<T>(bounce_buffer, device_span<T const>(this->data(), 1), stream); | ||
| cuda_memcpy<T>(bounce_buffer, device_span<T const>{data(), 1}, stream); | ||
| return std::move(bounce_buffer[0]); | ||
| } | ||
|
|
||
| void set_value_async(T const& value, rmm::cuda_stream_view stream) | ||
| { | ||
| bounce_buffer[0] = value; | ||
| cuda_memcpy_async<T>(device_span<T>(this->data(), 1), bounce_buffer, stream); | ||
| cuda_memcpy_async<T>(device_span<T>{data(), 1}, bounce_buffer, stream); | ||
| } | ||
|
|
||
| void set_value_async(T&& value, rmm::cuda_stream_view stream) | ||
| { | ||
| bounce_buffer[0] = std::move(value); | ||
| cuda_memcpy_async<T>(device_span<T>{this->data(), 1}, bounce_buffer, stream); | ||
| cuda_memcpy_async<T>(device_span<T>{data(), 1}, bounce_buffer, stream); | ||
| } | ||
|
|
||
| void set_value_to_zero_async(rmm::cuda_stream_view stream) { set_value_async(T{}, stream); } | ||
|
|
||
| [[nodiscard]] T* data() noexcept { return _storage.data(); } | ||
|
|
||
| [[nodiscard]] T const* data() const noexcept { return _storage.data(); } | ||
|
|
||
| private: | ||
| rmm::device_uvector<T> _storage; | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should we just go straight to
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I chose not to, I'd prefer to do all the |
||
| mutable cudf::detail::host_vector<T> bounce_buffer; | ||
| }; | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -14,7 +14,6 @@ | |
|
|
||
| #include <rmm/cuda_stream_view.hpp> | ||
| #include <rmm/device_buffer.hpp> | ||
| #include <rmm/device_scalar.hpp> | ||
|
|
||
| #include <span> | ||
| #include <string_view> | ||
|
|
@@ -189,7 +188,7 @@ class fixed_width_scalar : public scalar { | |
| [[nodiscard]] T const* data() const; | ||
|
|
||
| protected: | ||
| rmm::device_scalar<T> _data; ///< device memory containing the value | ||
| cudf::detail::device_scalar<T> _data; ///< device memory containing the value | ||
|
|
||
| /** | ||
| * @brief Construct a new fixed width scalar object. | ||
|
|
@@ -212,7 +211,7 @@ class fixed_width_scalar : public scalar { | |
| * @param stream CUDA stream used for device memory operations. | ||
| * @param mr Device memory resource to use for device memory allocation. | ||
| */ | ||
| fixed_width_scalar(rmm::device_scalar<T>&& data, | ||
| fixed_width_scalar(cudf::detail::device_scalar<T>&& data, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This certainly gives me some pause. We have essentially turned this constructor from public to internal since it requires an internal class to call it. |
||
| bool is_valid = true, | ||
| rmm::cuda_stream_view stream = cudf::get_default_stream(), | ||
| rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); | ||
|
|
@@ -274,7 +273,7 @@ class numeric_scalar : public detail::fixed_width_scalar<T> { | |
| * @param stream CUDA stream used for device memory operations. | ||
| * @param mr Device memory resource to use for device memory allocation. | ||
| */ | ||
| numeric_scalar(rmm::device_scalar<T>&& data, | ||
| numeric_scalar(cudf::detail::device_scalar<T>&& data, | ||
| bool is_valid = true, | ||
| rmm::cuda_stream_view stream = cudf::get_default_stream(), | ||
| rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); | ||
|
|
@@ -366,7 +365,7 @@ class fixed_point_scalar : public scalar { | |
| * @param stream CUDA stream used for device memory operations. | ||
| * @param mr Device memory resource to use for device memory allocation. | ||
| */ | ||
| fixed_point_scalar(rmm::device_scalar<rep_type>&& data, | ||
| fixed_point_scalar(cudf::detail::device_scalar<rep_type>&& data, | ||
| numeric::scale_type scale, | ||
| bool is_valid = true, | ||
| rmm::cuda_stream_view stream = cudf::get_default_stream(), | ||
|
|
@@ -402,7 +401,7 @@ class fixed_point_scalar : public scalar { | |
| [[nodiscard]] rep_type const* data() const; | ||
|
|
||
| protected: | ||
| rmm::device_scalar<rep_type> _data; ///< device memory containing the value | ||
| cudf::detail::device_scalar<rep_type> _data; ///< device memory containing the value | ||
| }; | ||
|
|
||
| /** | ||
|
|
@@ -476,7 +475,7 @@ class string_scalar : public scalar { | |
| * @param stream CUDA stream used for device memory operations. | ||
| * @param mr Device memory resource to use for device memory allocation. | ||
| */ | ||
| string_scalar(rmm::device_scalar<value_type>& data, | ||
| string_scalar(cudf::detail::device_scalar<value_type>& data, | ||
| bool is_valid = true, | ||
| rmm::cuda_stream_view stream = cudf::get_default_stream(), | ||
| rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); | ||
|
|
@@ -585,7 +584,7 @@ class chrono_scalar : public detail::fixed_width_scalar<T> { | |
| * @param stream CUDA stream used for device memory operations. | ||
| * @param mr Device memory resource to use for device memory allocation. | ||
| */ | ||
| chrono_scalar(rmm::device_scalar<T>&& data, | ||
| chrono_scalar(cudf::detail::device_scalar<T>&& data, | ||
| bool is_valid = true, | ||
| rmm::cuda_stream_view stream = cudf::get_default_stream(), | ||
| rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -93,7 +93,7 @@ std::pair<rmm::device_uvector<size_type>, bool> compute_single_pass_aggs( | |
| rmm::device_uvector<size_type> block_cardinality(grid_size, stream); | ||
|
|
||
| // Flag indicating whether a global memory aggregation fallback is required or not. | ||
| rmm::device_scalar<cuda::std::atomic_flag> needs_global_memory_fallback(stream); | ||
| rmm::device_uvector<cuda::std::atomic_flag> needs_global_memory_fallback(1, stream); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Probably worth adding or just moving the |
||
| CUDF_CUDA_TRY(cudaMemsetAsync( | ||
| needs_global_memory_fallback.data(), 0, sizeof(cuda::std::atomic_flag), stream.value())); | ||
|
|
||
|
|
@@ -110,8 +110,8 @@ std::pair<rmm::device_uvector<size_type>, bool> compute_single_pass_aggs( | |
|
|
||
| auto const needs_fallback = [&] { | ||
| cuda::std::atomic_flag h_needs_fallback; | ||
| // Cannot use `device_scalar::value` as it requires a copy constructor, which | ||
| // `atomic_flag` doesn't have. | ||
| // Cannot use a value-returning helper because atomic_flag is not copy-constructible; | ||
| // copy the raw bytes back to host instead. | ||
| CUDF_CUDA_TRY(cudf::detail::memcpy_async(&h_needs_fallback, | ||
| needs_global_memory_fallback.data(), | ||
| sizeof(cuda::std::atomic_flag), | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is the
device_uvectorimportant information in the developer guide?I saw Bradley's comment about
cuda::bufferand it seems this would need to be kept insync with internal/private data members of the class.