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, adn code coupling experience.

First step - Do Now

• 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 Dropping Support for AoS #459

Planned

• To simplify code-coupling of different AMReX codes and to reduce compile-time, we plan to 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.
• Consider transitioning soon, after Specialize a Polymorphic ParticleContainer #428 was implemented.
• First tests show that this change is indeed performance-neutral, but maintenance-positive.

[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: AMReX-Codes/amrex#4603