Update Halo class to allow halo exchanges for device arrays

[brian-oneill]

This PR updates the Halo to allow for halo exchanges of arrays allocated in device memory space as well as host memory space. With this update, Omega can take advantage of GPU aware MPI implementations.

Changes include:

• Replacing SendBuffer, RecvBuffer and certain other member vectors of Neighbor and ExchList classes with arrays on both host and device
• Replace packBuffer and unpackBuffer overloaded functions with specialized function templates for supported array ranks
• Removing deepCopy before and after halo exchanges in classes and unit tests where no longer necessary
• Add device array tests to Halo unit test
• Add OMEGA_MPI_ON_DEVICE build flag (on by default)

Successfully built and passed unit tests with OMEGA_MPI_ON_DEVICE both on and off on Chrysalis (intel), Perlmutter CPU (intel) & GPU (nvidiagpu), and Frontier CPU (crayclang) & GPU (crayclanggpu)

Checklist

• Documentation:
  • User's Guide has been updated
  • Developer's Guide has been updated
  • Documentation has been built locally and changes look as expected
• Testing
  • Unit tests have passed. Please provide a relevant CDash build entry for verification.

[rljacob]

Curious why you test Perlmutter GPU with nvidia instead of gnu (which is what E3SM and SCREAM test with).

[grnydawn]

Curious why you test Perlmutter GPU with nvidia instead of gnu (which is what E3SM and SCREAM test with).

@rljacob , I think we want Omega to support both NVIDIA and GNU compilers on Perlmutter GPU nodes. Since Perlmutter uses NVIDIA GPUs, I think that we typically test Omega with the NVIDIA compiler first and the GNU compiler next. However, if E3SM and SCREAM consider GNU as the primary compiler on Perlmutter, we may also adopt the same compiler preference.

[rljacob]

We have not actually done a performance comparison between nvidia, gnu and intel on perlmutter gpus. But we have seen nvidia have trouble with some of the Fortran code. gnu is preferred unless there is evidence another one is better.

[mark-petersen]

Confirmed that PR passes unit tests on Frontier with cpu and gpu. Since the unit tests show this is working correctly because they create arrays with unique values per cell, do a halo exchange, and then compute the error.

Awaiting timing tests from Kieran Ringel for performance comparison between this new halo exchange on device and the previous halo exchange on host (gpu versus cpu).

[grnydawn]

Please see my comment on each items

[mwarusz]

I found one correctness bug and I think some generic Kokkos utilities could be moved to OmegaKokkos.h but the changes in general look good to me. I performed some initial scaling experiments on Frontier and Perlmutter. In general, the performance improved, although the results on Frontier for large gpu counts are a bit surprising (the original code is slightly faster). I also wasn't able to get GPU-aware MPI to work on Perlmutter.

[mwarusz]

Timing results for this PR with OMEGA_MPI_ON_DEVICE turned on and off (indicted in second half of name in the legend)

@kieran-ringel Note that depending on what exactly you measured this bug #163 (comment) might have affected these results, since it causes state and tracer halo exchanges to exchange host arrays only.

[kieran-ringel]

Updating timing with updated exchange device arrays in State and Tracer exchangeHalo functions

(note, x axis is actually number of nodes, not GPUs)

[grnydawn]

Approved. Tests passed on Frontier, Perlmutter, and Chrysalis using various compilers. The only exception is a segmentation fault failure in TIMESTEPPING_TEST on Perlmutter-GPU with the gnugpu compiler. This issue appears to be isolated to this specific case and could be addressed in a separate task.

[brian-oneill]

@grnydawn I had ran the tests successfully with gnugpu after the discussion above and didn't see any issue, but I built in debug mode. I'm seeing this seg fault now building in release. I'll take a look and see if I can track it down.

[kieran-ringel]

I ran the same performance tests as above on Perlmutter Approving PR based on performance results on these 2 machines

[grnydawn]

@brian-oneill , It is strange that I encountered a failure while running Omega built in "Debug" mode. However, when I retried the TIMESTEPPER_TEST this morning, I did not experience any failures in either "Debug" or "Release" mode.

FYI, yesterday, I was able to locate the source lines where the failure occurs. It segfaulted in Halo.h near line 850.

         parallelFor(
             {NK, NTotList, NJ}, KOKKOS_LAMBDA(int K, int IExch, int J) {
                auto Val       = Array(K, LocIndex(IExch), J);
                const R8 RVal  = reinterpret_cast<R8 &>(Val);
                const I4 IBuff = (K * NTotList + IExch) * NJ + J;
                LocBuff(IBuff) = RVal;
             });

[mark-petersen]

Thank you @brian-oneill! Approving based on testing above, particularly the plot of Halo exchange speed-up time by @kieran-ringel. The timing results on perlmutter are confusing, but I think that is due to the architecture and not an issue with this PR.

[mark-petersen]

Merging. We will keep an eye out for the failure @grnydawn mentions above, but since it was not repeatable, I don't want to hold up this merge.