span.h

前往此文件文档。

 
#ifndef XGBOOST_SPAN_H_
#define XGBOOST_SPAN_H_
 
#include <xgboost/base.h>
 
#include <cstddef> // size_t
#include <cstdio>
#include <iterator>
#include <limits> // numeric_limits
#include <type_traits>
#include <utility> // for move
#include <vector> // for vector
 
#if defined(__CUDACC__)
#include <cuda_runtime.h>
#endif // defined(__CUDACC__)
 
#if defined(_MSC_VER) && _MSC_VER < 1910
 
#define __span_noexcept
 
#pragma push_macro("constexpr")
#define constexpr /*constexpr*/
 
#else
 
#define __span_noexcept noexcept
 
#endif // defined(_MSC_VER) && _MSC_VER < 1910
 
namespace xgboost::common {
 
#if defined(__CUDA_ARCH__)
// 在设备代码中通常无法使用日志记录工具。
 
#if defined(_MSC_VER)
 
// Windows CUDA 没有 __assert_fail。
#define CUDA_KERNEL_CHECK(cond) \
 do { \
 if (XGBOOST_EXPECT(!(cond), false)) { \
 asm("trap;"); \
    }                                     \
 } while (0)
 
#else // defined(_MSC_VER)
 
#define __ASSERT_STR_HELPER(x) #x
 
#define CUDA_KERNEL_CHECK(cond) \
 (XGBOOST_EXPECT((cond), true) \
 ? static_cast<void>(0) \
 : __assert_fail(__ASSERT_STR_HELPER((cond)), __FILE__, __LINE__, __PRETTY_FUNCTION__))
 
#endif // defined(_MSC_VER)
 
#define KERNEL_CHECK(cond) CUDA_KERNEL_CHECK(cond)
 
#define SPAN_CHECK(cond) KERNEL_CHECK(cond)
 
#else // ------------------------------ 非CUDA ----------------------------
 
#if defined(XGBOOST_STRICT_R_MODE) && XGBOOST_STRICT_R_MODE == 1
 
#define KERNEL_CHECK(cond)
 
#define SPAN_CHECK(cond) KERNEL_CHECK(cond)
 
#else
 
#if defined(__SYCL_DEVICE_ONLY__)
 
// SYCL不支持终止
#define SYCL_KERNEL_CHECK(cond)
 
#define KERNEL_CHECK(cond) SYCL_KERNEL_CHECK(cond)
 
#else // defined(__SYCL_DEVICE_ONLY__)
#define KERNEL_CHECK(cond) (XGBOOST_EXPECT((cond), true) ? static_cast<void>(0) : std::terminate())
#endif // defined(__SYCL_DEVICE_ONLY__)
 
#define SPAN_CHECK(cond) KERNEL_CHECK(cond)
 
#endif // defined(XGBOOST_STRICT_R_MODE)
 
#endif // __CUDA_ARCH__
 
#define SPAN_LT(lhs, rhs) SPAN_CHECK((lhs) < (rhs))
 
namespace detail {
using ptrdiff_t = typename std::conditional_t< // NOLINT
 std::is_same_v<std::ptrdiff_t, std::int64_t>, std::ptrdiff_t, std::int64_t>;
} // namespace detail
 
#if defined(_MSC_VER) && _MSC_VER < 1910
constexpr const std::size_t
dynamic_extent = std::numeric_limits<std::size_t>::max(); // NOLINT
#else
constexpr std::size_t dynamic_extent = std::numeric_limits<std::size_t>::max(); // NOLINT
#endif // defined(_MSC_VER) && _MSC_VER < 1910
 
enum class byte : unsigned char {}; // NOLINT
 
template <class ElementType, std::size_t Extent>
class Span;
 
namespace detail {
 
template <typename SpanType, bool IsConst>
class SpanIterator {
 using ElementType = typename SpanType::element_type;
 
 public
 using iterator_category = std::random_access_iterator_tag; // NOLINT
 using value_type = typename SpanType::value_type; // NOLINT
 using difference_type = detail::ptrdiff_t; // NOLINT
 
 using reference = typename std::conditional< // NOLINT
 IsConst, const ElementType, ElementType>::type&;
 using pointer = typename std::add_pointer<reference>::type; // NOLINT
 
 constexpr SpanIterator() = default;
 
 XGBOOST_DEVICE constexpr SpanIterator(
 const SpanType* _span,
 typename SpanType::index_type _idx) __span_noexcept
 span_(_span), index_(_idx) {}
 
 friend SpanIterator<SpanType, true>;
 template <bool B, typename std::enable_if_t<!B && IsConst>* = nullptr>
 XGBOOST_DEVICE constexpr SpanIterator( // NOLINT
 const SpanIterator<SpanType, B>& other_) __span_noexcept
 : SpanIterator(other_.span_, other_.index_) {}
 
 XGBOOST_DEVICE reference operator*() const {
 SPAN_CHECK(index_ < span_->size());
 return *(span_->data() + index_);
  }
 XGBOOST_DEVICE reference operator[](difference_type n) const {
 return *(*this + n);
  }
 
 XGBOOST_DEVICE pointer operator->() const {
 SPAN_CHECK(index_ != span_->size());
 return span_->data() + index_;
  }
 
 XGBOOST_DEVICE SpanIterator& operator++() {
 SPAN_CHECK(index_ != span_->size());
 index_++;
 return *this;
  }
 
 XGBOOST_DEVICE SpanIterator operator++(int) {
 auto ret = *this;
 ++(*this);
 return ret;
  }
 
 XGBOOST_DEVICE SpanIterator& operator--() {
 SPAN_CHECK(index_ != 0 && index_ <= span_->size());
 index_--;
 return *this;
  }
 
 XGBOOST_DEVICE SpanIterator operator--(int) {
 auto ret = *this;
 --(*this);
 return ret;
  }
 
 XGBOOST_DEVICE SpanIterator operator+(difference_type n) const {
 auto ret = *this;
 return ret += n;
  }
 
 XGBOOST_DEVICE SpanIterator& operator+=(difference_type n) {
 SPAN_CHECK((index_ + n) <= span_->size());
 index_ += n;
 return *this;
  }
 
 XGBOOST_DEVICE difference_type operator-(SpanIterator rhs) const {
 SPAN_CHECK(span_ == rhs.span_);
 return index_ - rhs.index_;
  }
 
 XGBOOST_DEVICE SpanIterator operator-(difference_type n) const {
 auto ret = *this;
 return ret -= n;
  }
 
 XGBOOST_DEVICE SpanIterator& operator-=(difference_type n) {
 return *this += -n;
  }
 
 // 友元
 XGBOOST_DEVICE constexpr friend bool operator==(
 SpanIterator _lhs, SpanIterator _rhs) __span_noexcept {
 return _lhs.span_ == _rhs.span_ && _lhs.index_ == _rhs.index_;
  }
 
 XGBOOST_DEVICE constexpr friend bool operator!=(
 SpanIterator _lhs, SpanIterator _rhs) __span_noexcept {
 return !(_lhs == _rhs);
  }
 
 XGBOOST_DEVICE constexpr friend bool operator<(
 SpanIterator _lhs, SpanIterator _rhs) __span_noexcept {
 return _lhs.index_ < _rhs.index_;
  }
 
 XGBOOST_DEVICE constexpr friend bool operator<=(
 SpanIterator _lhs, SpanIterator _rhs) __span_noexcept {
 return !(_rhs < _lhs);
  }
 
 XGBOOST_DEVICE constexpr friend bool operator>(
 SpanIterator _lhs, SpanIterator _rhs) __span_noexcept {
 return _rhs < _lhs;
  }
 
 XGBOOST_DEVICE constexpr friend bool operator>=(
 SpanIterator _lhs, SpanIterator _rhs) __span_noexcept {
 return !(_rhs > _lhs);
  }
 
 protected
 const SpanType *span_ { nullptr };
 typename SpanType::index_type index_ { 0 };
};
 
 
// 这里倾向于使用constexpr而不是结构体来完成以下元编程。
// 但请记住，我们这里支持MSVC 2013。
 
template <std::size_t Extent, std::size_t Offset, std::size_t Count>
struct ExtentValue : public std::integral_constant<
 std::size_t, Count != dynamic_extent ?
 Count : (Extent != dynamic_extent ? Extent - Offset : Extent)> {};
 
template <typename T, std::size_t Extent>
struct ExtentAsBytesValue : public std::integral_constant<
 std::size_t,
 Extent == dynamic_extent ?
 Extent : sizeof(T) * Extent> {};
 
template <std::size_t From, std::size_t To>
struct IsAllowedExtentConversion : public std::integral_constant<
 bool, From == To || From == dynamic_extent || To == dynamic_extent> {};
 
template <class From, class To>
struct IsAllowedElementTypeConversion
 : public std::integral_constant<bool, std::is_convertible_v<From (*)[], To (*)[]>> {}; // NOLINT
 
template <class T>
struct IsSpanOracle : std::false_type {};
 
template <class T, std::size_t Extent>
struct IsSpanOracle<Span<T, Extent>> : std::true_type {};
 
template <class T>
struct IsSpan : public IsSpanOracle<typename std::remove_cv_t<T>> {};
 
// 在此处重新实现标准算法以适应CUDA。
template <typename T>
struct Less {
 XGBOOST_DEVICE constexpr bool operator()(const T& _x, const T& _y) const {
 return _x < _y;
  }
};
 
template <typename T>
struct Greater {
 XGBOOST_DEVICE constexpr bool operator()(const T& _x, const T& _y) const {
 return _x > _y;
  }
};
 
template <class InputIt1, class InputIt2,
 class Compare =
 detail::Less<decltype(std::declval<InputIt1>().operator*())>>
XGBOOST_DEVICE bool LexicographicalCompare(InputIt1 first1, InputIt1 last1,
 InputIt2 first2, InputIt2 last2) {
 Compare comp;
 for (; first1 != last1 && first2 != last2; ++first1, ++first2) {
 if (comp(*first1, *first2)) {
 return true;
    }
 if (comp(*first2, *first1)) {
 return false;
    }
  }
 return first1 == last1 && first2 != last2;
}
 
} // namespace detail
 
 
template <typename T,
 std::size_t Extent = dynamic_extent>
class Span {
 public
 using element_type = T; // NOLINT
 using value_type = typename std::remove_cv<T>::type; // NOLINT
 using index_type = std::size_t; // NOLINT
 using difference_type = detail::ptrdiff_t; // NOLINT
 using pointer = T*; // NOLINT
 using reference = T&; // NOLINT
 
 using iterator = detail::SpanIterator<Span<T, Extent>, false>; // NOLINT
 using const_iterator = const detail::SpanIterator<Span<T, Extent>, true>; // NOLINT
 using reverse_iterator = std::reverse_iterator<iterator>; // NOLINT
 using const_reverse_iterator = const std::reverse_iterator<const_iterator>; // NOLINT
 
 // 构造函数
 constexpr Span() = default;
 
 XGBOOST_DEVICE Span(pointer _ptr, index_type _count)
 size_(_count), data_(_ptr) {
 SPAN_CHECK(!(Extent != dynamic_extent && _count != Extent));
 SPAN_CHECK(_ptr || _count == 0);
  }
 
 XGBOOST_DEVICE Span(pointer _first, pointer _last)
 size_(_last - _first), data_(_first) {
 SPAN_CHECK(data_ || size_ == 0);
  }
 
 template <std::size_t N>
 XGBOOST_DEVICE constexpr Span(element_type (&arr)[N]) // NOLINT
 __span_noexcept : size_(N), data_(&arr[0]) {}
 
 template <class Container,
 class = typename std::enable_if_t<
 !std::is_const_v<element_type> && !detail::IsSpan<Container>::value &&
 std::is_convertible_v<typename Container::pointer, pointer> &&
 std::is_convertible_v<typename Container::pointer,
 decltype(std::declval<Container>().data())>>>
 Span(Container& _cont) // NOLINT
 : size_(_cont.size()), data_(_cont.data()) {
 static_assert(!detail::IsSpan<Container>::value, "调用了错误的Span构造函数。");
  }
 
 template <class Container,
 class = typename std::enable_if_t<
 std::is_const_v<element_type> && !detail::IsSpan<Container>::value &&
 std::is_convertible_v<typename Container::pointer, pointer> &&
 std::is_convertible_v<typename Container::pointer,
 decltype(std::declval<Container>().data())>>>
 Span(const Container& _cont) // NOLINT
 : size_(_cont.size()), data_(_cont.data()) {
 static_assert(!detail::IsSpan<Container>::value, "调用了错误的Span构造函数。");
  }
 
 template <class U, std::size_t OtherExtent,
 class = typename std::enable_if_t<
 detail::IsAllowedElementTypeConversion<U, T>::value &&
 detail::IsAllowedExtentConversion<OtherExtent, Extent>::value>>
 XGBOOST_DEVICE constexpr Span(const Span<U, OtherExtent>& _other) // NOLINT
 __span_noexcept : size_(_other.size()),
 data_(_other.data()) {}
 
 constexpr Span(Span const& _other) noexcept(true) = default;
 constexpr Span& operator=(Span const& _other) noexcept(true) = default;
 constexpr Span(Span&& _other) noexcept(true) = default;
 constexpr Span& operator=(Span&& _other) noexcept(true) = default;
 ~Span() noexcept(true) = default;
 
 XGBOOST_DEVICE constexpr iterator begin() const __span_noexcept { // NOLINT
 return {this, 0};
  }
 
 XGBOOST_DEVICE constexpr iterator end() const __span_noexcept { // NOLINT
 return {this, size()};
  }
 
 XGBOOST_DEVICE constexpr const_iterator cbegin() const __span_noexcept { // NOLINT
 return {this, 0};
  }
 
 XGBOOST_DEVICE constexpr const_iterator cend() const __span_noexcept { // NOLINT
 return {this, size()};
  }
 
 constexpr reverse_iterator rbegin() const __span_noexcept { // NOLINT
 return reverse_iterator{end()};
  }
 
 constexpr reverse_iterator rend() const __span_noexcept { // NOLINT
 return reverse_iterator{begin()};
  }
 
 XGBOOST_DEVICE constexpr const_reverse_iterator crbegin() const __span_noexcept { // NOLINT
 return const_reverse_iterator{cend()};
  }
 
 XGBOOST_DEVICE constexpr const_reverse_iterator crend() const __span_noexcept { // NOLINT
 return const_reverse_iterator{cbegin()};
  }
 
 // 元素访问
 
 XGBOOST_DEVICE reference front() const { // NOLINT
 return (*this)[0];
  }
 
 XGBOOST_DEVICE reference back() const { // NOLINT
 return (*this)[size() - 1];
  }
 
 XGBOOST_DEVICE reference operator[](index_type _idx) const {
 SPAN_LT(_idx, size());
 return data()[_idx];
  }
 
 XGBOOST_DEVICE reference operator()(index_type _idx) const {
 return this->operator[](_idx);
  }
 
 XGBOOST_DEVICE constexpr pointer data() const __span_noexcept { // NOLINT
 return data_;
  }
 
 // 观察者
 XGBOOST_DEVICE constexpr index_type size() const __span_noexcept { // NOLINT
 return size_;
  }
 XGBOOST_DEVICE constexpr index_type size_bytes() const __span_noexcept { // NOLINT
 return size() * sizeof(T);
  }
 
 XGBOOST_DEVICE constexpr bool empty() const __span_noexcept { // NOLINT
 return size() == 0;
  }
 
 // 子视图
 template <std::size_t Count>
 XGBOOST_DEVICE Span<element_type, Count> first() const { // NOLINT
 SPAN_CHECK(Count <= size());
 return {data(), Count};
  }
 
 XGBOOST_DEVICE Span<element_type, dynamic_extent> first( // NOLINT
 std::size_t _count) const {
 SPAN_CHECK(_count <= size());
 return {data(), _count};
  }
 
 template <std::size_t Count>
 XGBOOST_DEVICE Span<element_type, Count> last() const { // NOLINT
 SPAN_CHECK(Count <= size());
 return {data() + size() - Count, Count};
  }
 
 XGBOOST_DEVICE Span<element_type, dynamic_extent> last( // NOLINT
 std::size_t _count) const {
 SPAN_CHECK(_count <= size());
 return subspan(size() - _count, _count);
  }
 
 template <std::size_t Offset,
 std::size_t Count = dynamic_extent>
 XGBOOST_DEVICE auto subspan() const -> // NOLINT
 Span<element_type,
 detail::ExtentValue<Extent, Offset, Count>::value> {
 SPAN_CHECK((Count == dynamic_extent) ?
 (Offset <= size()) : (Offset + Count <= size()));
 return {data() + Offset, Count == dynamic_extent ? size() - Offset : Count};
  }
 
 XGBOOST_DEVICE Span<element_type, dynamic_extent> subspan( // NOLINT
 index_type _offset,
 index_type _count = dynamic_extent) const {
 SPAN_CHECK((_count == dynamic_extent) ?
 (_offset <= size()) : (_offset + _count <= size()));
 return {data() + _offset, _count ==
 dynamic_extent ? size() - _offset : _count};
  }
 
 private
 index_type size_ { 0 };
 pointer data_ { nullptr };
};
 
template <class T, std::size_t X, class U, std::size_t Y>
XGBOOST_DEVICE bool operator==(Span<T, X> l, Span<U, Y> r) {
 if (l.size() != r.size()) {
 return false;
  }
 for (auto l_beg = l.cbegin(), r_beg = r.cbegin(); l_beg != l.cend();
 ++l_beg, ++r_beg) {
 if (*l_beg != *r_beg) {
 return false;
    }
  }
 return true;
}
 
template <class T, std::size_t X, class U, std::size_t Y>
XGBOOST_DEVICE constexpr bool operator!=(Span<T, X> l, Span<U, Y> r) {
 return !(l == r);
}
 
template <class T, std::size_t X, class U, std::size_t Y>
XGBOOST_DEVICE constexpr bool operator<(Span<T, X> l, Span<U, Y> r) {
 return detail::LexicographicalCompare(l.begin(), l.end(),
 r.begin(), r.end());
}
 
template <class T, std::size_t X, class U, std::size_t Y>
XGBOOST_DEVICE constexpr bool operator<=(Span<T, X> l, Span<U, Y> r) {
 return !(l > r);
}
 
template <class T, std::size_t X, class U, std::size_t Y>
XGBOOST_DEVICE constexpr bool operator>(Span<T, X> l, Span<U, Y> r) {
 return detail::LexicographicalCompare<
 typename Span<T, X>::iterator, typename Span<U, Y>::iterator,
 detail::Greater<typename Span<T, X>::element_type>>(l.begin(), l.end(),
 r.begin(), r.end());
}
 
template <class T, std::size_t X, class U, std::size_t Y>
XGBOOST_DEVICE constexpr bool operator>=(Span<T, X> l, Span<U, Y> r) {
 return !(l < r);
}
 
template <class T, std::size_t E>
XGBOOST_DEVICE auto as_bytes(Span<T, E> s) __span_noexcept -> // NOLINT
 Span<const byte, detail::ExtentAsBytesValue<T, E>::value> {
 return {reinterpret_cast<const byte*>(s.data()), s.size_bytes()};
}
 
template <class T, std::size_t E>
XGBOOST_DEVICE auto as_writable_bytes(Span<T, E> s) __span_noexcept -> // NOLINT
 Span<byte, detail::ExtentAsBytesValue<T, E>::value> {
 return {reinterpret_cast<byte*>(s.data()), s.size_bytes()};
}
 
template <typename It>
class IterSpan {
 public
 using value_type = typename std::iterator_traits<It>::value_type; // NOLINT
 using index_type = std::size_t; // NOLINT
 using iterator = It; // NOLINT
 
 private
 It it_;
 index_type size_{0};
 
 public
 IterSpan() = default;
 XGBOOST_DEVICE IterSpan(It it, index_type size) : it_{std::move(it)}, size_{size} {}
 XGBOOST_DEVICE explicit IterSpan(common::Span<It, dynamic_extent> span)
 : it_{span.data()}, size_{span.size()} {}
 
 [[nodiscard]] XGBOOST_DEVICE index_type size() const noexcept { return size_; } // NOLINT
 [[nodiscard]] XGBOOST_DEVICE decltype(auto) operator[](index_type i) const { return it_[i]; }
 [[nodiscard]] XGBOOST_DEVICE decltype(auto) operator[](index_type i) { return it_[i]; }
 [[nodiscard]] XGBOOST_DEVICE bool empty() const noexcept { return size() == 0; } // NOLINT
 [[nodiscard]] XGBOOST_DEVICE It data() const noexcept { return it_; } // NOLINT
 [[nodiscard]] XGBOOST_DEVICE IterSpan<It> subspan( // NOLINT
 index_type _offset, index_type _count = dynamic_extent) const {
 SPAN_CHECK((_count == dynamic_extent) ? (_offset <= size()) : (_offset + _count <= size()));
 return {data() + _offset, _count == dynamic_extent ? size() - _offset : _count};
  }
 [[nodiscard]] XGBOOST_DEVICE constexpr iterator begin() const noexcept { // NOLINT
 return it_;
  }
 [[nodiscard]] XGBOOST_DEVICE constexpr iterator end() const noexcept { // NOLINT
 return it_ + size();
  }
};
 
template <typename T>
Span(std::vector<T> const&) -> Span<T const>;
 
template <typename T>
Span(std::vector<T>&) -> Span<T>;
} // namespace xgboost::common
 
 
#if defined(_MSC_VER) &&_MSC_VER < 1910
#undef constexpr
#pragma pop_macro("constexpr")
#undef __span_noexcept
#endif // _MSC_VER < 1910
 
#endif // XGBOOST_SPAN_H_