Skip to content

Commit

Permalink
Rename + flexiable indexing.
Browse files Browse the repository at this point in the history
  • Loading branch information
@trivialfis
trivialfis committed Dec 17, 2021
1 parent d0b2558 commit 92f8871
Show file tree
Hide file tree
Showing 4 changed files with 109 additions and 108 deletions.
71 changes: 36 additions & 35 deletions cpp/include/raft/common/span.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ namespace raft::common {
constexpr std::size_t dynamic_extent = std::numeric_limits<std::size_t>::max();

template <class ElementType, bool is_device, std::size_t Extent>
class span_t;
class span;

namespace detail {
/*!
Expand Down Expand Up @@ -72,7 +72,7 @@ struct is_span_oracle_t : std::false_type {
};

template <class T, bool is_device, std::size_t Extent>
struct is_span_oracle_t<span_t<T, is_device, Extent>> : std::true_type {
struct is_span_oracle_t<span<T, is_device, Extent>> : std::true_type {
};

template <class T>
Expand All @@ -99,7 +99,7 @@ __host__ __device__ constexpr auto lexicographical_compare(InputIt1 first1,
* most of the methods have bound check on debug build.
*/
template <typename T, bool is_device, std::size_t Extent = dynamic_extent>
class span_t {
class span {
public:
using element_type = T;
using value_type = typename std::remove_cv<T>::type;
Expand All @@ -116,21 +116,21 @@ class span_t {
using const_reverse_iterator = thrust::reverse_iterator<const_iterator>;

// constructors
constexpr span_t() noexcept = default;
constexpr span() noexcept = default;

constexpr span_t(pointer _ptr, size_type _count) noexcept : size_(_count), data_(_ptr)
constexpr span(pointer _ptr, size_type _count) noexcept : size_(_count), data_(_ptr)
{
assert(!(Extent != dynamic_extent && _count != Extent));
assert(_ptr || _count == 0);
}

constexpr span_t(pointer _first, pointer _last) noexcept : size_(_last - _first), data_(_first)
constexpr span(pointer _first, pointer _last) noexcept : size_(_last - _first), data_(_first)
{
assert(data_ || size_ == 0);
}

template <std::size_t N>
constexpr span_t(element_type (&arr)[N]) noexcept : size_(N), data_(&arr[0])
constexpr span(element_type (&arr)[N]) noexcept : size_(N), data_(&arr[0])
{
}

Expand All @@ -139,14 +139,14 @@ class span_t {
class = typename std::enable_if<
detail::is_allowed_element_type_conversion_t<U, T>::value &&
detail::is_allowed_extent_conversion_t<OtherExtent, Extent>::value>>
constexpr span_t(const span_t<U, is_device, OtherExtent>& _other) noexcept
constexpr span(const span<U, is_device, OtherExtent>& _other) noexcept
: size_(_other.size()), data_(_other.data())
{
}

constexpr span_t(const span_t& _other) noexcept : size_(_other.size()), data_(_other.data()) {}
constexpr span(const span& _other) noexcept : size_(_other.size()), data_(_other.data()) {}

constexpr auto operator=(const span_t& _other) noexcept -> span_t&
constexpr auto operator=(const span& _other) noexcept -> span&
{
size_ = _other.size();
data_ = _other.data();
Expand Down Expand Up @@ -186,7 +186,8 @@ class span_t {

constexpr auto back() const -> reference { return (*this)[size() - 1]; }

constexpr auto operator[](size_type _idx) const -> reference
template <typename Index>
constexpr auto operator[](Index _idx) const -> reference
{
assert(_idx < size());
return data()[_idx];
Expand All @@ -205,26 +206,26 @@ class span_t {

// Subviews
template <std::size_t Count>
constexpr auto first() const -> span_t<element_type, is_device, Count>
constexpr auto first() const -> span<element_type, is_device, Count>
{
assert(Count <= size());
return {data(), Count};
}

constexpr auto first(std::size_t _count) const -> span_t<element_type, is_device, dynamic_extent>
constexpr auto first(std::size_t _count) const -> span<element_type, is_device, dynamic_extent>
{
assert(_count <= size());
return {data(), _count};
}

template <std::size_t Count>
constexpr auto last() const -> span_t<element_type, is_device, Count>
constexpr auto last() const -> span<element_type, is_device, Count>
{
assert(Count <= size());
return {data() + size() - Count, Count};
}

constexpr auto last(std::size_t _count) const -> span_t<element_type, is_device, dynamic_extent>
constexpr auto last(std::size_t _count) const -> span<element_type, is_device, dynamic_extent>
{
assert(_count <= size());
return subspan(size() - _count, _count);
Expand All @@ -236,14 +237,14 @@ class span_t {
*/
template <std::size_t Offset, std::size_t Count = dynamic_extent>
constexpr auto subspan() const
-> span_t<element_type, is_device, detail::extent_value_t<Extent, Offset, Count>::value>
-> span<element_type, is_device, detail::extent_value_t<Extent, Offset, Count>::value>
{
assert((Count == dynamic_extent) ? (Offset <= size()) : (Offset + Count <= size()));
return {data() + Offset, Count == dynamic_extent ? size() - Offset : Count};
}

constexpr auto subspan(size_type _offset, size_type _count = dynamic_extent) const
-> span_t<element_type, is_device, dynamic_extent>
-> span<element_type, is_device, dynamic_extent>
{
assert((_count == dynamic_extent) ? (_offset <= size()) : (_offset + _count <= size()));
return {data() + _offset, _count == dynamic_extent ? size() - _offset : _count};
Expand All @@ -255,13 +256,13 @@ class span_t {
};

template <typename T, size_t extent = dynamic_extent>
using host_span_t = span_t<T, false, extent>;
using host_span = span<T, false, extent>;

template <typename T, size_t extent = dynamic_extent>
using device_span_t = span_t<T, true, extent>;
using device_span = span<T, true, extent>;

template <class T, std::size_t X, class U, std::size_t Y, bool is_device>
constexpr auto operator==(span_t<T, is_device, X> l, span_t<U, is_device, Y> r) -> bool
constexpr auto operator==(span<T, is_device, X> l, span<U, is_device, Y> r) -> bool
{
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) {
Expand All @@ -271,53 +272,53 @@ constexpr auto operator==(span_t<T, is_device, X> l, span_t<U, is_device, Y> r)
}

template <class T, std::size_t X, class U, std::size_t Y, bool is_device>
constexpr auto operator!=(span_t<T, is_device, X> l, span_t<U, is_device, Y> r)
constexpr auto operator!=(span<T, is_device, X> l, span<U, is_device, Y> r)
{
return !(l == r);
}

template <class T, std::size_t X, class U, std::size_t Y, bool is_device>
constexpr auto operator<(span_t<T, is_device, X> l, span_t<U, is_device, Y> r)
constexpr auto operator<(span<T, is_device, X> l, span<U, is_device, Y> r)
{
return detail::lexicographical_compare<
typename span_t<T, is_device, X>::iterator,
typename span_t<U, is_device, Y>::iterator,
thrust::less<typename span_t<T, is_device, X>::element_type>>(
typename span<T, is_device, X>::iterator,
typename span<U, is_device, Y>::iterator,
thrust::less<typename span<T, is_device, X>::element_type>>(
l.begin(), l.end(), r.begin(), r.end());
}

template <class T, std::size_t X, class U, std::size_t Y, bool is_device>
constexpr auto operator<=(span_t<T, is_device, X> l, span_t<U, is_device, Y> r)
constexpr auto operator<=(span<T, is_device, X> l, span<U, is_device, Y> r)
{
return !(l > r);
}

template <class T, std::size_t X, class U, std::size_t Y, bool is_device>
constexpr auto operator>(span_t<T, is_device, X> l, span_t<U, is_device, Y> r)
constexpr auto operator>(span<T, is_device, X> l, span<U, is_device, Y> r)
{
return detail::lexicographical_compare<
typename span_t<T, is_device, X>::iterator,
typename span_t<U, is_device, Y>::iterator,
thrust::greater<typename span_t<T, is_device, X>::element_type>>(
typename span<T, is_device, X>::iterator,
typename span<U, is_device, Y>::iterator,
thrust::greater<typename span<T, is_device, X>::element_type>>(
l.begin(), l.end(), r.begin(), r.end());
}

template <class T, std::size_t X, class U, std::size_t Y, bool is_device>
constexpr auto operator>=(span_t<T, is_device, X> l, span_t<U, is_device, Y> r)
constexpr auto operator>=(span<T, is_device, X> l, span<U, is_device, Y> r)
{
return !(l < r);
}

template <class T, bool is_device, std::size_t E>
auto as_bytes(span_t<T, is_device, E> s) noexcept
-> span_t<const std::byte, is_device, detail::extent_as_bytes_value_t<T, E>::value>
auto as_bytes(span<T, is_device, E> s) noexcept
-> span<const std::byte, is_device, detail::extent_as_bytes_value_t<T, E>::value>
{
return {reinterpret_cast<const std::byte*>(s.data()), s.size_bytes()};
}

template <class T, bool is_device, std::size_t E>
auto as_writable_bytes(span_t<T, is_device, E> s) noexcept
-> span_t<std::byte, is_device, detail::extent_as_bytes_value_t<T, E>::value>
auto as_writable_bytes(span<T, is_device, E> s) noexcept
-> span<std::byte, is_device, detail::extent_as_bytes_value_t<T, E>::value>
{
return {reinterpret_cast<std::byte*>(s.data()), s.size_bytes()};
}
Expand Down
Loading

0 comments on commit 92f8871

Please sign in to comment.