Suggestions for efficient scheme for elementwise reduction across ranks on MultiFabs?

[ElloiseFangelLloyd]

Hi all

I'm asking for suggestions on how to implement a particular operation. I appreciate any and all help :)

I have e.g. 8 MultiFabs which exist completely separately on 8 different ranks. As such they do not have the same DistributionMapping - each mf exists only on its respective rank. Each mf has one component and three dimensions.

I'd like the master rank to have a MultiFab which is the elementwise sum of the 8 other mfs.

From what I can see, ParallelDescriptor::ReduceRealSum() can perform these elementwise reductions, but first requires that the data be flattened, potentially into a Vector or similar. This is an easy enough thing to implement - but it seems inefficient to have to flatten each MultiFab each time and to introduce not-insignificant memory overhead.

If this is the best approach, then that's fine, but I am wondering if anyone has any suggestions for other ideas?
I had also wondered about doing a more manual MPI_Gather or simply iterating with MFIter and doing an elementwise addition, but this would scale badly over large numbers of GPUs.

Any feedback is appreciated! :)

[WeiqunZhang]

Do these MultiFabs have the same BoxArray? How many boxes do they have? Do they have ghost cells? If so, are the number of ghost cells the same? I suppose the BoxArrays are the same and the number of ghost cells are the same too, otherwise it probably does not make sense. If that is the case and there are not that many boxes, you can do

for (MFiter mfi(mf); mfi.isValid(); ++mfi) {
    auto& fab = mf[mfi];
    ParallellReduce::Sum(fab.dataPtr(), int(fab.size()), 0, ParallelDescriptor::Communicator()); 
}

If you have lots of boxes, you can allocate a single chunk of memory for a MultiFab. For example, MultiFab mf(ba,dm,ncomp,nghost,MFInfo{}.SetAllocSingleChunk(true). Then you can do

ParallellReduce::Sum(mf.SingleChunkPtr1, int(mf.singleChunkSize()/sizeof(Real)), 0, ParallelDescriptor::Communicator()); 

Here we assume that GPU aware MPI is used. If not, you will have to memcpy data from multifab to pinned memory first, before you can MPI Reduce.

However, the question is why you need to do this in the first place.

[ElloiseFangelLloyd]

Yes, you are completely right that they have the same BoxArray, ncomp, ngrow, there are very few boxes in the domain, and we are using GPU-aware MPI.

Your first code snippet was much more performant than my MPI implementation! Thanks for that. I'm just now trying to improve the performance a bit here and there.

Re: why we would do this: in hindsight I should have been more clear. :) I understand (but please correct me otherwise) that the AMReX distribution strategy is spatial. Most of my simulation time is spent processing particles (non-AMReX datastructures) that move through the spatial domain very quickly, so we would have to very often transfer particles between devices when they move between local spatial domains. With our setup this seemed like a solution doomed to scale poorly, and the AMReX Poisson solver is highly efficient even just on one GPU, we are instead choosing not to distribute the Poisson solve, so we have no spatial dependence in the particle distribution. On each rank, we deposit those particles onto the MultiFab grids, transfer them to the master, and let the master solve Poisson.

Our physical setup has some similarities with the ones WarpX can do, but we have additional physics that meant it made most sense for us to build our own datastructures etc. and then use AMReX for its solvers.