Skip to content

Commit 03d9e37

Browse files
tpadioleautpadioleau
committed
Slice from a DiscreteVector
1 parent a1c87dd commit 03d9e37

3 files changed

Lines changed: 90 additions & 30 deletions

File tree

‎src/ddc/chunk.hpp‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,20 @@ class Chunk : public ChunkCommon<ElementType, SupportType, Kokkos::layout_right>
149149
return *this;
150150
}
151151

152+
/// Slice out some dimensions
153+
template <class... QueryDDims>
154+
auto operator[](DiscreteVector<QueryDDims...> const& slice_spec) const
155+
{
156+
return view_type(*this)[slice_spec];
157+
}
158+
159+
/// Slice out some dimensions
160+
template <class... QueryDDims>
161+
auto operator[](DiscreteVector<QueryDDims...> const& slice_spec)
162+
{
163+
return span_view()[slice_spec];
164+
}
165+
152166
/// Slice out some dimensions
153167
template <class... QueryDDims>
154168
auto operator[](DiscreteElement<QueryDDims...> const& slice_spec) const

‎src/ddc/chunk_span.hpp‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -265,16 +265,15 @@ class ChunkSpan : public ChunkCommon<ElementType, SupportType, LayoutStridedPoli
265265
/** Slice out some dimensions
266266
*/
267267
template <class... QueryDDims>
268-
KOKKOS_FUNCTION constexpr auto operator[](
269-
DiscreteElement<QueryDDims...> const& slice_spec) const
268+
KOKKOS_FUNCTION constexpr auto operator[](DiscreteVector<QueryDDims...> const& slice_spec) const
270269
{
271270
using detail::TypeSeq;
272-
using QueryDDom = detail::Rebind<SupportType, TypeSeq<QueryDDims...>>::type;
273-
KOKKOS_ASSERT(QueryDDom(this->m_domain).contains(slice_spec))
271+
KOKKOS_ASSERT(
272+
((DiscreteVector<QueryDDims>(slice_spec)
273+
< DiscreteVector<QueryDDims>(this->m_domain.extents()))
274+
&& ...))
274275
Slicer<to_type_seq_t<SupportType>> const slicer;
275-
auto subview = slicer(
276-
this->allocation_mdspan(),
277-
QueryDDom(this->m_domain).distance_from_front(slice_spec));
276+
auto subview = slicer(this->allocation_mdspan(), slice_spec);
278277
using layout_type = decltype(subview)::layout_type;
279278
using extents_type = decltype(subview)::extents_type;
280279
using OutTypeSeqDDims
@@ -300,6 +299,17 @@ class ChunkSpan : public ChunkCommon<ElementType, SupportType, LayoutStridedPoli
300299
}
301300
}
302301

302+
/** Slice out some dimensions
303+
*/
304+
template <class... QueryDDims>
305+
KOKKOS_FUNCTION constexpr auto operator[](
306+
DiscreteElement<QueryDDims...> const& slice_spec) const
307+
{
308+
using QueryDDom = detail::Rebind<SupportType, detail::TypeSeq<QueryDDims...>>::type;
309+
KOKKOS_ASSERT(QueryDDom(this->m_domain).contains(slice_spec))
310+
return (*this)[QueryDDom(this->m_domain).distance_from_front(slice_spec)];
311+
}
312+
303313
/** Restrict to a subdomain, only valid when SupportType is a DiscreteDomain
304314
*/
305315
template <class... QueryDDims>

‎tests/chunk.cpp‎

Lines changed: 59 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -496,8 +496,6 @@ TEST(Chunk2DTest, Cview)
496496

497497
TEST(Chunk2DTest, SliceCoordX)
498498
{
499-
DElemX const slice_x_val(lbound_x + 1);
500-
501499
ChunkXY<double> chunk(dom_x_y);
502500
ChunkXY<double> const& chunk_cref = chunk;
503501
for (DElemX const ix : chunk.domain<DDimX>()) {
@@ -506,20 +504,35 @@ TEST(Chunk2DTest, SliceCoordX)
506504
}
507505
}
508506

509-
ddc::ChunkSpan const chunk_y = chunk_cref[slice_x_val];
510-
EXPECT_TRUE(
511-
(std::is_same_v<std::decay_t<decltype(chunk_y)>::layout_type, Kokkos::layout_right>));
512-
EXPECT_EQ(chunk_y.extent<DDimY>(), chunk.extent<DDimY>());
513-
for (DElemY const iy : chunk_cref.domain<DDimY>()) {
514-
// we expect complete equality, not EXPECT_DOUBLE_EQ: these are copy
515-
EXPECT_EQ(chunk_y(iy), chunk_cref(slice_x_val, iy));
507+
{
508+
DElemX const slice_x_val(lbound_x + 1);
509+
ddc::ChunkSpan const chunk_y = chunk_cref[slice_x_val];
510+
EXPECT_TRUE((std::is_same_v<
511+
std::decay_t<decltype(chunk_y)>::layout_type,
512+
Kokkos::layout_right>));
513+
EXPECT_EQ(chunk_y.extent<DDimY>(), chunk.extent<DDimY>());
514+
for (DElemY const iy : chunk_cref.domain<DDimY>()) {
515+
// we expect complete equality, not EXPECT_DOUBLE_EQ: these are copy
516+
EXPECT_EQ(chunk_y(iy), chunk_cref(slice_x_val, iy));
517+
}
518+
}
519+
520+
{
521+
DVectX const slice_x_val(1);
522+
ddc::ChunkSpan const chunk_y = chunk_cref[slice_x_val];
523+
EXPECT_TRUE((std::is_same_v<
524+
std::decay_t<decltype(chunk_y)>::layout_type,
525+
Kokkos::layout_right>));
526+
EXPECT_EQ(chunk_y.extent<DDimY>(), chunk.extent<DDimY>());
527+
for (DElemY const iy : chunk_cref.domain<DDimY>()) {
528+
// we expect complete equality, not EXPECT_DOUBLE_EQ: these are copy
529+
EXPECT_EQ(chunk_y(iy), chunk_cref(lbound_x + slice_x_val, iy));
530+
}
516531
}
517532
}
518533

519534
TEST(Chunk2DTest, SliceCoordY)
520535
{
521-
DElemY const slice_y_val(lbound_y + 1);
522-
523536
ChunkXY<double> chunk(dom_x_y);
524537
ChunkXY<double> const& chunk_cref = chunk;
525538
for (DElemX const ix : chunk.domain<DDimX>()) {
@@ -528,13 +541,30 @@ TEST(Chunk2DTest, SliceCoordY)
528541
}
529542
}
530543

531-
ddc::ChunkSpan const chunk_x = chunk_cref[slice_y_val];
532-
EXPECT_TRUE(
533-
(std::is_same_v<std::decay_t<decltype(chunk_x)>::layout_type, Kokkos::layout_stride>));
534-
EXPECT_EQ(chunk_x.extent<DDimX>(), chunk.extent<DDimX>());
535-
for (DElemX const ix : chunk_cref.domain<DDimX>()) {
536-
// we expect complete equality, not EXPECT_DOUBLE_EQ: these are copy
537-
EXPECT_EQ(chunk_x(ix), chunk_cref(ix, slice_y_val));
544+
{
545+
DElemY const slice_y_val(lbound_y + 1);
546+
ddc::ChunkSpan const chunk_x = chunk_cref[slice_y_val];
547+
EXPECT_TRUE((std::is_same_v<
548+
std::decay_t<decltype(chunk_x)>::layout_type,
549+
Kokkos::layout_stride>));
550+
EXPECT_EQ(chunk_x.extent<DDimX>(), chunk.extent<DDimX>());
551+
for (DElemX const ix : chunk_cref.domain<DDimX>()) {
552+
// we expect complete equality, not EXPECT_DOUBLE_EQ: these are copy
553+
EXPECT_EQ(chunk_x(ix), chunk_cref(ix, slice_y_val));
554+
}
555+
}
556+
557+
{
558+
DVectY const slice_y_val(1);
559+
ddc::ChunkSpan const chunk_x = chunk_cref[slice_y_val];
560+
EXPECT_TRUE((std::is_same_v<
561+
std::decay_t<decltype(chunk_x)>::layout_type,
562+
Kokkos::layout_stride>));
563+
EXPECT_EQ(chunk_x.extent<DDimX>(), chunk.extent<DDimX>());
564+
for (DElemX const ix : chunk_cref.domain<DDimX>()) {
565+
// we expect complete equality, not EXPECT_DOUBLE_EQ: these are copy
566+
EXPECT_EQ(chunk_x(ix), chunk_cref(ix, lbound_y + slice_y_val));
567+
}
538568
}
539569
}
540570

@@ -543,11 +573,17 @@ TEST(Chunk2DTest, SliceCoordXOutOfBounds)
543573
#if !defined(NDEBUG) // The assertion is only checked if NDEBUG isn't defined
544574
ChunkXY<double> chunk(dom_x_y);
545575

546-
char const* const death_msg
547-
= R"rgx(.ssert.*QueryDDom\(this->m_domain\).contains\(slice_spec\))rgx";
548-
549-
EXPECT_DEATH(chunk[dom_x.front() - 1], death_msg);
550-
EXPECT_DEATH(chunk[dom_x.back() + 1], death_msg);
576+
{
577+
char const* const death_msg
578+
= R"rgx(.ssert.*QueryDDom\(this->m_domain\).contains\(slice_spec\))rgx";
579+
EXPECT_DEATH(chunk[dom_x.front() - 1], death_msg);
580+
EXPECT_DEATH(chunk[dom_x.back() + 1], death_msg);
581+
}
582+
{
583+
char const* const death_msg
584+
= R"rgx(.ssert.*\(DiscreteVector<QueryDDims>\(slice_spec\) < DiscreteVector<QueryDDims>\(this->m_domain.extents\(\)\)\) && ...\))rgx";
585+
EXPECT_DEATH(chunk[dom_x.extents()], death_msg);
586+
}
551587
#else
552588
GTEST_SKIP();
553589
#endif

0 commit comments

Comments
 (0)