DDC 0.7.0
Loading...
Searching...
No Matches
chunk_span.hpp
1// Copyright (C) The DDC development team, see COPYRIGHT.md file
2//
3// SPDX-License-Identifier: MIT
4
5#pragma once
6
7#include <array>
8#include <cassert>
9#include <cstddef>
10#include <type_traits>
11#include <utility>
12
13#include <Kokkos_Core.hpp>
14
15#include "detail/kokkos.hpp"
16#include "detail/type_seq.hpp"
17#include "detail/type_traits.hpp"
18
19#include "chunk_common.hpp"
20#include "discrete_domain.hpp"
21#include "discrete_element.hpp"
22
23namespace ddc {
24
25template <class, class, class>
26class Chunk;
27
28template <
29 class ElementType,
30 class SupportType,
31 class LayoutStridedPolicy = Kokkos::layout_right,
32 class MemorySpace = Kokkos::DefaultHostExecutionSpace::memory_space>
33class ChunkSpan;
34
35template <class ElementType, class SupportType, class LayoutStridedPolicy, class MemorySpace>
36inline constexpr bool
38 = true;
39
40template <class ElementType, class SupportType, class LayoutStridedPolicy, class MemorySpace>
41inline constexpr bool
43 = true;
44
45template <class ElementType, class SupportType, class LayoutStridedPolicy, class MemorySpace>
46class ChunkSpan : public ChunkCommon<ElementType, SupportType, LayoutStridedPolicy>
47{
48 static_assert(
49 std::is_same_v<LayoutStridedPolicy, Kokkos::layout_left>
50 || std::is_same_v<LayoutStridedPolicy, Kokkos::layout_right>
51 || std::is_same_v<LayoutStridedPolicy, Kokkos::layout_stride>,
52 "ChunkSpan only supports layout_left, layout_right or layout_stride");
53
54protected:
55 using base_type = ChunkCommon<ElementType, SupportType, LayoutStridedPolicy>;
56
57public:
58 /// type of a span of this full chunk
59 using span_type = ChunkSpan<ElementType, SupportType, LayoutStridedPolicy, MemorySpace>;
60
61 /// type of a view of this full chunk
62 using view_type = ChunkSpan<ElementType const, SupportType, LayoutStridedPolicy, MemorySpace>;
63
64 using discrete_domain_type = typename base_type::discrete_domain_type;
65
66 using memory_space = MemorySpace;
67
68 /// The dereferenceable part of the co-domain but with a different domain, starting at 0
69 using allocation_mdspan_type = typename base_type::allocation_mdspan_type;
70
71 using const_allocation_mdspan_type = typename base_type::const_allocation_mdspan_type;
72
73 using discrete_element_type = typename discrete_domain_type::discrete_element_type;
74
75 using discrete_vector_type = typename discrete_domain_type::discrete_vector_type;
76
77 using extents_type = typename base_type::extents_type;
78
79 using layout_type = typename base_type::layout_type;
80
81 using accessor_type = typename base_type::accessor_type;
82
83 using mapping_type = typename base_type::mapping_type;
84
85 using element_type = typename base_type::element_type;
86
87 using value_type = typename base_type::value_type;
88
89 using size_type = typename base_type::size_type;
90
91 using data_handle_type = typename base_type::data_handle_type;
92
93 using reference = typename base_type::reference;
94
95 template <class, class, class, class>
96 friend class ChunkSpan;
97
98protected:
99 template <class QueryDDim, class... ODDims>
100 KOKKOS_FUNCTION static constexpr auto get_slicer_for(DiscreteVector<ODDims...> const& c)
101 {
102 DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
103 if constexpr (in_tags_v<QueryDDim, detail::TypeSeq<ODDims...>>) {
104 return c.template get<QueryDDim>();
105 } else {
106 return Kokkos::full_extent;
107 }
108 DDC_IF_NVCC_THEN_POP
109 }
110
111 template <class QueryDDim, class... ODDims, class... OODDims>
112 KOKKOS_FUNCTION static constexpr auto get_slicer_for(
113 DiscreteDomain<ODDims...> const& c,
114 DiscreteDomain<OODDims...> const& origin)
115 {
116 DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
117 if constexpr (in_tags_v<QueryDDim, detail::TypeSeq<ODDims...>>) {
118 return std::pair<std::size_t, std::size_t>(
119 front<QueryDDim>(c) - front<QueryDDim>(origin),
120 back<QueryDDim>(c) + 1 - front<QueryDDim>(origin));
121 } else {
122 return Kokkos::full_extent;
123 }
124 DDC_IF_NVCC_THEN_POP
125 }
126
127 template <class TypeSeq>
128 struct slicer;
129
130 template <class... DDims>
131 struct slicer<detail::TypeSeq<DDims...>>
132 {
133 template <class... ODDims>
134 KOKKOS_FUNCTION constexpr auto operator()(
135 allocation_mdspan_type const& span,
136 DiscreteVector<ODDims...> const& c) const
137 {
138 return Kokkos::submdspan(span, get_slicer_for<DDims>(c)...);
139 }
140
141 template <class... ODDims, class... OODDims>
142 KOKKOS_FUNCTION constexpr auto operator()(
143 allocation_mdspan_type const& span,
144 DiscreteDomain<ODDims...> const& c,
145 DiscreteDomain<OODDims...> const& origin) const
146 {
147 return Kokkos::submdspan(span, get_slicer_for<DDims>(c, origin)...);
148 }
149 };
150
151public:
152 /// Empty ChunkSpan
153 KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan() = default;
154
155 /** Constructs a new ChunkSpan by copy, yields a new view to the same data
156 * @param other the ChunkSpan to copy
157 */
158 KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan(ChunkSpan const& other) = default;
159
160 /** Constructs a new ChunkSpan by move
161 * @param other the ChunkSpan to move
162 */
163 KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan(ChunkSpan&& other) noexcept = default;
164
165 /** Forbids to construct a ChunkSpan from a rvalue of type Chunk.
166 */
167 template <
168 class OElementType,
169 class Allocator,
170 class = std::enable_if_t<std::is_same_v<typename Allocator::memory_space, MemorySpace>>>
171 ChunkSpan(Chunk<OElementType, SupportType, Allocator>&& other) noexcept = delete;
172
173 /** Constructs a new ChunkSpan from a Chunk, yields a new view to the same data
174 * @param other the Chunk to view
175 */
176 template <
177 class OElementType,
178 class Allocator,
179 class = std::enable_if_t<std::is_same_v<typename Allocator::memory_space, MemorySpace>>>
180 KOKKOS_FUNCTION constexpr explicit ChunkSpan(
181 Chunk<OElementType, SupportType, Allocator>& other) noexcept
182 : base_type(other.m_allocation_mdspan, other.m_domain)
183 {
184 }
185
186 /** Constructs a new ChunkSpan from a Chunk, yields a new view to the same data
187 * @param other the Chunk to view
188 */
189 // Disabled by SFINAE in the case of `ElementType` is not `const` to avoid write access
190 template <
191 class OElementType,
192 class SFINAEElementType = ElementType,
193 class = std::enable_if_t<std::is_const_v<SFINAEElementType>>,
194 class Allocator,
195 class = std::enable_if_t<std::is_same_v<typename Allocator::memory_space, MemorySpace>>>
196 KOKKOS_FUNCTION constexpr explicit ChunkSpan(
197 Chunk<OElementType, SupportType, Allocator> const& other) noexcept
198 : base_type(other.m_allocation_mdspan, other.m_domain)
199 {
200 }
201
202 /** Constructs a new ChunkSpan by copy of a chunk, yields a new view to the same data
203 * @param other the ChunkSpan to move
204 */
205 template <class OElementType>
206 KOKKOS_FUNCTION constexpr explicit ChunkSpan(
207 ChunkSpan<OElementType, SupportType, layout_type, MemorySpace> const& other) noexcept
208 : base_type(other.m_allocation_mdspan, other.m_domain)
209 {
210 }
211
212 /** Constructs a new ChunkSpan from scratch
213 * @param ptr the allocation pointer to the data
214 * @param domain the domain that sustains the view
215 */
216 template <
217 class Mapping = mapping_type,
218 std::enable_if_t<std::is_constructible_v<Mapping, extents_type>, int> = 0>
219 KOKKOS_FUNCTION constexpr ChunkSpan(ElementType* const ptr, SupportType const& domain)
220 : base_type(ptr, domain)
221 {
222 }
223
224 /** Constructs a new ChunkSpan from scratch
225 * @param allocation_mdspan the allocation mdspan to the data
226 * @param domain the domain that sustains the view
227 */
228 KOKKOS_FUNCTION constexpr ChunkSpan(
229 allocation_mdspan_type allocation_mdspan,
230 SupportType const& domain)
231 : base_type(allocation_mdspan, domain)
232 {
233 for (std::size_t i = 0; i < SupportType::rank(); ++i) {
234 assert(allocation_mdspan.extent(i)
235 == static_cast<std::size_t>(detail::array(domain.extents())[i]));
236 }
237 }
238
239 /** Constructs a new ChunkSpan from scratch
240 * @param view the Kokkos view
241 * @param domain the domain that sustains the view
242 */
243 template <class KokkosView, class = std::enable_if_t<Kokkos::is_view_v<KokkosView>>>
244 KOKKOS_FUNCTION constexpr ChunkSpan(KokkosView const& view, SupportType const& domain) noexcept
245 : ChunkSpan(
246 detail::build_mdspan(view, std::make_index_sequence<SupportType::rank()> {}),
247 domain)
248 {
249 }
250
251 KOKKOS_DEFAULTED_FUNCTION ~ChunkSpan() noexcept = default;
252
253 /** Copy-assigns a new value to this ChunkSpan, yields a new view to the same data
254 * @param other the ChunkSpan to copy
255 * @return *this
256 */
257 KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan& operator=(ChunkSpan const& other) = default;
258
259 /** Move-assigns a new value to this ChunkSpan
260 * @param other the ChunkSpan to move
261 * @return *this
262 */
263 KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan& operator=(ChunkSpan&& other) noexcept = default;
264
265 /** Slice out some dimensions
266 */
267 template <class... QueryDDims>
268 KOKKOS_FUNCTION constexpr auto operator[](
269 DiscreteElement<QueryDDims...> const& slice_spec) const
270 {
271 assert(select<QueryDDims...>(this->m_domain).contains(slice_spec));
272 slicer<to_type_seq_t<SupportType>> const slicer;
273 auto subview = slicer(
274 this->allocation_mdspan(),
275 ddc::DiscreteDomain<QueryDDims...>(this->m_domain).distance_from_front(slice_spec));
276 using layout_type = typename decltype(subview)::layout_type;
277 using extents_type = typename decltype(subview)::extents_type;
278 using detail::TypeSeq;
279 using OutTypeSeqDDims
280 = type_seq_remove_t<to_type_seq_t<SupportType>, TypeSeq<QueryDDims...>>;
281 using OutDDom = typename detail::RebindDomain<SupportType, OutTypeSeqDDims>::type;
282 if constexpr (
283 std::is_same_v<layout_type, Kokkos::Experimental::layout_left_padded<>>
284 || std::is_same_v<layout_type, Kokkos::Experimental::layout_right_padded<>>) {
285 Kokkos::layout_stride::mapping<extents_type> const mapping_stride(subview.mapping());
286 Kokkos::mdspan<ElementType, extents_type, Kokkos::layout_stride> const
287 a(subview.data_handle(), mapping_stride);
288 return ChunkSpan<
289 ElementType,
290 OutDDom,
291 Kokkos::layout_stride,
292 memory_space>(a, OutDDom(this->m_domain));
293 } else {
294 return ChunkSpan<
295 ElementType,
296 OutDDom,
297 layout_type,
298 memory_space>(subview, OutDDom(this->m_domain));
299 }
300 }
301
302 /** Restrict to a subdomain, only valid when SupportType is a DiscreteDomain
303 */
304 template <
305 class... QueryDDims,
306 class SFINAESupportType = SupportType,
307 std::enable_if_t<is_discrete_domain_v<SFINAESupportType>, std::nullptr_t> = nullptr>
308 KOKKOS_FUNCTION constexpr auto operator[](DiscreteDomain<QueryDDims...> const& odomain) const
309 {
310 assert(odomain.empty()
311 || (DiscreteDomain<QueryDDims...>(this->m_domain).contains(odomain.front())
312 && DiscreteDomain<QueryDDims...>(this->m_domain).contains(odomain.back())));
313 slicer<to_type_seq_t<SupportType>> const slicer;
314 auto subview = slicer(this->allocation_mdspan(), odomain, this->m_domain);
315 using layout_type = typename decltype(subview)::layout_type;
316 using extents_type = typename decltype(subview)::extents_type;
317 if constexpr (
318 std::is_same_v<layout_type, Kokkos::Experimental::layout_left_padded<>>
319 || std::is_same_v<layout_type, Kokkos::Experimental::layout_right_padded<>>) {
320 Kokkos::layout_stride::mapping<extents_type> const mapping_stride(subview.mapping());
321 Kokkos::mdspan<ElementType, extents_type, Kokkos::layout_stride> const
322 a(subview.data_handle(), mapping_stride);
323 return ChunkSpan<
324 ElementType,
325 decltype(this->m_domain.restrict_with(odomain)),
326 Kokkos::layout_stride,
327 memory_space>(a, this->m_domain.restrict_with(odomain));
328 } else {
329 return ChunkSpan<
330 ElementType,
331 decltype(this->m_domain.restrict_with(odomain)),
332 layout_type,
333 memory_space>(subview, this->m_domain.restrict_with(odomain));
334 }
335 }
336
337 /** Element access using a list of DiscreteElement
338 * @param delems discrete elements
339 * @return reference to this element
340 */
341 template <
342 class... DElems,
343 std::enable_if_t<detail::all_of_v<is_discrete_element_v<DElems>...>, int> = 0>
344 KOKKOS_FUNCTION constexpr reference operator()(DElems const&... delems) const noexcept
345 {
346 static_assert(
347 SupportType::rank() == (0 + ... + DElems::size()),
348 "Invalid number of dimensions");
349 assert(this->m_domain.contains(delems...));
350 return DDC_MDSPAN_ACCESS_OP(
351 this->m_allocation_mdspan,
352 detail::array(this->m_domain.distance_from_front(delems...)));
353 }
354
355 /** Element access using a list of DiscreteVector
356 * @param dvects discrete vectors
357 * @return reference to this element
358 */
359 template <
360 class... DVects,
361 std::enable_if_t<detail::all_of_v<is_discrete_vector_v<DVects>...>, int> = 0,
362 std::enable_if_t<sizeof...(DVects) != 0, int> = 0>
363 KOKKOS_FUNCTION constexpr reference operator()(DVects const&... dvects) const noexcept
364 {
365 static_assert(
366 SupportType::rank() == (0 + ... + DVects::size()),
367 "Invalid number of dimensions");
368 return DDC_MDSPAN_ACCESS_OP(
369 this->m_allocation_mdspan,
370 detail::array(discrete_vector_type(dvects...)));
371 }
372
373 /** Access to the underlying allocation pointer
374 * @return allocation pointer
375 */
376 KOKKOS_FUNCTION constexpr ElementType* data_handle() const
377 {
378 return base_type::data_handle();
379 }
380
381 /** Provide a mdspan on the memory allocation
382 * @return allocation mdspan
383 */
384 KOKKOS_FUNCTION constexpr allocation_mdspan_type allocation_mdspan() const
385 {
386 return base_type::allocation_mdspan();
387 }
388
389 /** Provide a mdspan on the memory allocation
390 * @return allocation mdspan
391 */
392 KOKKOS_FUNCTION constexpr auto allocation_kokkos_view() const
393 {
394 auto s = this->allocation_mdspan();
395 auto kokkos_layout = detail::build_kokkos_layout(
396 s.extents(),
397 s.mapping(),
398 std::make_index_sequence<SupportType::rank()> {});
399 return Kokkos::View<
400 detail::mdspan_to_kokkos_element_t<ElementType, SupportType::rank()>,
401 decltype(kokkos_layout),
402 MemorySpace>(s.data_handle(), kokkos_layout);
403 }
404
405 KOKKOS_FUNCTION constexpr view_type span_cview() const
406 {
407 return view_type(*this);
408 }
409
410 KOKKOS_FUNCTION constexpr span_type span_view() const
411 {
412 return *this;
413 }
414};
415
416template <class DataType, class... Properties, class SupportType>
418 Kokkos::View<DataType, Properties...> const& view,
419 SupportType domain)
420 -> ChunkSpan<
421 detail::kokkos_to_mdspan_element_t<
422 typename Kokkos::View<DataType, Properties...>::data_type>,
423 SupportType,
424 detail::kokkos_to_mdspan_layout_t<
425 typename Kokkos::View<DataType, Properties...>::array_layout>,
426 typename Kokkos::View<DataType, Properties...>::memory_space>;
427
428template <class ElementType, class SupportType, class Allocator>
429ChunkSpan(Chunk<ElementType, SupportType, Allocator>& other) -> ChunkSpan<
430 ElementType,
431 SupportType,
432 Kokkos::layout_right,
433 typename Allocator::memory_space>;
434
435template <class ElementType, class SupportType, class Allocator>
436ChunkSpan(Chunk<ElementType, SupportType, Allocator> const& other) -> ChunkSpan<
437 ElementType const,
438 SupportType,
439 Kokkos::layout_right,
440 typename Allocator::memory_space>;
441
442template <
443 class ElementType,
444 class SupportType,
445 class LayoutStridedPolicy = Kokkos::layout_right,
446 class MemorySpace = Kokkos::HostSpace>
447using ChunkView = ChunkSpan<ElementType const, SupportType, LayoutStridedPolicy, MemorySpace>;
448
449} // namespace ddc
friend class ChunkSpan
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan(ChunkSpan &&other) noexcept=default
Constructs a new ChunkSpan by move.
KOKKOS_FUNCTION constexpr view_type span_cview() const
KOKKOS_FUNCTION constexpr ChunkSpan(Chunk< OElementType, SupportType, Allocator > const &other) noexcept
Constructs a new ChunkSpan from a Chunk, yields a new view to the same data.
KOKKOS_FUNCTION constexpr ChunkSpan(allocation_mdspan_type allocation_mdspan, SupportType const &domain)
Constructs a new ChunkSpan from scratch.
KOKKOS_FUNCTION constexpr auto operator[](DiscreteElement< QueryDDims... > const &slice_spec) const
Slice out some dimensions.
KOKKOS_FUNCTION constexpr ChunkSpan(KokkosView const &view, SupportType const &domain) noexcept
Constructs a new ChunkSpan from scratch.
ChunkSpan(Chunk< OElementType, SupportType, Allocator > &&other) noexcept=delete
Forbids to construct a ChunkSpan from a rvalue of type Chunk.
KOKKOS_FUNCTION constexpr auto operator[](DiscreteDomain< QueryDDims... > const &odomain) const
Restrict to a subdomain, only valid when SupportType is a DiscreteDomain.
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan(ChunkSpan const &other)=default
Constructs a new ChunkSpan by copy, yields a new view to the same data.
static KOKKOS_FUNCTION constexpr auto get_slicer_for(DiscreteDomain< ODDims... > const &c, DiscreteDomain< OODDims... > const &origin)
KOKKOS_FUNCTION constexpr ChunkSpan(ElementType *const ptr, SupportType const &domain)
Constructs a new ChunkSpan from scratch.
static KOKKOS_FUNCTION constexpr auto get_slicer_for(DiscreteVector< ODDims... > const &c)
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan()=default
Empty ChunkSpan.
KOKKOS_FUNCTION constexpr allocation_mdspan_type allocation_mdspan() const
Provide a mdspan on the memory allocation.
KOKKOS_FUNCTION constexpr ElementType * data_handle() const
Access to the underlying allocation pointer.
KOKKOS_FUNCTION constexpr auto allocation_kokkos_view() const
Provide a mdspan on the memory allocation.
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan & operator=(ChunkSpan const &other)=default
Copy-assigns a new value to this ChunkSpan, yields a new view to the same data.
KOKKOS_FUNCTION constexpr ChunkSpan(Chunk< OElementType, SupportType, Allocator > &other) noexcept
Constructs a new ChunkSpan from a Chunk, yields a new view to the same data.
KOKKOS_FUNCTION constexpr ChunkSpan(ChunkSpan< OElementType, SupportType, layout_type, MemorySpace > const &other) noexcept
Constructs a new ChunkSpan by copy of a chunk, yields a new view to the same data.
KOKKOS_DEFAULTED_FUNCTION ~ChunkSpan() noexcept=default
KOKKOS_FUNCTION constexpr span_type span_view() const
KOKKOS_FUNCTION constexpr reference operator()(DElems const &... delems) const noexcept
Element access using a list of DiscreteElement.
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan & operator=(ChunkSpan &&other) noexcept=default
Move-assigns a new value to this ChunkSpan.
KOKKOS_FUNCTION constexpr reference operator()(DVects const &... dvects) const noexcept
Element access using a list of DiscreteVector.
friend class Chunk
Definition chunk.hpp:81
friend class DiscreteDomain
KOKKOS_FUNCTION constexpr bool operator!=(DiscreteVector< OTags... > const &rhs) const noexcept
The top-level namespace of DDC.
constexpr bool enable_borrowed_chunk< ChunkSpan< ElementType, SupportType, LayoutStridedPolicy, MemorySpace > >
ChunkSpan(Chunk< ElementType, SupportType, Allocator > const &other) -> ChunkSpan< ElementType const, SupportType, Kokkos::layout_right, typename Allocator::memory_space >
ChunkSpan(Chunk< ElementType, SupportType, Allocator > &other) -> ChunkSpan< ElementType, SupportType, Kokkos::layout_right, typename Allocator::memory_space >
constexpr bool enable_chunk< ChunkSpan< ElementType, SupportType, LayoutStridedPolicy, MemorySpace > >
KOKKOS_DEDUCTION_GUIDE ChunkSpan(Kokkos::View< DataType, Properties... > const &view, SupportType domain) -> ChunkSpan< detail::kokkos_to_mdspan_element_t< typename Kokkos::View< DataType, Properties... >::data_type >, SupportType, detail::kokkos_to_mdspan_layout_t< typename Kokkos::View< DataType, Properties... >::array_layout >, typename Kokkos::View< DataType, Properties... >::memory_space >