Fix: GPU runs w/o GPU-Aware MPI

cmake/dependencies/ABLASTR.cmake

@@ -186,7 +186,7 @@ set(ImpactX_ablastr_branch "759c5cb11a417187edc764b853cfb68018b91fce"
 set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git"
     CACHE STRING
     "Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)")
-set(ImpactX_amrex_branch "c5d27c8ff1b7b32238ebc75a9642a3fc36e65e9c"
+set(ImpactX_amrex_branch "5fbc49bce6ccaaefd771c7594e24e0ef46d5e01c"
     CACHE STRING
     "Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)")
 

src/particles/spacecharge/Deposit1D.cpp

@@ -15,6 +15,7 @@
 #include <AMReX_REAL.H>
 #include <AMReX_BLProfiler.H>
 #include <AMReX_GpuContainers.H>
+#include <AMReX_GpuParallelReduce.H>
 #include <AMReX_ParmParse.H>
 
 #include <cmath>
@@ -47,8 +48,7 @@ namespace impactx::particles::spacecharge
         // Sum up all partial charge histograms to each MPI process to calculate
         // the global charge slope.
         amrex::ParallelAllReduce::Sum(
-            charge_distribution.data(),
-            charge_distribution.size(),
+            charge_distribution,
             amrex::ParallelDescriptor::Communicator()
         );
 

src/particles/spacecharge/GatherAndPush.cpp

@@ -20,6 +20,7 @@
 #include <AMReX_SPACE.H>      // for AMREX_D_DECL
 #include <AMReX_GpuContainers.H>
 #include <AMReX_ParmParse.H>
+#include <AMReX_Reduce.H>
 
 namespace impactx::particles::spacecharge
 {
@@ -62,7 +63,7 @@ namespace impactx::particles::spacecharge
             charge_distribution = Deposit1D( pc, bin_min, bin_max, num_bins);
             bool const GetNumberDensity = true;
             impactx::particles::wakefields::DerivativeCharge1D(charge_distribution, charge_distribution_slope, bin_size, GetNumberDensity);
-            Qb_abs = bin_size * std::accumulate(charge_distribution.begin(), charge_distribution.end(), 0.0_rt);
+            Qb_abs = bin_size * amrex::Reduce::Sum(charge_distribution.size(), charge_distribution.data());
         }
 
         amrex::Real const * const beam_profile = charge_distribution.data();

src/particles/spacecharge/Gauss2p5dPush.cpp

@@ -15,6 +15,7 @@
 #include <AMReX_REAL.H>
 #include <AMReX_BLProfiler.H>
 #include <AMReX_GpuContainers.H>
+#include <AMReX_GpuParallelReduce.H>
 #include <AMReX_ParmParse.H>
 #include <ablastr/warn_manager/WarnManager.H>
 
@@ -221,8 +222,7 @@ namespace impactx::particles::spacecharge
         // Sum up all partial charge histograms to each MPI process to calculate
         // the global charge slope.
         amrex::ParallelAllReduce::Sum(
-            charge_distribution.data(),
-            charge_distribution.size(),
+            charge_distribution,
             amrex::ParallelDescriptor::Communicator()
         );
 

src/particles/wakefields/HandleWakefield.H

@@ -16,6 +16,10 @@
 #include "WakeConvolution.H"
 #include "WakePush.H"
 
+#include <AMReX_GpuContainers.H>
+#include <AMReX_GpuDevice.H>
+#include <AMReX_GpuParallelReduce.H>
+
 #include <cmath>
 #include <stdexcept>
 #include <vector>
@@ -101,13 +105,15 @@ namespace impactx::particles::wakefields
             // Sum up all partial charge histograms to one MPI process to calculate the global wakefield.
             // Once calculated, we will distribute convolved_wakefield back to every MPI process.
             amrex::ParallelReduce::Sum(
-                charge_distribution.data(),
-                charge_distribution.size(),
+                charge_distribution,
                 amrex::ParallelDescriptor::IOProcessorNumber(),
                 amrex::ParallelDescriptor::Communicator()
             );
 
-            amrex::Gpu::DeviceVector<amrex::Real> convolved_wakefield;
+            // Pre-size on every rank to the (global) wake length.  On the IO
+            // processor this is overwritten by the convolution result of the
+            // same length.
+            amrex::Gpu::DeviceVector<amrex::Real> convolved_wakefield(num_bins, 0.0);
             if (amrex::ParallelDescriptor::IOProcessor()) {
                 // Call the mean transverse position function
                 impactx::particles::wakefields::MeanTransversePosition(particle_container, mean_x, mean_y, bin_min,
@@ -137,21 +143,27 @@ namespace impactx::particles::wakefields
                 convolved_wakefield = impactx::particles::wakefields::convolve_fft(slopes, wake_function, bin_size);
             }
 
-            // Broadcast the global wakefield to every MPI rank
+            // Broadcast the global wakefield from the IO processor (which alone
+            // computed it) to every MPI rank. Receivers are pre-sized above.
             amrex::ParallelDescriptor::Bcast(
-                convolved_wakefield.data(),
-                convolved_wakefield.size(),
-                amrex::ParallelDescriptor::IOProcessorNumber()
+                convolved_wakefield,
+                amrex::ParallelDescriptor::IOProcessorNumber(),
+                amrex::ParallelDescriptor::Communicator()
             );
 
             // Check convolution output
             if (print_wakefield && amrex::ParallelDescriptor::IOProcessor())
             {
+                // convolved_wakefield is in (non-managed) device memory. Copy it
+                // to the host before iterating it on the host for printing.
+                amrex::Gpu::HostVector<amrex::Real> h_wakefield(convolved_wakefield.size());
+                amrex::Gpu::dtoh_memcpy(h_wakefield.data(), convolved_wakefield.data(),
+                                        convolved_wakefield.size() * sizeof(amrex::Real));
                 std::cout << "Convolved wakefield: ";
                 std::ofstream outfile("convolved_wakefield.txt");
-                for (double const i : convolved_wakefield) {
-                    std::cout << i << " ";
-                    outfile << i << std::endl;
+                for (amrex::Real const val : h_wakefield) {
+                    std::cout << val << " ";
+                    outfile << val << std::endl;
                 }
                 std::cout << std::endl;
                 outfile.close();