span.hpp
1 /*
2  * Copyright (c) 2020-2022, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <rmm/device_buffer.hpp>
20 #include <rmm/device_uvector.hpp>
21 #include <rmm/device_vector.hpp>
22 
23 #include <thrust/detail/raw_pointer_cast.h>
24 #include <thrust/device_vector.h>
25 #include <thrust/host_vector.h>
26 #include <thrust/memory.h>
27 
28 #include <cstddef>
29 #include <limits>
30 #include <type_traits>
31 
32 namespace cudf {
33 
34 constexpr std::size_t dynamic_extent = std::numeric_limits<std::size_t>::max();
35 
36 namespace detail {
37 
41 template <typename T, std::size_t Extent, typename Derived>
42 class span_base {
43  static_assert(Extent == dynamic_extent, "Only dynamic extent is supported");
44 
45  public:
46  using element_type = T;
47  using value_type = std::remove_cv<T>;
48  using size_type = std::size_t;
49  using difference_type = std::ptrdiff_t;
50  using pointer = T*;
51  using iterator = T*;
52  using const_pointer = T const*;
53  using reference = T&;
54  using const_reference = T const&;
55 
56  static constexpr std::size_t extent = Extent;
57 
58  constexpr span_base() noexcept {}
59  constexpr span_base(pointer data, size_type size) : _data(data), _size(size) {}
60  // constexpr span_base(pointer begin, pointer end) : _data(begin), _size(end - begin) {}
61  constexpr span_base(span_base const& other) noexcept = default;
62  constexpr span_base& operator=(span_base const& other) noexcept = default;
63 
64  // not noexcept due to undefined behavior when size = 0
65  constexpr reference front() const { return _data[0]; }
66  // not noexcept due to undefined behavior when size = 0
67  constexpr reference back() const { return _data[_size - 1]; }
68  // not noexcept due to undefined behavior when idx < 0 || idx >= size
69  constexpr reference operator[](size_type idx) const { return _data[idx]; }
70 
71  constexpr iterator begin() const noexcept { return _data; }
72  constexpr iterator end() const noexcept { return _data + _size; }
73  constexpr pointer data() const noexcept { return _data; }
74 
75  [[nodiscard]] constexpr size_type size() const noexcept { return _size; }
76  [[nodiscard]] constexpr size_type size_bytes() const noexcept { return sizeof(T) * _size; }
77  [[nodiscard]] constexpr bool empty() const noexcept { return _size == 0; }
78 
84  constexpr Derived first(size_type count) const noexcept { return Derived(_data, count); }
85 
91  constexpr Derived last(size_type count) const noexcept
92  {
93  return Derived(_data + _size - count, count);
94  }
95 
96  constexpr Derived subspan(size_type offset, size_type count) const noexcept
97  {
98  return Derived(_data + offset, count);
99  }
100 
101  private:
102  pointer _data{nullptr};
103  size_type _size{0};
104 };
105 
106 } // namespace detail
107 
108 // ===== host_span =================================================================================
109 
110 template <typename T>
111 struct is_host_span_supported_container : std::false_type {
112 };
113 
114 template <typename T, typename Alloc>
116  std::vector<T, Alloc>> : std::true_type {
117 };
118 
119 template <typename T, typename Alloc>
121  thrust::host_vector<T, Alloc>> : std::true_type {
122 };
123 
124 template <typename T, typename Alloc>
126  std::basic_string<T, std::char_traits<T>, Alloc>> : std::true_type {
127 };
128 
129 template <typename T, std::size_t Extent = cudf::dynamic_extent>
130 struct host_span : public cudf::detail::span_base<T, Extent, host_span<T, Extent>> {
132  using base::base;
133 
134  constexpr host_span() noexcept : base() {} // required to compile on centos
135 
136  // Constructor from container
137  template <
138  typename C,
139  // Only supported containers of types convertible to T
140  std::enable_if_t<is_host_span_supported_container<C>::value &&
141  std::is_convertible_v<std::remove_pointer_t<decltype(thrust::raw_pointer_cast(
142  std::declval<C&>().data()))> (*)[],
143  T (*)[]>>* = nullptr>
144  constexpr host_span(C& in) : base(in.data(), in.size())
145  {
146  }
147 
148  // Constructor from const container
149  template <
150  typename C,
151  // Only supported containers of types convertible to T
152  std::enable_if_t<is_host_span_supported_container<C>::value &&
153  std::is_convertible_v<std::remove_pointer_t<decltype(thrust::raw_pointer_cast(
154  std::declval<C&>().data()))> (*)[],
155  T (*)[]>>* = nullptr>
156  constexpr host_span(C const& in) : base(in.data(), in.size())
157  {
158  }
159 
160  // Copy construction to support const conversion
161  template <typename OtherT,
162  std::size_t OtherExtent,
163  std::enable_if_t<(Extent == OtherExtent || Extent == dynamic_extent) &&
164  std::is_convertible_v<OtherT (*)[], T (*)[]>,
165  void>* = nullptr>
166  constexpr host_span(const host_span<OtherT, OtherExtent>& other) noexcept
167  : base(other.data(), other.size())
168  {
169  }
170 };
171 
172 // ===== device_span ===============================================================================
173 
174 template <typename T>
175 struct is_device_span_supported_container : std::false_type {
176 };
177 
178 template <typename T, typename Alloc>
180  thrust::device_vector<T, Alloc>> : std::true_type {
181 };
182 
183 template <typename T>
185  rmm::device_vector<T>> : std::true_type {
186 };
187 
188 template <typename T>
190  rmm::device_uvector<T>> : std::true_type {
191 };
192 
193 template <typename T, std::size_t Extent = cudf::dynamic_extent>
194 struct device_span : public cudf::detail::span_base<T, Extent, device_span<T, Extent>> {
196  using base::base;
197 
198  constexpr device_span() noexcept : base() {} // required to compile on centos
199 
200  template <
201  typename C,
202  // Only supported containers of types convertible to T
203  std::enable_if_t<is_device_span_supported_container<C>::value &&
204  std::is_convertible_v<std::remove_pointer_t<decltype(thrust::raw_pointer_cast(
205  std::declval<C&>().data()))> (*)[],
206  T (*)[]>>* = nullptr>
207  constexpr device_span(C& in) : base(thrust::raw_pointer_cast(in.data()), in.size())
208  {
209  }
210 
211  template <
212  typename C,
213  // Only supported containers of types convertible to T
214  std::enable_if_t<is_device_span_supported_container<C>::value &&
215  std::is_convertible_v<std::remove_pointer_t<decltype(thrust::raw_pointer_cast(
216  std::declval<C&>().data()))> (*)[],
217  T (*)[]>>* = nullptr>
218  constexpr device_span(C const& in) : base(thrust::raw_pointer_cast(in.data()), in.size())
219  {
220  }
221 
222  template <typename OtherT,
223  std::size_t OtherExtent,
224  std::enable_if_t<(Extent == OtherExtent || Extent == dynamic_extent) &&
225  std::is_convertible_v<OtherT (*)[], T (*)[]>,
226  void>* = nullptr>
227  constexpr device_span(const device_span<OtherT, OtherExtent>& other) noexcept
228  : base(other.data(), other.size())
229  {
230  }
231 };
232 
233 namespace detail {
234 
240 template <typename T, template <typename, std::size_t> typename RowType>
241 class base_2dspan {
242  public:
243  using size_type = std::pair<size_t, size_t>;
244 
245  constexpr base_2dspan() noexcept = default;
246  constexpr base_2dspan(T* data, size_t rows, size_t columns) noexcept
247  : _data{data}, _size{rows, columns}
248  {
249  }
250  base_2dspan(T* data, size_type size) noexcept : _data{data}, _size{size} {}
251 
252  constexpr auto data() const noexcept { return _data; }
253  constexpr auto size() const noexcept { return _size; }
254  constexpr auto count() const noexcept { return size().first * size().second; }
255  [[nodiscard]] constexpr bool is_empty() const noexcept { return count() == 0; }
256 
257  static constexpr size_t flatten_index(size_t row, size_t column, size_type size) noexcept
258  {
259  return row * size.second + column;
260  }
261 
262  constexpr RowType<T, dynamic_extent> operator[](size_t row) const
263  {
264  return {this->data() + flatten_index(row, 0, this->size()), this->size().second};
265  }
266 
267  [[nodiscard]] constexpr RowType<T, dynamic_extent> front() const { return (*this)[0]; }
268  [[nodiscard]] constexpr RowType<T, dynamic_extent> back() const
269  {
270  return (*this)[size().first - 1];
271  }
272 
273  constexpr base_2dspan subspan(size_t first_row, size_t num_rows) const noexcept
274  {
275  return base_2dspan(
276  _data + flatten_index(first_row, 0, this->size()), num_rows, this->size().second);
277  }
278 
279  constexpr RowType<T, dynamic_extent> flat_view()
280  {
281  return {this->data(), this->size().first * this->size().second};
282  }
283 
284  template <typename OtherT,
285  template <typename, size_t>
286  typename OtherRowType,
287  std::enable_if_t<std::is_convertible_v<OtherRowType<OtherT, dynamic_extent>,
288  RowType<T, dynamic_extent>>,
289  void>* = nullptr>
290  constexpr base_2dspan(base_2dspan<OtherT, OtherRowType> const& other) noexcept
291  : _data{other.data()}, _size{other.size()}
292  {
293  }
294 
295  protected:
296  T* _data = nullptr;
297  size_type _size{0, 0};
298 };
299 
305 template <class T>
307 
313 template <class T>
315 
316 } // namespace detail
317 } // namespace cudf
cudf::detail::span_base::last
constexpr Derived last(size_type count) const noexcept
Obtains a subspan consisting of the last N elements of the sequence.
Definition: span.hpp:91
cudf::column
A container of nullable device data as a column of elements.
Definition: column.hpp:45
cudf::size_type
int32_t size_type
Row index type for columns and tables.
Definition: types.hpp:84
cudf::detail::span_base::first
constexpr Derived first(size_type count) const noexcept
Obtains a subspan consisting of the first N elements of the sequence.
Definition: span.hpp:84
cudf::detail::base_2dspan
Generic class for row-major 2D spans. Not compliant with STL container semantics/syntax.
Definition: span.hpp:241
cudf::host_span
Definition: span.hpp:130
rmm
cudf::detail::span_base
C++20 std::span with reduced feature set.
Definition: span.hpp:42
device_buffer.hpp
cudf::is_host_span_supported_container
Definition: span.hpp:111
cudf
cuDF interfaces
Definition: aggregation.hpp:34
cudf::device_span
Definition: span.hpp:194
cudf::is_device_span_supported_container
Definition: span.hpp:175