device_scalar.hpp

转到此文件的文档。

/*
 * 版权所有 (c) 2019-2025, NVIDIA CORPORATION.
 *
 * 根根据 Apache 许可，版本 2.0 (“许可”)获得许可；
 * 除非遵守许可，否则您不得使用此文件。
 * 您可以在以下位置获得许可的副本：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则，软件
 * 在许可下分发时按“原样”提供，
 * 不附带任何明示或暗示的保证或条件。
 * 请参阅许可了解管理权限和
 * 限制的特定语言。
 */

#pragma once

#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <type_traits>

namespace RMM_NAMESPACE {
template <typename T>
class device_scalar {
 public
 static_assert(std::is_trivially_copyable_v<T>, "Scalar type must be trivially copyable");

 using value_type = typename device_uvector<T>::value_type;
 using size_type = typename device_uvector<T>::size_type;
 using reference = typename device_uvector<T>::reference;
 using const_reference = typename device_uvector<T>::const_reference;
 using pointer =
 typename device_uvector<T>::pointer;
 using const_pointer = typename device_uvector<T>::const_pointer;

 RMM_EXEC_CHECK_DISABLE
 ~device_scalar() = default;

 RMM_EXEC_CHECK_DISABLE
 device_scalar(device_scalar&&) noexcept = default;

 device_scalar& operator=(device_scalar&&) noexcept = default;

 device_scalar(device_scalar const&) = delete;

 device_scalar& operator=(device_scalar const&) = delete;

 device_scalar() = delete;

 explicit device_scalar(cuda_stream_view stream,
 device_async_resource_ref mr = mr::get_current_device_resource_ref())
 : _storage{1, stream, mr}
 {
 }

 explicit device_scalar(value_type const& initial_value,
 cuda_stream_view stream,
 device_async_resource_ref mr = mr::get_current_device_resource_ref())
 : _storage{1, stream, mr}
 {
 set_value_async(initial_value, stream);
 }

 device_scalar(device_scalar const& other,
 cuda_stream_view stream,
 device_async_resource_ref mr = mr::get_current_device_resource_ref())
 : _storage{other._storage, stream, mr}
 {
 }

 [[nodiscard]] value_type value(cuda_stream_view stream) const
 {
 return _storage.front_element(stream);
 }

 void set_value_async(value_type const& value, cuda_stream_view stream)
 {
 _storage.set_element_async(0, value, stream);
 }

 // Disallow passing literals to set_value to avoid race conditions where the memory holding the
 // literal can be freed before the async memcpy / memset executes.
 void set_value_async(value_type&&, cuda_stream_view) = delete;

 void set_value_to_zero_async(cuda_stream_view stream)
 {
 _storage.set_element_to_zero_async(value_type{0}, stream);
 }

 [[nodiscard]] pointer data() noexcept { return static_cast<pointer>(_storage.data()); }

 [[nodiscard]] const_pointer data() const noexcept
 {
 return static_cast<const_pointer>(_storage.data());
 }

 [[nodiscard]] constexpr size_type size() const noexcept { return 1; }

 [[nodiscard]] cuda_stream_view stream() const noexcept { return _storage.stream(); }

 void set_stream(cuda_stream_view stream) noexcept { _storage.set_stream(stream); }

 private
 rmm::device_uvector<T> _storage;
};
 // 组结束
} // 命名空间 RMM_NAMESPACE