span.hpp

前往此文件的文档。

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

#pragma once

#include <cudf/detail/utilities/host_vector.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/export.hpp>

#include <rmm/device_buffer.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/device_vector.hpp>

#include <thrust/detail/raw_pointer_cast.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/memory.h>

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

namespace CUDF_EXPORT cudf {
constexpr std::size_t dynamic_extent = std::numeric_limits<std::size_t>::max();
 // 组结束
namespace detail {

template <typename T, std::size_t Extent, typename Derived>
class span_base {
 static_assert(Extent == dynamic_extent, "仅支持动态范围");

 public
 using element_type = T;
 using value_type = std::remove_cv<T>;
 using size_type = std::size_t;
 using difference_type = std::ptrdiff_t;
 using pointer = T*;
 using iterator = T*;
 using const_pointer = T const*;
 using reference = T&;
 using const_reference =
 T const&;

 static constexpr std::size_t extent = Extent;

 CUDF_HOST_DEVICE constexpr span_base() noexcept {} // 复制构造函数。
 CUDF_HOST_DEVICE constexpr span_base(pointer data, size_type size) : _data(data), _size(size) {}
 // constexpr span_base(pointer begin, pointer end) : _data(begin), _size(end - begin) {}
 CUDF_HOST_DEVICE constexpr span_base(span_base const&) noexcept = default;
 CUDF_HOST_DEVICE constexpr span_base& operator=(span_base const&) noexcept = default;

 [[nodiscard]] CUDF_HOST_DEVICE constexpr iterator begin() const noexcept { return _data; }
 [[nodiscard]] CUDF_HOST_DEVICE constexpr iterator end() const noexcept { return _data + _size; }
 [[nodiscard]] CUDF_HOST_DEVICE constexpr pointer data() const noexcept { return _data; }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr size_type size() const noexcept { return _size; }
 [[nodiscard]] CUDF_HOST_DEVICE constexpr size_type size_bytes() const noexcept
 {
 return sizeof(T) * _size;
 }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr bool empty() const noexcept { return _size == 0; }

 [[nodiscard]] constexpr Derived first(size_type count) const noexcept
 {
 return Derived(_data, count);
 }

 [[nodiscard]] constexpr Derived last(size_type count) const noexcept
 {
 return Derived(_data + _size - count, count);
 }

 protected
 pointer _data{nullptr};
 size_type _size{0};
};

} // namespace detail

// ===== host_span =================================================================================

template <typename T>
struct is_host_span_supported_container : std::false_type {};

template <typename T, typename Alloc>
struct is_host_span_supported_container< //
 std::vector<T, Alloc>> : std::true_type {};

template <typename T, typename Alloc>
struct is_host_span_supported_container< //
 thrust::host_vector<T, Alloc>> : std::true_type {};

template <typename T, typename Alloc>
struct is_host_span_supported_container< //
 std::basic_string<T, std::char_traits<T>, Alloc>> : std::true_type {};

template <typename T, std::size_t Extent = cudf::dynamic_extent>
struct host_span : public cudf::detail::span_base<T, Extent, host_span<T, Extent>> {
 using base = cudf::detail::span_base<T, Extent, host_span<T, Extent>>;
 using base::base;

 constexpr host_span() noexcept : base() {} // centos 编译所需

 CUDF_HOST_DEVICE constexpr host_span(T* data, std::size_t size, bool is_device_accessible)
 : base(data, size), _is_device_accessible{is_device_accessible}
 {
 }

 template <typename C,
 // Only supported containers of types convertible to T
 std::enable_if_t<is_host_span_supported_container<C>::value &&
 std::is_convertible_v<
 std::remove_pointer_t<decltype(thrust::raw_pointer_cast( // NOLINT
 std::declval<C&>().data()))> (*)[],
 T (*)[]>>* = nullptr> // NOLINT
 constexpr host_span(C& in) : base(thrust::raw_pointer_cast(in.data()), in.size())
 {
 }

 template <typename C,
 // Only supported containers of types convertible to T
 std::enable_if_t<is_host_span_supported_container<C>::value &&
 std::is_convertible_v<
 std::remove_pointer_t<decltype(thrust::raw_pointer_cast( // NOLINT
 std::declval<C&>().data()))> (*)[],
 T (*)[]>>* = nullptr> // NOLINT
 constexpr host_span(C const& in) : base(thrust::raw_pointer_cast(in.data()), in.size())
 {
 }

 template <typename OtherT,
 // Only supported containers of types convertible to T
 std::enable_if_t<std::is_convertible_v<OtherT (*)[], T (*)[]>>* = nullptr> // NOLINT
 constexpr host_span(cudf::detail::host_vector<OtherT>& in)
 : base(in.data(), in.size()), _is_device_accessible{in.get_allocator().is_device_accessible()}
 {
 }

 template <typename OtherT,
 // Only supported containers of types convertible to T
 std::enable_if_t<std::is_convertible_v<OtherT (*)[], T (*)[], void>>* = nullptr> // NOLINT
 constexpr host_span(cudf::detail::host_vector<OtherT> const& in)
 : base(in.data(), in.size()), _is_device_accessible{in.get_allocator().is_device_accessible()}
 {
 }

 // 拷贝构造以支持 const 转换
 template <typename OtherT,
 std::size_t OtherExtent,
 std::enable_if_t<(Extent == OtherExtent || Extent == dynamic_extent) &&
 std::is_convertible_v<OtherT (*)[], T (*)[], // NOLINT
 void>* = nullptr>
 constexpr host_span(host_span<OtherT, OtherExtent> const& other) noexcept
 : base(other.data(), other.size()), _is_device_accessible{other.is_device_accessible()}
 {
 }
 // 如果 idx < 0 或 idx >= size，则行为未定义，因此不是 noexcept
 constexpr typename base::reference operator[](typename base::size_type idx) const
 {
 static_assert(sizeof(idx) >= sizeof(size_t), "index type must not be smaller than size_t");
 return this->_data[idx];
 }

 // 如果 size = 0，则行为未定义，因此不是 noexcept
 [[nodiscard]] constexpr typename base::reference front() const { return this->_data[0]; }
 // 如果 size = 0，则行为未定义，因此不是 noexcept
 [[nodiscard]] constexpr typename base::reference back() const
 {
 return this->_data[this->_size - 1];
 }

 [[nodiscard]] bool is_device_accessible() const { return _is_device_accessible; }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr host_span subspan(
 typename base::size_type offset, typename base::size_type count) const noexcept
 {
 return host_span{this->data() + offset, count, _is_device_accessible};
 }

 private
 bool _is_device_accessible{false};
};

// ===== device_span ===============================================================================

template <typename T>
struct is_device_span_supported_container : std::false_type {};

template <typename T, typename Alloc>
struct is_device_span_supported_container< //
 thrust::device_vector<T, Alloc>> : std::true_type {};

template <typename T>
struct is_device_span_supported_container< //
 rmm::device_vector<T>> : std::true_type {};

template <typename T>
struct is_device_span_supported_container< //
 rmm::device_uvector<T>> : std::true_type {};

template <typename T, std::size_t Extent = cudf::dynamic_extent>
struct device_span : public cudf::detail::span_base<T, Extent, device_span<T, Extent>> {
 using base = cudf::detail::span_base<T, Extent, device_span<T, Extent>>;
 using base::base;

 CUDF_HOST_DEVICE constexpr device_span() noexcept : base() {} // centos 编译所需

 template <typename C,
 // Only supported containers of types convertible to T
 std::enable_if_t<is_device_span_supported_container<C>::value &&
 std::is_convertible_v<
 std::remove_pointer_t<decltype(thrust::raw_pointer_cast( // NOLINT
 std::declval<C&>().data()))> (*)[],
 T (*)[]>>* = nullptr> // NOLINT
 constexpr device_span(C& in) : base(thrust::raw_pointer_cast(in.data()), in.size())
 {
 }

 template <typename C,
 // Only supported containers of types convertible to T
 std::enable_if_t<is_device_span_supported_container<C>::value &&
 std::is_convertible_v<
 std::remove_pointer_t<decltype(thrust::raw_pointer_cast( // NOLINT
 std::declval<C&>().data()))> (*)[],
 T (*)[]>>* = nullptr> // NOLINT
 constexpr device_span(C const& in) : base(thrust::raw_pointer_cast(in.data()), in.size())
 {
 }

 // 拷贝构造以支持 const 转换
 template <typename OtherT,
 std::size_t OtherExtent,
 std::enable_if_t<(Extent == OtherExtent || Extent == dynamic_extent) &&
 std::is_convertible_v<OtherT (*)[], T (*)[], // NOLINT
 void>* = nullptr>
 CUDF_HOST_DEVICE constexpr device_span(device_span<OtherT, OtherExtent> const& other) noexcept
 : base(other.data(), other.size())
 {
 }

 // 如果 idx < 0 或 idx >= size，则行为未定义，因此不是 noexcept
 __device__ constexpr typename base::reference operator[](typename base::size_type idx) const
 {
 static_assert(sizeof(idx) >= sizeof(size_t), "index type must not be smaller than size_t");
 return this->_data[idx];
 }

 // 如果 size = 0，则行为未定义，因此不是 noexcept
 [[nodiscard]] __device__ constexpr typename base::reference front() const
 {
 return this->_data[0];
 }
 // 如果 size = 0，则行为未定义，因此不是 noexcept
 [[nodiscard]] __device__ constexpr typename base::reference back() const
 {
 return this->_data[this->_size - 1];
 }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr device_span subspan(
 typename base::size_type offset, typename base::size_type count) const noexcept
 {
 return device_span{this->data() + offset, count};
 }
}; // 组结束

namespace detail {

template <typename T, template <typename, std::size_t> typename RowType>
class base_2dspan {
 public
 using size_type =
 std::pair<size_t, size_t>

 constexpr base_2dspan() noexcept = default;
 constexpr base_2dspan(RowType<T, dynamic_extent> flat_view, size_t columns)
 : _flat{flat_view}, _size{columns == 0 ? 0 : flat_view.size() / columns, columns}
 {
#ifndef __CUDA_ARCH__
 CUDF_EXPECTS(_size.first * _size.second == flat_view.size(), "无效的二维跨度大小");
#endif
 }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr auto data() const noexcept { return _flat.data(); }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr auto size() const noexcept { return _size; }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr auto count() const noexcept { return _flat.size(); }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr bool is_empty() const noexcept { return count() == 0; }

 CUDF_HOST_DEVICE constexpr RowType<T, dynamic_extent> operator[](std::size_t row) const
 {
 return _flat.subspan(row * _size.second, _size.second);
 }

 [[nodiscard]] CUDF_HOST_DEVICE constexpr RowType<T, dynamic_extent> flat_view() const
 {
 return _flat;
 }

 template <typename OtherT,
 template <typename, size_t>
 typename OtherRowType,
 std::enable_if_t<std::is_convertible_v<OtherRowType<OtherT, dynamic_extent>,
 RowType<T, dynamic_extent>>,
 void>* = nullptr>
 constexpr base_2dspan(base_2dspan<OtherT, OtherRowType> const& other) noexcept
 : _flat{other.flat_view()}, _size{other.size()}
 {
 }

 protected
 RowType<T, dynamic_extent> _flat;
 size_type _size{0, 0};
};

template <class T>
using host_2dspan = base_2dspan<T, host_span>;

template <class T>
using device_2dspan = base_2dspan<T, device_span>;

} // namespace detail
} // namespace CUDF_EXPORT cudf