device_uvector.hpp

转到此文件的文档。

/*
 * Copyright (c) 2020-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/exec_check_disable.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/device_buffer.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cstddef>
#include <type_traits>
#include <utility>

namespace RMM_NAMESPACE {
template <typename T>
class device_uvector {
 static_assert(std::is_trivially_copyable_v<T>,
 "device_uvector 仅支持可平凡复制的类型。");

 public
 using value_type = T;
 using size_type = std::size_t;
 using reference = value_type&
 using const_reference = value_type const&
 using pointer = value_type*;
 using const_pointer = value_type const*;
 using iterator = pointer;
 using const_iterator = const_pointer;

 RMM_EXEC_CHECK_DISABLE
 ~device_uvector() = default;

 RMM_EXEC_CHECK_DISABLE
 device_uvector(device_uvector&&) noexcept = default;

 RMM_EXEC_CHECK_DISABLE
 device_uvector& operator=(device_uvector&&) noexcept =
 default;

 device_uvector(device_uvector const&) = delete;

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

 device_uvector() = delete;

 explicit device_uvector(std::size_t size,
 cuda_stream_view stream,
 device_async_resource_ref mr = mr::get_current_device_resource_ref())
 : _storage{elements_to_bytes(size), stream, mr}
 {
 }

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

 [[nodiscard]] pointer element_ptr(std::size_t element_index) noexcept
 {
 assert(element_index < size() && "尝试访问越界元素。");
 return data() + element_index;
 }

 [[nodiscard]] const_pointer element_ptr(std::size_t element_index) const noexcept
 {
 assert(element_index < size() && "尝试访问越界元素。");
 return data() + element_index;
 }

 void set_element_async(std::size_t element_index,
 value_type const& value,
 cuda_stream_view stream)
 {
 RMM_EXPECTS(
 element_index < size(), "尝试访问越界元素。", rmm::out_of_range);

 if constexpr (std::is_same_v<value_type, bool>) {
 RMM_CUDA_TRY(
 cudaMemsetAsync(element_ptr(element_index), value, sizeof(value), stream.value()));
 return;
 }

 if constexpr (std::is_fundamental_v<value_type>) {
 if (value == value_type{0}) {
 set_element_to_zero_async(element_index, stream);
 return;
 }
 }

 RMM_CUDA_TRY(cudaMemcpyAsync(
 element_ptr(element_index), &value, sizeof(value), cudaMemcpyDefault, stream.value()));
 }

 // 我们删除右值引用重载，以防止在字面值或隐式临时值被删除或超出范围后对其进行异步复制。
 // We delete the r-value reference overload to prevent asynchronously copying from a literal or
 // implicit temporary value after it is deleted or goes out of scope.
 void set_element_async(std::size_t, value_type const&&, cuda_stream_view) = delete;
 void set_element_to_zero_async(std::size_t element_index, cuda_stream_view stream)
 {
 RMM_EXPECTS(
 element_index < size(), "尝试访问越界元素。", rmm::out_of_range);
 RMM_CUDA_TRY(
 cudaMemsetAsync(element_ptr(element_index), 0, sizeof(value_type), stream.value()));
 }

 void set_element(std::size_t element_index, T const& value, cuda_stream_view stream)
 {
 set_element_async(element_index, value, stream);
 stream.synchronize_no_throw();
 }

 [[nodiscard]] value_type element(std::size_t element_index, cuda_stream_view stream) const
 {
 RMM_EXPECTS(
 element_index < size(), "尝试访问越界元素。", rmm::out_of_range);
 value_type value;
 RMM_CUDA_TRY(cudaMemcpyAsync(
 &value, element_ptr(element_index), sizeof(value), cudaMemcpyDefault, stream.value()));
 stream.synchronize();
 return value;
 }

 [[nodiscard]] value_type front_element(cuda_stream_view stream) const
 {
 return element(0, stream);
 }

 [[nodiscard]] value_type back_element(cuda_stream_view stream) const
 {
 return element(size() - 1, stream);
 }

 void reserve(std::size_t new_capacity, cuda_stream_view stream)
 {
 _storage.reserve(elements_to_bytes(new_capacity), stream);
 }

 void resize(std::size_t new_size, cuda_stream_view stream)
 {
 _storage.resize(elements_to_bytes(new_size), stream);
 }

 void shrink_to_fit(cuda_stream_view stream) { _storage.shrink_to_fit(stream); }

 device_buffer release() noexcept { return std::move(_storage); }

 [[nodiscard]] std::size_t capacity() const noexcept
 {
 return bytes_to_elements(_storage.capacity());
 }

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

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

 [[nodiscard]] iterator begin() noexcept { return data(); }

 [[nodiscard]] const_iterator cbegin() const noexcept { return data(); }

 [[nodiscard]] const_iterator begin() const noexcept { return cbegin(); }

 [[nodiscard]] iterator end() noexcept { return data() + size(); }

 [[nodiscard]] const_iterator cend() const noexcept { return data() + size(); }

 [[nodiscard]] const_iterator end() const noexcept { return cend(); }

 [[nodiscard]] std::size_t size() const noexcept { return bytes_to_elements(_storage.size()); }

 [[nodiscard]] std::int64_t ssize() const noexcept
 {
 assert(size() < static_cast<std::size_t>(std::numeric_limits<int64_t>::max()) &&
 "大小溢出有符号整数");
 return static_cast<int64_t>(size());
 }

 [[nodiscard]] bool is_empty() const noexcept { return size() == 0; }

 [[nodiscard]] rmm::device_async_resource_ref memory_resource() const noexcept
 {
 return _storage.memory_resource();
 }

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

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

 private
 device_buffer _storage{};

 [[nodiscard]] std::size_t constexpr elements_to_bytes(std::size_t num_elements) const noexcept
 {
 return num_elements * sizeof(value_type);
 }

 [[nodiscard]] std::size_t constexpr bytes_to_elements(std::size_t num_bytes) const noexcept
 {
 return num_bytes / sizeof(value_type);
 }
};
 // 组结束
} // namespace RMM_NAMESPACE