Particle Roadmap in AMReX/pyAMReX

[ax3l]

To make optimal use of pyAMReX, code bindings, performance per se, etc. we are undergoing a general modernization of particle data structures in AMReX.

I encourage all AMReX codes that use particles to follow the modernizations, which will not only provide more performance but also a more seamless Python, AI/ML, and code coupling experience.

First step (2023)

• Transition to PureSoA particles for easier layout, faster performance and potential better SIMD performance.
• The implementation is now available since 2023 and well-tested in BLAST codes [1], [2], [3].
• Status: transition now. We will drop support for AoS to focus development in pyAMReX, see #459

Second step (2025)

• To simplify code-coupling of different AMReX codes and to reduce compile-time, we transition to amrex::PolymorphicArenaAllocator instead of having separate specializations for host/pinned/device containers. This fixes a long-standing design issue in ParticleContainer that was solved better in AMReX MultiFab.
• Transition after #428 was implemented.
• First tests show that this change is indeed performance-neutral and reduces pyAMReX compile- and link time dramatically.

Third step - Soon (2025/26+)

• For maximum runtime compatibility of AMReX codes and better modularity, we will add a new runtime-component-only particle type (of pure SoA).
• The way AMReX particle compute pattern are performed (pointer unwrapping on the host before the kernel), there will be no performance impact (after https://github.com/AMReX-Codes/amrex/pull/4404 )
• This will reduce type complexity, and shrink compile- and link-time of specialized PC types to a minimum.
• With this, all AMReX codes can use the exact same PC type, increasing compatibility (handling at runtime) & aiding modular design.

[atmyers]

With PolymorphicArenaAllocator, we need to make sure users can only call setArena before any particles are added to the container. Otherwise I'm not sure what the desired behavior would be.

[BenWibking]

Are there any plans to update the AMReX documentation to document PureSoA particles?

We've been wanting to transition, but all of the existing documentation AFAIK (https://amrex-codes.github.io/amrex/docs_html/Particle_Chapter.html) is for AoS particles. Some code patterns are straightforward to translate, but once we run into one usage pattern that we can't find an equivalent for, we are stuck.

[ax3l]

With PolymorphicArenaAllocator, we need to make sure users can only call setArena` before any particles are added to the container. Otherwise I'm not sure what the desired behavior would be.

Agreed, we should throw a runtime error otherwise.

We could provide simple to(arena) copy helper functions, as in PyTorch: move memory but stay in the current variable. Easier/straight forward would also be something like copy_to(arena) that returns a copy in a new variable.

Are there any plans to update the AMReX documentation to document PureSoA particles?

Yes this is urgently needed. @atmyers can you please take this on? I can assist/review.

[AlexanderSinn]

This makes it easier to use PolymorphicArenaAllocator: https://github.com/AMReX-Codes/amrex/pull/4603

[ax3l]

I am starting off a couple PRs to get polymorphic PC into pyAMReX and WarpX/ImpactX. HiPACE++ (thanks to @AlexanderSinn) is already using it.

• [x] HiPACE++: https://github.com/Hi-PACE/hipace/pulls?q=is%3Apr+PolymorphicArenaAllocator+is%3Aclosed
• [ ] pyAMReX: https://github.com/AMReX-Codes/pyamrex/pull/428
• [ ] WarpX: https://github.com/BLAST-WarpX/warpx/pull/6374
• [ ] ImpactX: https://github.com/BLAST-ImpactX/impactx/pull/1210

[ax3l]

@atmyers updated the AMReX docs for pure SoA:

• https://github.com/AMReX-Codes/amrex/pull/4759

Andrew, the first 4 sub-sections of https://amrex-codes.github.io/amrex/docs_html/Particle_Chapter.html still look like they recommend the old mixed layout by default. Can you please rework them so users do not pick up the old format anymore for new projects?