Fix a gpu issue

If gpu aware mpi is not used, we need to use host memory for communication.
RevathiJambunathan · Jan 25, 2024 · c7a225e · c7a225e
1 parent edd6049
commit c7a225e
Showing 1 changed file with 28 additions and 3 deletions.
diff --git a/Source/Laser/LaserProfilesImpl/LaserProfileStationField.cpp b/Source/Laser/LaserProfilesImpl/LaserProfileStationField.cpp
@@ -143,11 +143,36 @@ void StationFieldLaserProfile::update (amrex::Real t)
             a(i,j,0) = (a(i,j,0)*(amrex::Real(1.0)-w) + a(i,j,1)*w);
 #endif
         });
+        amrex::Gpu::streamSynchronize();
     }
 
-    amrex::Box b = m_slice_fab->box();
-    b.setBig(AMREX_SPACEDIM-1,0);
-    amrex::ParallelDescriptor::Bcast(m_slice_fab->dataPtr(), b.numPts());
+    if (amrex::ParallelDescriptor::NProcs() > 1) {
+        amrex::Box b = m_slice_fab->box();
+        b.setBig(AMREX_SPACEDIM-1,0);
+#ifdef AMREX_USE_GPU
+        if ( ! amrex::ParallelDescriptor::UseGpuAwareMpi() &&
+             ! m_slice_fab->arena()->isHostAccessible() )
+        {
+            amrex::FArrayBox host_fab(b, 1, amrex::The_Pinned_Arena());
+            if (amrex::ParallelDescriptor::MyProc() == 0) {
+                amrex::Gpu::dtoh_memcpy_async(host_fab.dataPtr(),
+                                              m_slice_fab->dataPtr(),
+                                              host_fab.nBytes());
+                amrex::Gpu::streamSynchronize();
+            }
+            amrex::ParallelDescriptor::Bcast(host_fab.dataPtr(), b.numPts());
+            if (amrex::ParallelDescriptor::MyProc() != 0) {
+                amrex::Gpu::htod_memcpy_async(m_slice_fab->dataPtr(),
+                                              host_fab.dataPtr(),
+                                              host_fab.nBytes());
+                amrex::Gpu::streamSynchronize();
+            }
+        } else
+#endif
+        {
+            amrex::ParallelDescriptor::Bcast(m_slice_fab->dataPtr(), b.numPts());
+        }
+    }
 }
 
 void StationFieldLaserProfile::fill_amplitude (