Skip to content

Commit

Permalink
Add a for loop that is unrolled at compile time (AMReX-Codes#3674)
Browse files Browse the repository at this point in the history
## 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
  • Loading branch information
maxpkatz authored Dec 19, 2023
1 parent d1e55fb commit ef38229
Showing 1 changed file with 24 additions and 0 deletions.
24 changes: 24 additions & 0 deletions Src/Base/AMReX_Loop.H
Original file line number Diff line number Diff line change
Expand Up @@ -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<auto I, auto N, class F>
AMREX_GPU_HOST_DEVICE AMREX_INLINE
constexpr void constexpr_for (F&& f)
{
if constexpr (I < N) {
f(std::integral_constant<decltype(I), I>());
constexpr_for<I+1, N>(f);
}
}

#include <AMReX_Loop.nolint.H>

}
Expand Down

0 comments on commit ef38229

Please sign in to comment.