Skip to content

Commit

Permalink
Add TupleSplit (#3978)
Browse files Browse the repository at this point in the history
## Summary

This PR adds a `TupleSplit` function for `GpuTuple` similar
to `IntVectSplit` in #3969. It can be used as an inverse function
of `TupleCat` and `TypeMultiplier`.
Example:
```C++
auto tup = amrex::makeTuple(2,4,5,7,2.324,7,8,342.3f,4ull,1ll,-38,"test");
auto [t0,t1,t2,t3] = amrex::TupleSplit<3,3,4,2>(tup);
// t0 = GpuTuple( 2, 4, 5 )
// t1 = GpuTuple( 7, 2.324, 7 )
// t2 = GpuTuple( 8, 342.2999878, 4, 1 )
// t3 = GpuTuple( -38, test )
```

## Additional background

The current implementation does not move the tuple elements and makes a
copy instead.
  • Loading branch information
AlexanderSinn committed Jun 10, 2024
1 parent c40b1ba commit f42aa8c
Showing 1 changed file with 64 additions and 0 deletions.
64 changes: 64 additions & 0 deletions Src/Base/AMReX_Tuple.H
Original file line number Diff line number Diff line change
Expand Up @@ -315,6 +315,70 @@ TupleCat (TP1&& a, TP2&& b, TPs&&... args)
std::forward<TPs>(args)...);
}

// TupleSplit

namespace detail {

template<std::size_t...Is>
struct SplitIndexList {
template<std::size_t J>
AMREX_GPU_HOST_DEVICE
static constexpr std::size_t get () noexcept {
std::size_t arr[sizeof...(Is)] = {Is...};
return arr[J];
}

template<std::size_t J>
AMREX_GPU_HOST_DEVICE
static constexpr std::size_t get_exclusive_sum () noexcept {
std::size_t arr[sizeof...(Is)] = {Is...};
std::size_t sum = 0;
for (std::size_t k=0; k<J; ++k) {
sum += arr[k];
}
return sum;
}
};

template <std::size_t start, typename... Args, std::size_t... Is>
AMREX_GPU_HOST_DEVICE
constexpr auto
GetSubTuple (const GpuTuple<Args...>& tup, std::index_sequence<Is...>) noexcept
{
return makeTuple(amrex::get<start+Is>(tup)...);
}

template <typename... Args, std::size_t... Is, typename SIL>
AMREX_GPU_HOST_DEVICE
constexpr auto
TupleSplitImp (const GpuTuple<Args...>& tup, std::index_sequence<Is...>, SIL) noexcept
{
return makeTuple(
GetSubTuple<(SIL::template get_exclusive_sum<Is>())>(
tup,
std::make_index_sequence<SIL::template get<Is>()>()
)...
);
}
}

/**
* \brief Returns a GpuTuple of GpuTuples obtained by splitting the input GpuTuple
* according to the sizes specified by the template arguments.
*/
template <std::size_t... Is, typename... Args>
AMREX_GPU_HOST_DEVICE
constexpr auto
TupleSplit (const GpuTuple<Args...>& tup) noexcept
{
static_assert((0 + ... + Is) == sizeof...(Args), "Incorrect total size in TupleSplit");
return detail::TupleSplitImp(
tup,
std::make_index_sequence<sizeof...(Is)>(),
detail::SplitIndexList<Is...>()
);
}

// Apply

namespace detail {
Expand Down

0 comments on commit f42aa8c

Please sign in to comment.