13#include <Kokkos_Core.hpp>
15#include "ddc/chunk_common.hpp"
16#include "ddc/detail/kokkos.hpp"
17#include "ddc/detail/type_seq.hpp"
18#include "ddc/discrete_domain.hpp"
19#include "ddc/discrete_element.hpp"
23template <
class,
class,
class>
29 class LayoutStridedPolicy = Kokkos::layout_right,
30 class MemorySpace = Kokkos::DefaultHostExecutionSpace::memory_space>
43template <
class ElementType,
class... DDims,
class LayoutStridedPolicy,
class MemorySpace>
48 std::is_same_v<LayoutStridedPolicy, Kokkos::layout_left>
49 || std::is_same_v<LayoutStridedPolicy, Kokkos::layout_right>
50 || std::is_same_v<LayoutStridedPolicy, Kokkos::layout_stride>,
51 "ChunkSpan only supports layout_left, layout_right or layout_stride");
57 using typename base_type::internal_mdspan_type;
71 using discrete_domain_type =
typename base_type::discrete_domain_type;
73 using memory_space = MemorySpace;
76 using allocation_mdspan_type =
typename base_type::allocation_mdspan_type;
78 using const_allocation_mdspan_type =
typename base_type::const_allocation_mdspan_type;
80 using discrete_element_type =
typename discrete_domain_type::discrete_element_type;
82 using extents_type =
typename base_type::extents_type;
84 using layout_type =
typename base_type::layout_type;
86 using accessor_type =
typename base_type::accessor_type;
88 using mapping_type =
typename base_type::mapping_type;
90 using element_type =
typename base_type::element_type;
92 using value_type =
typename base_type::value_type;
94 using size_type =
typename base_type::size_type;
96 using data_handle_type =
typename base_type::data_handle_type;
98 using reference =
typename base_type::reference;
100 template <
class,
class,
class,
class>
105 allocation_mdspan_type
const& allocation_mdspan,
106 discrete_domain_type
const& domain)
108 if (!domain.empty()) {
109 extents_type
const extents_s((front<DDims>(domain) + extents<DDims>(domain)).uid()...);
110 std::array<std::size_t,
sizeof...(DDims)>
const strides_s {
111 allocation_mdspan.mapping().stride(
112 type_seq_rank_v<DDims, detail::TypeSeq<DDims...>>)...};
113 Kokkos::layout_stride::mapping<extents_type>
const mapping_s(extents_s, strides_s);
114 return internal_mdspan_type(
115 allocation_mdspan.data_handle() - mapping_s(front<DDims>(domain).uid()...),
119 return internal_mdspan_type(allocation_mdspan);
122 template <
class QueryDDim,
class... ODDims>
125 DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
126 if constexpr (in_tags_v<QueryDDim, detail::TypeSeq<ODDims...>>) {
127 return (uid<QueryDDim>(c) - front<QueryDDim>(
this->m_domain).uid());
129 return Kokkos::full_extent;
134 template <
class QueryDDim,
class... ODDims>
137 DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
138 if constexpr (in_tags_v<QueryDDim, detail::TypeSeq<ODDims...>>) {
139 return std::pair<std::size_t, std::size_t>(
140 front<QueryDDim>(c) - front<QueryDDim>(
this->m_domain),
141 back<QueryDDim>(c) + 1 - front<QueryDDim>(
this->m_domain));
143 return Kokkos::full_extent;
153
154
158
159
163
167 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
168 ChunkSpan(
Chunk<OElementType, discrete_domain_type, Allocator>&& other)
noexcept =
delete;
171
172
176 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
178 Chunk<OElementType, discrete_domain_type, Allocator>& other)
noexcept
179 : base_type(other.m_internal_mdspan, other.m_domain)
184
185
189 class SFINAEElementType = ElementType,
190 class = std::enable_if_t<std::is_const_v<SFINAEElementType>>,
192 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
194 Chunk<OElementType, discrete_domain_type, Allocator>
const& other)
noexcept
195 : base_type(other.m_internal_mdspan, other.m_domain)
200
201
202 template <
class OElementType>
204 ChunkSpan<OElementType, discrete_domain_type, layout_type, MemorySpace>
const&
206 : base_type(other.m_internal_mdspan, other.m_domain)
211
212
213
215 class Mapping = mapping_type,
216 std::enable_if_t<std::is_constructible_v<Mapping, extents_type>,
int> = 0>
218 : base_type(ptr, domain)
223
224
225
227 allocation_mdspan_type allocation_mdspan,
228 discrete_domain_type
const& domain)
231 assert(((allocation_mdspan.extent(type_seq_rank_v<DDims, detail::TypeSeq<DDims...>>)
232 ==
static_cast<std::size_t>(domain.
template extent<DDims>().value()))
237
238
239
240 template <
class KokkosView,
class = std::enable_if_t<Kokkos::is_view_v<KokkosView>>>
242 KokkosView
const& view,
243 discrete_domain_type
const& domain)
noexcept
245 detail::build_mdspan(view, std::make_index_sequence<
sizeof...(DDims)> {}),
253
254
255
259
260
261
265
266 template <
class... QueryDDims>
268 DiscreteElement<QueryDDims...>
const& slice_spec)
const
270 auto subview = Kokkos::submdspan(
allocation_mdspan(), get_slicer_for<DDims>(slice_spec)...);
271 using layout_type =
typename decltype(subview)::layout_type;
272 using extents_type =
typename decltype(subview)::extents_type;
273 using detail::TypeSeq;
274 using OutTypeSeqDDims = type_seq_remove_t<TypeSeq<DDims...>, TypeSeq<QueryDDims...>>;
275 using OutDDom = detail::convert_type_seq_to_discrete_domain_t<OutTypeSeqDDims>;
277 std::is_same_v<layout_type, Kokkos::Experimental::layout_left_padded<>>
278 || std::is_same_v<layout_type, Kokkos::Experimental::layout_right_padded<>>) {
279 Kokkos::layout_stride::mapping<extents_type>
const mapping_stride(subview.mapping());
280 Kokkos::mdspan<ElementType, extents_type, Kokkos::layout_stride>
const
281 a(subview.data_handle(), mapping_stride);
285 Kokkos::layout_stride,
286 memory_space>(a, OutDDom(
this->m_domain));
292 memory_space>(subview, OutDDom(
this->m_domain));
297
298 template <
class... QueryDDims>
302 using layout_type =
typename decltype(subview)::layout_type;
303 using extents_type =
typename decltype(subview)::extents_type;
305 std::is_same_v<layout_type, Kokkos::Experimental::layout_left_padded<>>
306 || std::is_same_v<layout_type, Kokkos::Experimental::layout_right_padded<>>) {
307 Kokkos::layout_stride::mapping<extents_type>
const mapping_stride(subview.mapping());
308 Kokkos::mdspan<ElementType, extents_type, Kokkos::layout_stride>
const
309 a(subview.data_handle(), mapping_stride);
312 decltype(
this->m_domain.restrict_with(odomain)),
313 Kokkos::layout_stride,
314 memory_space>(a,
this->m_domain.restrict_with(odomain));
318 decltype(
this->m_domain.restrict_with(odomain)),
320 memory_space>(subview,
this->m_domain.restrict_with(odomain));
325
326
327
328 template <
class... DElems>
332 sizeof...(DDims) == (0 + ... + DElems::size()),
333 "Invalid number of dimensions");
334 static_assert((is_discrete_element_v<DElems> && ...),
"Expected DiscreteElements");
335 assert(((DiscreteElement<DDims>(take<DDims>(delems...)) >= front<DDims>(
this->m_domain))
337 assert(((DiscreteElement<DDims>(take<DDims>(delems...)) <= back<DDims>(
this->m_domain))
339 return DDC_MDSPAN_ACCESS_OP(
this->m_internal_mdspan, uid<DDims>(take<DDims>(delems...))...);
343
344
347 return base_type::data_handle();
351
352
355 return base_type::allocation_mdspan();
359
360
363 auto s =
this->allocation_mdspan();
364 auto kokkos_layout = detail::build_kokkos_layout(
367 std::make_index_sequence<
sizeof...(DDims)> {});
369 detail::mdspan_to_kokkos_element_t<ElementType,
sizeof...(DDims)>,
370 decltype(kokkos_layout),
371 MemorySpace>(s.data_handle(), kokkos_layout);
376 return view_type(*
this);
385template <
class ElementType,
class SupportType,
class LayoutStridedPolicy,
class MemorySpace>
389 std::is_same_v<LayoutStridedPolicy, Kokkos::layout_left>
390 || std::is_same_v<LayoutStridedPolicy, Kokkos::layout_right>
391 || std::is_same_v<LayoutStridedPolicy, Kokkos::layout_stride>,
392 "ChunkSpan only supports layout_left, layout_right or layout_stride");
395 using base_type =
ChunkCommon<ElementType, SupportType, LayoutStridedPolicy>;
399 using span_type =
ChunkSpan<ElementType, SupportType, LayoutStridedPolicy, MemorySpace>;
402 using view_type =
ChunkSpan<ElementType
const, SupportType, LayoutStridedPolicy, MemorySpace>;
404 using discrete_domain_type =
typename base_type::discrete_domain_type;
406 using memory_space = MemorySpace;
409 using allocation_mdspan_type =
typename base_type::allocation_mdspan_type;
411 using const_allocation_mdspan_type =
typename base_type::const_allocation_mdspan_type;
413 using discrete_element_type =
typename discrete_domain_type::discrete_element_type;
415 using extents_type =
typename base_type::extents_type;
417 using layout_type =
typename base_type::layout_type;
419 using accessor_type =
typename base_type::accessor_type;
421 using mapping_type =
typename base_type::mapping_type;
423 using element_type =
typename base_type::element_type;
425 using value_type =
typename base_type::value_type;
427 using size_type =
typename base_type::size_type;
429 using data_handle_type =
typename base_type::data_handle_type;
431 using reference =
typename base_type::reference;
433 template <
class,
class,
class,
class>
437 template <
class QueryDDim,
class... ODDims>
440 DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
441 if constexpr (in_tags_v<QueryDDim, detail::TypeSeq<ODDims...>>) {
442 return c.
template get<QueryDDim>();
444 return Kokkos::full_extent;
449 template <
class TypeSeq>
452 template <
class... DDims>
453 struct slicer<detail::TypeSeq<DDims...>>
455 template <
class... ODDims>
456 KOKKOS_FUNCTION
constexpr auto operator()(
457 allocation_mdspan_type
const& span,
460 return Kokkos::submdspan(span, get_slicer_for<DDims>(c)...);
469
470
474
475
479
483 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
484 ChunkSpan(
Chunk<OElementType, SupportType, Allocator>&& other)
noexcept =
delete;
487
488
492 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
494 Chunk<OElementType, SupportType, Allocator>& other)
noexcept
495 : base_type(other.m_allocation_mdspan, other.m_domain)
500
501
505 class SFINAEElementType = ElementType,
506 class = std::enable_if_t<std::is_const_v<SFINAEElementType>>,
508 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
510 Chunk<OElementType, SupportType, Allocator>
const& other)
noexcept
511 : base_type(other.m_allocation_mdspan, other.m_domain)
516
517
518 template <
class OElementType>
520 ChunkSpan<OElementType, SupportType, layout_type, MemorySpace>
const& other)
noexcept
521 : base_type(other.m_allocation_mdspan, other.m_domain)
526
527
528
530 class Mapping = mapping_type,
531 std::enable_if_t<std::is_constructible_v<Mapping, extents_type>,
int> = 0>
533 : base_type(ptr, domain)
538
539
540
542 allocation_mdspan_type allocation_mdspan,
543 SupportType
const& domain)
544 : base_type(allocation_mdspan, domain)
546 for (std::size_t i = 0; i < SupportType::rank(); ++i) {
547 assert(allocation_mdspan.extent(i)
548 ==
static_cast<std::size_t>(detail::array(domain.extents())[i]));
553
554
555
556 template <
class KokkosView,
class = std::enable_if_t<Kokkos::is_view_v<KokkosView>>>
559 detail::build_mdspan(view, std::make_index_sequence<SupportType::rank()> {}),
567
568
569
573
574
575
579
580 template <
class... QueryDDims>
582 DiscreteElement<QueryDDims...>
const& slice_spec)
const
584 slicer<to_type_seq_t<SupportType>>
const slicer;
586 = slicer(
this->allocation_mdspan(),
this->m_domain.distance_from_front(slice_spec));
587 using layout_type =
typename decltype(subview)::layout_type;
588 using extents_type =
typename decltype(subview)::extents_type;
589 using detail::TypeSeq;
590 using OutTypeSeqDDims
591 = type_seq_remove_t<to_type_seq_t<SupportType>, TypeSeq<QueryDDims...>>;
592 using OutDDom =
typename detail::RebindDomain<SupportType, OutTypeSeqDDims>::type;
594 std::is_same_v<layout_type, Kokkos::Experimental::layout_left_padded<>>
595 || std::is_same_v<layout_type, Kokkos::Experimental::layout_right_padded<>>) {
596 Kokkos::layout_stride::mapping<extents_type>
const mapping_stride(subview.mapping());
597 Kokkos::mdspan<ElementType, extents_type, Kokkos::layout_stride>
const
598 a(subview.data_handle(), mapping_stride);
602 Kokkos::layout_stride,
603 memory_space>(a, OutDDom(
this->m_domain));
609 memory_space>(subview, OutDDom(
this->m_domain));
614
615
616
617 template <
class... DElems>
621 SupportType::rank() == (0 + ... + DElems::size()),
622 "Invalid number of dimensions");
623 static_assert((is_discrete_element_v<DElems> && ...),
"Expected DiscreteElements");
624 assert(
this->m_domain.contains(delems...));
625 return DDC_MDSPAN_ACCESS_OP(
626 this->m_allocation_mdspan,
627 detail::array(
this->m_domain.distance_from_front(delems...)));
631
632
635 return base_type::data_handle();
639
640
643 return base_type::allocation_mdspan();
647
648
651 auto s =
this->allocation_mdspan();
652 auto kokkos_layout = detail::build_kokkos_layout(
655 std::make_index_sequence<SupportType::rank()> {});
657 detail::mdspan_to_kokkos_element_t<ElementType, SupportType::rank()>,
658 decltype(kokkos_layout),
659 MemorySpace>(s.data_handle(), kokkos_layout);
664 return view_type(*
this);
673template <
class DataType,
class... Properties,
class SupportType>
675 Kokkos::View<DataType, Properties...>
const& view,
678 detail::kokkos_to_mdspan_element_t<
679 typename Kokkos::View<DataType, Properties...>::
data_type>,
681 detail::kokkos_to_mdspan_layout_t<
682 typename Kokkos::View<DataType, Properties...>::
array_layout>,
683 typename Kokkos::View<DataType, Properties...>::
memory_space>;
685template <
class ElementType,
class SupportType,
class Allocator>
690 Kokkos::layout_right,
693template <
class ElementType,
class SupportType,
class Allocator>
698 Kokkos::layout_right,
704 class LayoutStridedPolicy = Kokkos::layout_right,
705 class MemorySpace = Kokkos::HostSpace>
706using ChunkView =
ChunkSpan<ElementType
const, SupportType, LayoutStridedPolicy, MemorySpace>;
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 auto allocation_kokkos_view() const
Provide a mdspan on the memory allocation.
KOKKOS_FUNCTION constexpr ElementType * data_handle() const
Access to the underlying allocation pointer.
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan()=default
Empty ChunkSpan.
KOKKOS_FUNCTION constexpr reference operator()(DElems const &... delems) const noexcept
Element access using a list of DiscreteElement.
KOKKOS_FUNCTION constexpr ChunkSpan(Chunk< OElementType, discrete_domain_type, Allocator > &other) noexcept
Constructs a new ChunkSpan from a Chunk, yields a new view to the same data.
ChunkSpan(Chunk< OElementType, discrete_domain_type, Allocator > &&other) noexcept=delete
Forbids to construct a ChunkSpan from a rvalue of type Chunk.
KOKKOS_FUNCTION constexpr auto get_slicer_for(DiscreteElement< ODDims... > const &c) const
KOKKOS_FUNCTION constexpr span_type span_view() const
KOKKOS_FUNCTION constexpr allocation_mdspan_type allocation_mdspan() const
Provide a mdspan on the memory allocation.
KOKKOS_FUNCTION constexpr auto operator[](DiscreteDomain< QueryDDims... > const &odomain) const
Restrict to a subdomain.
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan(ChunkSpan &&other) noexcept=default
Constructs a new ChunkSpan by move.
KOKKOS_FUNCTION constexpr ChunkSpan(KokkosView const &view, discrete_domain_type const &domain) noexcept
Constructs a new ChunkSpan from scratch.
static KOKKOS_FUNCTION internal_mdspan_type build_internal_mdspan(allocation_mdspan_type const &allocation_mdspan, discrete_domain_type const &domain)
KOKKOS_DEFAULTED_FUNCTION ~ChunkSpan() noexcept=default
KOKKOS_FUNCTION constexpr ChunkSpan(Chunk< OElementType, discrete_domain_type, Allocator > const &other) noexcept
Constructs a new ChunkSpan from a Chunk, yields a new view to the same data.
KOKKOS_FUNCTION constexpr auto operator[](DiscreteElement< QueryDDims... > const &slice_spec) const
Slice out some dimensions.
KOKKOS_FUNCTION constexpr ChunkSpan(ElementType *const ptr, discrete_domain_type const &domain)
Constructs a new ChunkSpan from scratch.
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan & operator=(ChunkSpan &&other) noexcept=default
Move-assigns a new value to this ChunkSpan.
KOKKOS_FUNCTION constexpr ChunkSpan(ChunkSpan< OElementType, discrete_domain_type, layout_type, MemorySpace > const &other) noexcept
Constructs a new ChunkSpan by copy of a chunk, yields a new view to the same data.
KOKKOS_FUNCTION constexpr ChunkSpan(allocation_mdspan_type allocation_mdspan, discrete_domain_type const &domain)
Constructs a new ChunkSpan from scratch.
KOKKOS_FUNCTION constexpr auto get_slicer_for(DiscreteDomain< ODDims... > const &c) const
KOKKOS_FUNCTION constexpr view_type span_cview() const
KOKKOS_DEFAULTED_FUNCTION constexpr ChunkSpan(ChunkSpan const &other)=default
Constructs a new ChunkSpan by copy, yields a new view to the same data.
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_DEFAULTED_FUNCTION constexpr ChunkSpan(ChunkSpan const &other)=default
Constructs a new ChunkSpan by copy, yields a new view to the same data.
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_FUNCTION constexpr reference operator()(DElems const &... delems) const noexcept
Element access using a list of DiscreteElement.
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_DEFAULTED_FUNCTION constexpr ChunkSpan & operator=(ChunkSpan &&other) noexcept=default
Move-assigns a new value to this ChunkSpan.
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 > &other) -> ChunkSpan< ElementType, SupportType, Kokkos::layout_right, typename Allocator::memory_space >
constexpr bool enable_chunk< ChunkSpan< ElementType, SupportType, LayoutStridedPolicy, MemorySpace > >
ChunkSpan(Chunk< ElementType, SupportType, Allocator > const &other) -> ChunkSpan< const ElementType, SupportType, Kokkos::layout_right, typename Allocator::memory_space >
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 >