Optimizing GPU acceleration of binary collision algorithms.

[mhaseeb123]

[Edit by @RemiLehe]: The aim of this PR is to speed up binary collisions on GPU by exposing more parallelism: instead of looping with one GPU thread per cell, we loop with one GPU thread per "number of independent pair" (i.e. pairs that do not touch the same macroparticles, so that there is no race condition), where the number of independent pairs is determined by the lower number of macroparticle of either species, within each cell.

[Updated]: This PR includes optimized GPU implementation for the two particle collision algorithms (Coulomb and Nuclear).

Average Performance Improvement: 4x

Results for a specific collisions-heavy use-case from Dave inputs_1d_H1_lassen2.txt

Built with: cmake -DWarpX_DIMS=1 -DWarpX_COMPUTE=CUDA on Perlmutter

New kernel:

MultiParticleContainer::doCollisions()                   100      11.89      11.89      11.89  81.47%
Avg. per step = 0.1375590809 s
Total Time                     : 14.58777054

WarpX dev branch:

MultiParticleContainer::doCollisions()                   100      67.96      67.96      67.96  96.49%
Avg. per step = 0.6979252147
Total Time                     : 70.43095279

Error Rate

Slightly higher for a few Azure tests leading to failure of CI pipelines. All good otherwise.

TODO (@RemiLehe):

• Add some more comments
• Benchmark on CPU + possibly use #ifdef to use the new code only with GPU

[RemiLehe]

Thanks for this PR!
I tried to run the Deuterium_Tritium_Fusion_3D locally (on a Macbook), with:

./run_test.sh Deuterium_Tritium_Fusion_3D

but I got the following error:

--- INFO    : Writing plotfile Deuterium_Tritium_Fusion_3D_plt000000
STEP 1 starts ...

   WARNING: Test stderr:
SIGILL Invalid, privileged, or ill-formed instruction
See Backtrace.0.0 file for details
SIGILL Invalid, privileged, or ill-formed instruction
See Backtrace.1.0 file for details
Abort(4) on node 0 (rank 0 in comm 496): application called MPI_Abort(comm=0x84000002, 4) - process 0
Abort(4) on node 1 (rank 1 in comm 496): application called MPI_Abort(comm=0x84000001, 4) - process 1

Note that this error only appears when using amrex.fpe_trap_invalid=1 (which is automatically used with the run_test.sh utility)

[roelof-groenewald]

@RemiLehe I ran the CPU performance tests as we discussed (using cases 1 - 3 of the Turner benchmarks). The short and good news is that the results matched expectation and performance was NOT worse on CPU due to these changes. So, as discussed, this PR can be merged 🎉

Here are the inclusive timings (note that MultiParticleContainer::doCollisions() include both MCC and DSMC collisions involved in these tests, but we don't have a timer for just DSMC):

• test 1:

  • devel: MultiParticleContainer::doCollisions() 512000 3.723 52.39 145.4 13.71%
  • this PR: MultiParticleContainer::doCollisions() 512000 3.155 35.78 124.8 12.19%

• test 2:

  • devel: MultiParticleContainer::doCollisions() 4096000 390.6 636.2 796 13.30%
  • this PR: MultiParticleContainer::doCollisions() 4096000 275.8 523 746.5 12.48%

• test 3:

  • devel: MultiParticleContainer::doCollisions() 8192000 721.1 2407 3471 19.19%
  • this PR: MultiParticleContainer::doCollisions() 8192000 497 2099 3287 17.91%

[roelof-groenewald]

Here are also the physics results (for posterity) with the changes in this PR included:

[RemiLehe]

@roelof-groenewald Thanks a lot for doing these tests! This is great news!