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(
110 (front<DDims>(domain) +
ddc::extents<DDims>(domain)).uid()...);
111 std::array<std::size_t,
sizeof...(DDims)>
const strides_s {
112 allocation_mdspan.mapping().stride(
113 type_seq_rank_v<DDims, detail::TypeSeq<DDims...>>)...};
114 Kokkos::layout_stride::mapping<extents_type>
const mapping_s(extents_s, strides_s);
115 return internal_mdspan_type(
116 allocation_mdspan.data_handle() - mapping_s(front<DDims>(domain).uid()...),
120 return internal_mdspan_type(allocation_mdspan);
123 template <
class QueryDDim,
class... ODDims>
126 DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
127 if constexpr (in_tags_v<QueryDDim, detail::TypeSeq<ODDims...>>) {
128 return (uid<QueryDDim>(c) - front<QueryDDim>(
this->m_domain).uid());
130 return Kokkos::full_extent;
135 template <
class QueryDDim,
class... ODDims>
138 DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
139 if constexpr (in_tags_v<QueryDDim, detail::TypeSeq<ODDims...>>) {
140 return std::pair<std::size_t, std::size_t>(
141 front<QueryDDim>(c) - front<QueryDDim>(
this->m_domain),
142 back<QueryDDim>(c) + 1 - front<QueryDDim>(
this->m_domain));
144 return Kokkos::full_extent;
154
155
159
160
164
168 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
169 ChunkSpan(
Chunk<OElementType, discrete_domain_type, Allocator>&& other)
noexcept =
delete;
172
173
177 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
179 Chunk<OElementType, discrete_domain_type, Allocator>& other)
noexcept
180 : base_type(other.m_internal_mdspan, other.m_domain)
185
186
190 class SFINAEElementType = ElementType,
191 class = std::enable_if_t<std::is_const_v<SFINAEElementType>>,
193 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
195 Chunk<OElementType, discrete_domain_type, Allocator>
const& other)
noexcept
196 : base_type(other.m_internal_mdspan, other.m_domain)
201
202
203 template <
class OElementType>
205 ChunkSpan<OElementType, discrete_domain_type, layout_type, MemorySpace>
const&
207 : base_type(other.m_internal_mdspan, other.m_domain)
212
213
214
216 class Mapping = mapping_type,
217 std::enable_if_t<std::is_constructible_v<Mapping, extents_type>,
int> = 0>
219 : base_type(ptr, domain)
224
225
226
228 allocation_mdspan_type allocation_mdspan,
229 discrete_domain_type
const& domain)
232 assert(((allocation_mdspan.extent(type_seq_rank_v<DDims, detail::TypeSeq<DDims...>>)
233 ==
static_cast<std::size_t>(domain.
template extent<DDims>().value()))
238
239
240
241 template <
class KokkosView,
class = std::enable_if_t<Kokkos::is_view_v<KokkosView>>>
243 KokkosView
const& view,
244 discrete_domain_type
const& domain)
noexcept
246 detail::build_mdspan(view, std::make_index_sequence<
sizeof...(DDims)> {}),
254
255
256
260
261
262
266
267 template <
class... QueryDDims>
269 DiscreteElement<QueryDDims...>
const& slice_spec)
const
271 auto subview = Kokkos::submdspan(
allocation_mdspan(), get_slicer_for<DDims>(slice_spec)...);
272 using layout_type =
typename decltype(subview)::layout_type;
273 using extents_type =
typename decltype(subview)::extents_type;
274 using detail::TypeSeq;
275 using OutTypeSeqDDims = type_seq_remove_t<TypeSeq<DDims...>, TypeSeq<QueryDDims...>>;
276 using OutDDom = detail::convert_type_seq_to_discrete_domain_t<OutTypeSeqDDims>;
278 std::is_same_v<layout_type, Kokkos::Experimental::layout_left_padded<>>
279 || std::is_same_v<layout_type, Kokkos::Experimental::layout_right_padded<>>) {
280 Kokkos::layout_stride::mapping<extents_type>
const mapping_stride(subview.mapping());
281 Kokkos::mdspan<ElementType, extents_type, Kokkos::layout_stride>
const
282 a(subview.data_handle(), mapping_stride);
286 Kokkos::layout_stride,
287 memory_space>(a, OutDDom(
this->m_domain));
293 memory_space>(subview, OutDDom(
this->m_domain));
298
299 template <
class... QueryDDims>
303 using layout_type =
typename decltype(subview)::layout_type;
304 using extents_type =
typename decltype(subview)::extents_type;
306 std::is_same_v<layout_type, Kokkos::Experimental::layout_left_padded<>>
307 || std::is_same_v<layout_type, Kokkos::Experimental::layout_right_padded<>>) {
308 Kokkos::layout_stride::mapping<extents_type>
const mapping_stride(subview.mapping());
309 Kokkos::mdspan<ElementType, extents_type, Kokkos::layout_stride>
const
310 a(subview.data_handle(), mapping_stride);
313 decltype(
this->m_domain.restrict_with(odomain)),
314 Kokkos::layout_stride,
315 memory_space>(a,
this->m_domain.restrict_with(odomain));
319 decltype(
this->m_domain.restrict_with(odomain)),
321 memory_space>(subview,
this->m_domain.restrict_with(odomain));
326
327
328
329 template <
class... DElems>
333 sizeof...(DDims) == (0 + ... + DElems::size()),
334 "Invalid number of dimensions");
335 static_assert((is_discrete_element_v<DElems> && ...),
"Expected DiscreteElements");
336 assert(((DiscreteElement<DDims>(take<DDims>(delems...)) >= front<DDims>(
this->m_domain))
338 assert(((DiscreteElement<DDims>(take<DDims>(delems...)) <= back<DDims>(
this->m_domain))
340 return DDC_MDSPAN_ACCESS_OP(
this->m_internal_mdspan, uid<DDims>(take<DDims>(delems...))...);
344
345
348 return base_type::data_handle();
352
353
356 return base_type::allocation_mdspan();
360
361
364 auto s =
this->allocation_mdspan();
365 auto kokkos_layout = detail::build_kokkos_layout(
368 std::make_index_sequence<
sizeof...(DDims)> {});
370 detail::mdspan_to_kokkos_element_t<ElementType,
sizeof...(DDims)>,
371 decltype(kokkos_layout),
372 MemorySpace>(s.data_handle(), kokkos_layout);
377 return view_type(*
this);
386template <
class ElementType,
class SupportType,
class LayoutStridedPolicy,
class MemorySpace>
390 std::is_same_v<LayoutStridedPolicy, Kokkos::layout_left>
391 || std::is_same_v<LayoutStridedPolicy, Kokkos::layout_right>
392 || std::is_same_v<LayoutStridedPolicy, Kokkos::layout_stride>,
393 "ChunkSpan only supports layout_left, layout_right or layout_stride");
396 using base_type =
ChunkCommon<ElementType, SupportType, LayoutStridedPolicy>;
400 using span_type =
ChunkSpan<ElementType, SupportType, LayoutStridedPolicy, MemorySpace>;
403 using view_type =
ChunkSpan<ElementType
const, SupportType, LayoutStridedPolicy, MemorySpace>;
405 using discrete_domain_type =
typename base_type::discrete_domain_type;
407 using memory_space = MemorySpace;
410 using allocation_mdspan_type =
typename base_type::allocation_mdspan_type;
412 using const_allocation_mdspan_type =
typename base_type::const_allocation_mdspan_type;
414 using discrete_element_type =
typename discrete_domain_type::discrete_element_type;
416 using extents_type =
typename base_type::extents_type;
418 using layout_type =
typename base_type::layout_type;
420 using accessor_type =
typename base_type::accessor_type;
422 using mapping_type =
typename base_type::mapping_type;
424 using element_type =
typename base_type::element_type;
426 using value_type =
typename base_type::value_type;
428 using size_type =
typename base_type::size_type;
430 using data_handle_type =
typename base_type::data_handle_type;
432 using reference =
typename base_type::reference;
434 template <
class,
class,
class,
class>
438 template <
class QueryDDim,
class... ODDims>
441 DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
442 if constexpr (in_tags_v<QueryDDim, detail::TypeSeq<ODDims...>>) {
443 return c.
template get<QueryDDim>();
445 return Kokkos::full_extent;
450 template <
class TypeSeq>
453 template <
class... DDims>
454 struct slicer<detail::TypeSeq<DDims...>>
456 template <
class... ODDims>
457 KOKKOS_FUNCTION
constexpr auto operator()(
458 allocation_mdspan_type
const& span,
461 return Kokkos::submdspan(span, get_slicer_for<DDims>(c)...);
470
471
475
476
480
484 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
485 ChunkSpan(
Chunk<OElementType, SupportType, Allocator>&& other)
noexcept =
delete;
488
489
493 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
495 Chunk<OElementType, SupportType, Allocator>& other)
noexcept
496 : base_type(other.m_allocation_mdspan, other.m_domain)
501
502
506 class SFINAEElementType = ElementType,
507 class = std::enable_if_t<std::is_const_v<SFINAEElementType>>,
509 class = std::enable_if_t<std::is_same_v<
typename Allocator::memory_space, MemorySpace>>>
511 Chunk<OElementType, SupportType, Allocator>
const& other)
noexcept
512 : base_type(other.m_allocation_mdspan, other.m_domain)
517
518
519 template <
class OElementType>
521 ChunkSpan<OElementType, SupportType, layout_type, MemorySpace>
const& other)
noexcept
522 : base_type(other.m_allocation_mdspan, other.m_domain)
527
528
529
531 class Mapping = mapping_type,
532 std::enable_if_t<std::is_constructible_v<Mapping, extents_type>,
int> = 0>
534 : base_type(ptr, domain)
539
540
541
543 allocation_mdspan_type allocation_mdspan,
544 SupportType
const& domain)
545 : base_type(allocation_mdspan, domain)
547 for (std::size_t i = 0; i < SupportType::rank(); ++i) {
548 assert(allocation_mdspan.extent(i)
549 ==
static_cast<std::size_t>(detail::array(domain.extents())[i]));
554
555
556
557 template <
class KokkosView,
class = std::enable_if_t<Kokkos::is_view_v<KokkosView>>>
560 detail::build_mdspan(view, std::make_index_sequence<SupportType::rank()> {}),
568
569
570
574
575
576
580
581 template <
class... QueryDDims>
583 DiscreteElement<QueryDDims...>
const& slice_spec)
const
585 slicer<to_type_seq_t<SupportType>>
const slicer;
587 = slicer(
this->allocation_mdspan(),
this->m_domain.distance_from_front(slice_spec));
588 using layout_type =
typename decltype(subview)::layout_type;
589 using extents_type =
typename decltype(subview)::extents_type;
590 using detail::TypeSeq;
591 using OutTypeSeqDDims
592 = type_seq_remove_t<to_type_seq_t<SupportType>, TypeSeq<QueryDDims...>>;
593 using OutDDom =
typename detail::RebindDomain<SupportType, OutTypeSeqDDims>::type;
595 std::is_same_v<layout_type, Kokkos::Experimental::layout_left_padded<>>
596 || std::is_same_v<layout_type, Kokkos::Experimental::layout_right_padded<>>) {
597 Kokkos::layout_stride::mapping<extents_type>
const mapping_stride(subview.mapping());
598 Kokkos::mdspan<ElementType, extents_type, Kokkos::layout_stride>
const
599 a(subview.data_handle(), mapping_stride);
603 Kokkos::layout_stride,
604 memory_space>(a, OutDDom(
this->m_domain));
610 memory_space>(subview, OutDDom(
this->m_domain));
615
616
617
618 template <
class... DElems>
622 SupportType::rank() == (0 + ... + DElems::size()),
623 "Invalid number of dimensions");
624 static_assert((is_discrete_element_v<DElems> && ...),
"Expected DiscreteElements");
625 assert(
this->m_domain.contains(delems...));
626 return DDC_MDSPAN_ACCESS_OP(
627 this->m_allocation_mdspan,
628 detail::array(
this->m_domain.distance_from_front(delems...)));
632
633
636 return base_type::data_handle();
640
641
644 return base_type::allocation_mdspan();
648
649
652 auto s =
this->allocation_mdspan();
653 auto kokkos_layout = detail::build_kokkos_layout(
656 std::make_index_sequence<SupportType::rank()> {});
658 detail::mdspan_to_kokkos_element_t<ElementType, SupportType::rank()>,
659 decltype(kokkos_layout),
660 MemorySpace>(s.data_handle(), kokkos_layout);
665 return view_type(*
this);
674template <
class DataType,
class... Properties,
class SupportType>
676 Kokkos::View<DataType, Properties...>
const& view,
679 detail::kokkos_to_mdspan_element_t<
680 typename Kokkos::View<DataType, Properties...>::
data_type>,
682 detail::kokkos_to_mdspan_layout_t<
683 typename Kokkos::View<DataType, Properties...>::
array_layout>,
684 typename Kokkos::View<DataType, Properties...>::
memory_space>;
686template <
class ElementType,
class SupportType,
class Allocator>
690 Kokkos::layout_right,
693template <
class ElementType,
class SupportType,
class Allocator>
697 Kokkos::layout_right,
703 class LayoutStridedPolicy = Kokkos::layout_right,
704 class MemorySpace = Kokkos::HostSpace>
705using 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 >