Particle load-balancing on startup

[n01r]

Certain setups, like laser solid density interactions, require very high numbers of particles and have very strongly localized density distributions.
It would be good have an option to load-balance the particles already before or during creation to avoid out-of-memory errors during initialization.

From meeting logs:

• Perhaps precompute where particles need to be initialized?
• With profiles / parsers, we currently do:
  • For a given GPU, go through cells, count how many particles go in each cell.
  • Then init in cells.
• Challenge: we do the domain decomposition before we init the particles, so for very localized, dense-plasma targets we quickly run out of memory during init and have to use way more GPUs than we need for later sim.
• Workflow to solve would be:
  • Store as before the number of particles that will go in each cell.
  • Build a cost function from that.
  • Redistribute blocks.
  • Init particles.

[lucafedeli88]

Hi @n01r, I think that this in an excellent idea and I also think that implementing what you suggest should not be too difficult. It would be beneficial for several of our simulation cases.

I was discussing offline this issue with @tmsclark2 and one of our considerations is that a possible further improvement would be to implement a mechanism to inform during the simulation the rank 0 process that at least one rank is running out of memory, which would result into an immediate call to the load balancer. This could be achieved with a reduce MPI call. However, since this is more complex to implement, I think that what you propose is already a very good starting point.

[ax3l]

Discussed on Mar 27th, 2024:
AMReX on the very first init: w/o calling load balance first, create a space filling curve distribution map.
MakeDistributionMap: Overwrite to not do a default with SFC.
Try to use a spaced out assignment as an opt-in option.

@WeiqunZhang will work on this as part of KISMET 🙏

[n01r]

That sounds great, thank you! 🙏