From ef38229189e3213f992a2e89dbe304fb49db9287 Mon Sep 17 00:00:00 2001 From: Max Katz Date: Mon, 18 Dec 2023 23:21:03 -0500 Subject: [PATCH] Add a for loop that is unrolled at compile time (#3674) ## Summary The constexpr_for function is fully unrolled at compile time. This is useful for relatively short loops where some of the functions inside the loop are known to possible to evaluate at compile time and may be relatively expensive, so evaluating them at compile time rather than at runtime may be beneficial for performance reasons. ## Additional background This has been used in AMReX-Astro/Microphysics successfully in the context of evaluating some nuclear reaction network quantities at compile time. ## Checklist The proposed changes: - [ ] fix a bug or incorrect behavior in AMReX - [x] add new capabilities to AMReX - [ ] changes answers in the test suite to more than roundoff level - [ ] are likely to significantly affect the results of downstream AMReX users - [ ] include documentation in the code and/or rst files, if appropriate --- Src/Base/AMReX_Loop.H | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/Src/Base/AMReX_Loop.H b/Src/Base/AMReX_Loop.H index 84b39107e45..19e1c3e5191 100644 --- a/Src/Base/AMReX_Loop.H +++ b/Src/Base/AMReX_Loop.H @@ -211,6 +211,30 @@ void LoopConcurrentOnCpu (Box const& bx, int ncomp, F&& f) noexcept }}}} } +// Implementation of "constexpr for" based on +// https://artificial-mind.net/blog/2020/10/31/constexpr-for +// +// Approximates what one would get from a compile-time +// unrolling of the loop +// for (int i = 0; i < N; ++i) { +// f(i); +// } +// +// The mechanism is recursive: we evaluate f(i) at the current +// i and then call the for loop at i+1. f() is a lambda function +// that provides the body of the loop and takes only an integer +// i as its argument. + +template +AMREX_GPU_HOST_DEVICE AMREX_INLINE +constexpr void constexpr_for (F&& f) +{ + if constexpr (I < N) { + f(std::integral_constant()); + constexpr_for(f); + } +} + #include }