AMReX-Codes / amrex

AMReX: Software Framework for Block Structured AMR
https://amrex-codes.github.io/amrex
Other
503 stars 336 forks source link

Add TupleSplit #3978

Closed AlexanderSinn closed 3 weeks ago

AlexanderSinn commented 3 weeks ago

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:

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.

Checklist

The proposed changes:

WeiqunZhang commented 3 weeks ago

A slightly shorter version could be

diff --git a/Src/Base/AMReX_Tuple.H b/Src/Base/AMReX_Tuple.H
index 080b5e349f..5dbf0c4022 100644
--- a/Src/Base/AMReX_Tuple.H
+++ b/Src/Base/AMReX_Tuple.H
@@ -324,19 +324,14 @@ namespace detail {
         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];
+            return std::get<J>(std::forward_as_tuple(Is...));
         }

-        template<std::size_t J>
+        template<std::size_t... Js>
         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;
+        static constexpr std::size_t get_exclusive_sum (std::index_sequence<Js...>) noexcept {
+            std::size_t arr[] = {Is...};
+            return (std::size_t(0) + ... + arr[Js]);
         }
     };

@@ -354,7 +349,7 @@ namespace detail {
     TupleSplitImp (const GpuTuple<Args...>& tup, std::index_sequence<Is...>, SIL) noexcept
     {
         return makeTuple(
-            GetSubTuple<(SIL::template get_exclusive_sum<Is>())>(
+            GetSubTuple<(SIL::get_exclusive_sum(std::make_index_sequence<Is>()))>(
                 tup,
                 std::make_index_sequence<SIL::template get<Is>()>()
             )...

But it does not matter. I was just having fun.

AlexanderSinn commented 3 weeks ago

The return (std::size_t(0) + ... + arr[Js]); is very nice.