Skip to content

Commit

Permalink
Almost implement Chunk and ChunkSpan with SupportType
Browse files Browse the repository at this point in the history
  • Loading branch information
@tpadioleau
tpadioleau committed Dec 5, 2024
1 parent 145765c commit d42ad0e
Show file tree
Hide file tree
Showing 5 changed files with 220 additions and 222 deletions.
39 changes: 19 additions & 20 deletions include/ddc/chunk.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,25 +24,24 @@ class Chunk;
template <class ElementType, class SupportType, class Allocator>
inline constexpr bool enable_chunk<Chunk<ElementType, SupportType, Allocator>> = true;

template <class ElementType, class... DDims, class Allocator>
class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
: public ChunkCommon<ElementType, DiscreteDomain<DDims...>, Kokkos::layout_right>
template <class ElementType, class SupportType, class Allocator>
class Chunk : public ChunkCommon<ElementType, SupportType, Kokkos::layout_right>
{
protected:
using base_type = ChunkCommon<ElementType, DiscreteDomain<DDims...>, Kokkos::layout_right>;
using base_type = ChunkCommon<ElementType, SupportType, Kokkos::layout_right>;

public:
/// type of a span of this full chunk
using span_type = ChunkSpan<
ElementType,
DiscreteDomain<DDims...>,
SupportType,
Kokkos::layout_right,
typename Allocator::memory_space>;

/// type of a view of this full chunk
using view_type = ChunkSpan<
ElementType const,
DiscreteDomain<DDims...>,
SupportType,
Kokkos::layout_right,
typename Allocator::memory_space>;

Expand Down Expand Up @@ -93,7 +92,7 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
/// Construct a labeled Chunk on a domain with uninitialized values
explicit Chunk(
std::string const& label,
discrete_domain_type const& domain,
SupportType const& domain,
Allocator allocator = Allocator())
: base_type(allocator.allocate(label, domain.size()), domain)
, m_allocator(std::move(allocator))
Expand All @@ -102,7 +101,7 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
}

/// Construct a Chunk on a domain with uninitialized values
explicit Chunk(discrete_domain_type const& domain, Allocator allocator = Allocator())
explicit Chunk(SupportType const& domain, Allocator allocator = Allocator())
: Chunk("no-label", domain, std::move(allocator))
{
}
Expand Down Expand Up @@ -189,7 +188,7 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
element_type const& operator()(DElems const&... delems) const noexcept
{
static_assert(
sizeof...(DDims) == (0 + ... + DElems::size()),
SupportType::rank() == (0 + ... + DElems::size()),
"Invalid number of dimensions");
static_assert((is_discrete_element_v<DElems> && ...), "Expected DiscreteElements");
assert(this->m_domain.is_inside(delems...));
Expand All @@ -206,7 +205,7 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
element_type& operator()(DElems const&... delems) noexcept
{
static_assert(
sizeof...(DDims) == (0 + ... + DElems::size()),
SupportType::rank() == (0 + ... + DElems::size()),
"Invalid number of dimensions");
static_assert((is_discrete_element_v<DElems> && ...), "Expected DiscreteElements");
assert(this->m_domain.is_inside(delems...));
Expand Down Expand Up @@ -264,9 +263,9 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
auto kokkos_layout = detail::build_kokkos_layout(
s.extents(),
s.mapping(),
std::make_index_sequence<sizeof...(DDims)> {});
std::make_index_sequence<SupportType::rank()> {});
return Kokkos::View<
detail::mdspan_to_kokkos_element_t<ElementType, sizeof...(DDims)>,
detail::mdspan_to_kokkos_element_t<ElementType, SupportType::rank()>,
decltype(kokkos_layout),
typename Allocator::memory_space>(s.data_handle(), kokkos_layout);
}
Expand All @@ -280,9 +279,9 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
auto kokkos_layout = detail::build_kokkos_layout(
s.extents(),
s.mapping(),
std::make_index_sequence<sizeof...(DDims)> {});
std::make_index_sequence<SupportType::rank()> {});
return Kokkos::View<
detail::mdspan_to_kokkos_element_t<const ElementType, sizeof...(DDims)>,
detail::mdspan_to_kokkos_element_t<const ElementType, SupportType::rank()>,
decltype(kokkos_layout),
typename Allocator::memory_space>(s.data_handle(), kokkos_layout);
}
Expand All @@ -303,13 +302,13 @@ class Chunk<ElementType, DiscreteDomain<DDims...>, Allocator>
}
};

template <class... DDims, class Allocator>
template <class SupportType, class Allocator>
Chunk(std::string const&,
DiscreteDomain<DDims...> const&,
Allocator) -> Chunk<typename Allocator::value_type, DiscreteDomain<DDims...>, Allocator>;
SupportType const&,
Allocator) -> Chunk<typename Allocator::value_type, SupportType, Allocator>;

template <class... DDims, class Allocator>
Chunk(DiscreteDomain<DDims...> const&,
Allocator) -> Chunk<typename Allocator::value_type, DiscreteDomain<DDims...>, Allocator>;
template <class SupportType, class Allocator>
Chunk(SupportType const&,
Allocator) -> Chunk<typename Allocator::value_type, SupportType, Allocator>;

} // namespace ddc
30 changes: 15 additions & 15 deletions include/ddc/chunk_common.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,26 +32,25 @@ KOKKOS_FUNCTION auto get_domain(ChunkType const& chunk) noexcept
template <class ElementType, class SupportType, class LayoutStridedPolicy>
class ChunkCommon;

template <class ElementType, class... DDims, class LayoutStridedPolicy>
class ChunkCommon<ElementType, DiscreteDomain<DDims...>, LayoutStridedPolicy>
template <class ElementType, class SupportType, class LayoutStridedPolicy>
class ChunkCommon
{
public:
using discrete_domain_type = DiscreteDomain<DDims...>;
using discrete_domain_type = SupportType;

#if defined(DDC_BUILD_DEPRECATED_CODE)
using mdomain_type [[deprecated("Use `discrete_domain_type` instead")]]
= DiscreteDomain<DDims...>;
using mdomain_type [[deprecated("Use `discrete_domain_type` instead")]] = SupportType;
#endif

/// The dereferenceable part of the co-domain but with a different domain, starting at 0
using allocation_mdspan_type = Kokkos::mdspan<
ElementType,
Kokkos::dextents<std::size_t, sizeof...(DDims)>,
Kokkos::dextents<std::size_t, SupportType::rank()>,
LayoutStridedPolicy>;

using const_allocation_mdspan_type = Kokkos::mdspan<
const ElementType,
Kokkos::dextents<std::size_t, sizeof...(DDims)>,
Kokkos::dextents<std::size_t, SupportType::rank()>,
LayoutStridedPolicy>;

using discrete_element_type = typename discrete_domain_type::discrete_element_type;
Expand Down Expand Up @@ -91,7 +90,7 @@ class ChunkCommon<ElementType, DiscreteDomain<DDims...>, LayoutStridedPolicy>
allocation_mdspan_type m_allocation_mdspan;

/// The mesh on which this chunk is defined
discrete_domain_type m_domain;
SupportType m_domain;

public:
static KOKKOS_FUNCTION constexpr int rank() noexcept
Expand Down Expand Up @@ -128,9 +127,9 @@ class ChunkCommon<ElementType, DiscreteDomain<DDims...>, LayoutStridedPolicy>
template <class Mapping = mapping_type>
static KOKKOS_FUNCTION constexpr std::
enable_if_t<std::is_constructible_v<Mapping, extents_type>, allocation_mdspan_type>
make_allocation_mdspan(ElementType* ptr, discrete_domain_type const& domain)
make_allocation_mdspan(ElementType* ptr, SupportType const& domain)
{
return allocation_mdspan_type(ptr, ::ddc::extents<DDims>(domain).value()...);
return allocation_mdspan_type(ptr, detail::array(domain.extents()));
}

public:
Expand All @@ -139,7 +138,7 @@ class ChunkCommon<ElementType, DiscreteDomain<DDims...>, LayoutStridedPolicy>
return m_allocation_mdspan.accessor();
}

KOKKOS_FUNCTION constexpr DiscreteVector<DDims...> extents() const noexcept
KOKKOS_FUNCTION constexpr SupportType::discrete_vector_type extents() const noexcept
{
return m_domain.extents();
}
Expand Down Expand Up @@ -178,13 +177,14 @@ class ChunkCommon<ElementType, DiscreteDomain<DDims...>, LayoutStridedPolicy>
template <class QueryDDim>
KOKKOS_FUNCTION constexpr size_type stride() const
{
return m_allocation_mdspan.stride(type_seq_rank_v<QueryDDim, detail::TypeSeq<DDims...>>);
return m_allocation_mdspan.stride(
type_seq_rank_v<QueryDDim, detail::ToTypeSeq<SupportType>>);
}

/** Provide access to the domain on which this chunk is defined
* @return the domain on which this chunk is defined
*/
KOKKOS_FUNCTION constexpr discrete_domain_type domain() const noexcept
KOKKOS_FUNCTION constexpr SupportType domain() const noexcept
{
return m_domain;
}
Expand All @@ -208,7 +208,7 @@ class ChunkCommon<ElementType, DiscreteDomain<DDims...>, LayoutStridedPolicy>
*/
KOKKOS_FUNCTION constexpr ChunkCommon(
allocation_mdspan_type allocation_mdspan,
discrete_domain_type const& domain) noexcept
SupportType const& domain) noexcept
: m_allocation_mdspan(std::move(allocation_mdspan))
, m_domain(domain)
{
Expand All @@ -221,7 +221,7 @@ class ChunkCommon<ElementType, DiscreteDomain<DDims...>, LayoutStridedPolicy>
template <
class Mapping = mapping_type,
std::enable_if_t<std::is_constructible_v<Mapping, extents_type>, int> = 0>
KOKKOS_FUNCTION constexpr ChunkCommon(ElementType* ptr, discrete_domain_type const& domain)
KOKKOS_FUNCTION constexpr ChunkCommon(ElementType* ptr, SupportType const& domain)
: m_allocation_mdspan(make_allocation_mdspan(ptr, domain))
, m_domain(domain)
{
Expand Down
Loading

0 comments on commit d42ad0e

Please sign in to comment.