fix gpu version of computing cloud cover (#1671)

erf-model · Jul 9, 2024 · f13af91 · f13af91
1 parent 58617d8
commit f13af91
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 15 deletions.
diff --git a/Source/ERF.H b/Source/ERF.H
@@ -161,7 +161,8 @@ public:
     void sample_points (int lev, amrex::Real time, amrex::IntVect cell, amrex::MultiFab& mf);
     void sample_lines (int lev, amrex::Real time, amrex::IntVect cell, amrex::MultiFab& mf);
 
-    void derive_diag_profiles (amrex::Gpu::HostVector<amrex::Real>& h_avg_u   , amrex::Gpu::HostVector<amrex::Real>& h_avg_v  ,
+    void derive_diag_profiles (amrex::Real time,
+                               amrex::Gpu::HostVector<amrex::Real>& h_avg_u   , amrex::Gpu::HostVector<amrex::Real>& h_avg_v  ,
                                amrex::Gpu::HostVector<amrex::Real>& h_avg_w   , amrex::Gpu::HostVector<amrex::Real>& h_avg_rho,
                                amrex::Gpu::HostVector<amrex::Real>& h_avg_th  , amrex::Gpu::HostVector<amrex::Real>& h_avg_ksgs,
                                amrex::Gpu::HostVector<amrex::Real>& h_avg_kturb,

diff --git a/Source/IO/ERF_Write1DProfiles.cpp b/Source/IO/ERF_Write1DProfiles.cpp
@@ -34,7 +34,8 @@ ERF::write_1D_profiles (Real time)
         Gpu::HostVector<Real> h_avg_sgshfx, h_avg_sgsdiss; // only output tau_{theta,w} and epsilon for now
 
         if (NumDataLogs() > 1) {
-            derive_diag_profiles(h_avg_u, h_avg_v, h_avg_w,
+            derive_diag_profiles(time,
+                                 h_avg_u, h_avg_v, h_avg_w,
                                  h_avg_rho, h_avg_th, h_avg_ksgs, h_avg_kturb,
                                  h_avg_qv, h_avg_qc, h_avg_qr, h_avg_wqv, h_avg_wqc, h_avg_wqr, h_avg_qi, h_avg_qs, h_avg_qg,
                                  h_avg_uu, h_avg_uv, h_avg_uw, h_avg_vv, h_avg_vw, h_avg_ww,
@@ -179,18 +180,21 @@ ERF::write_1D_profiles (Real time)
  * @param h_avg_pw Profile for pressure perturbation * z-velocity on Host
  */
 
-void ERF::derive_diag_profiles(Gpu::HostVector<Real>& h_avg_u   , Gpu::HostVector<Real>& h_avg_v  , Gpu::HostVector<Real>& h_avg_w,
-                          Gpu::HostVector<Real>& h_avg_rho , Gpu::HostVector<Real>& h_avg_th , Gpu::HostVector<Real>& h_avg_ksgs,
-                          Gpu::HostVector<Real>& h_avg_kturb, Gpu::HostVector<Real>& h_avg_qv  , Gpu::HostVector<Real>& h_avg_qc , Gpu::HostVector<Real>& h_avg_qr,
-                          Gpu::HostVector<Real>& h_avg_wqv , Gpu::HostVector<Real>& h_avg_wqc, Gpu::HostVector<Real>& h_avg_wqr,
-                          Gpu::HostVector<Real>& h_avg_qi  , Gpu::HostVector<Real>& h_avg_qs , Gpu::HostVector<Real>& h_avg_qg,
-                          Gpu::HostVector<Real>& h_avg_uu  , Gpu::HostVector<Real>& h_avg_uv , Gpu::HostVector<Real>& h_avg_uw,
-                          Gpu::HostVector<Real>& h_avg_vv  , Gpu::HostVector<Real>& h_avg_vw , Gpu::HostVector<Real>& h_avg_ww,
-                          Gpu::HostVector<Real>& h_avg_uth , Gpu::HostVector<Real>& h_avg_vth, Gpu::HostVector<Real>& h_avg_wth,
-                          Gpu::HostVector<Real>& h_avg_thth,
-                          Gpu::HostVector<Real>& h_avg_uiuiu  , Gpu::HostVector<Real>& h_avg_uiuiv , Gpu::HostVector<Real>& h_avg_uiuiw,
-                          Gpu::HostVector<Real>& h_avg_p,
-                          Gpu::HostVector<Real>& h_avg_pu  , Gpu::HostVector<Real>& h_avg_pv , Gpu::HostVector<Real>& h_avg_pw)
+void ERF::derive_diag_profiles(Real time,
+                               Gpu::HostVector<Real>& h_avg_u   , Gpu::HostVector<Real>& h_avg_v  , Gpu::HostVector<Real>& h_avg_w,
+                               Gpu::HostVector<Real>& h_avg_rho , Gpu::HostVector<Real>& h_avg_th , Gpu::HostVector<Real>& h_avg_ksgs,
+                               Gpu::HostVector<Real>& h_avg_kturb, Gpu::HostVector<Real>& h_avg_qv,
+                               Gpu::HostVector<Real>& h_avg_qc , Gpu::HostVector<Real>& h_avg_qr,
+                               Gpu::HostVector<Real>& h_avg_wqv , Gpu::HostVector<Real>& h_avg_wqc, Gpu::HostVector<Real>& h_avg_wqr,
+                               Gpu::HostVector<Real>& h_avg_qi  , Gpu::HostVector<Real>& h_avg_qs , Gpu::HostVector<Real>& h_avg_qg,
+                               Gpu::HostVector<Real>& h_avg_uu  , Gpu::HostVector<Real>& h_avg_uv , Gpu::HostVector<Real>& h_avg_uw,
+                               Gpu::HostVector<Real>& h_avg_vv  , Gpu::HostVector<Real>& h_avg_vw , Gpu::HostVector<Real>& h_avg_ww,
+                               Gpu::HostVector<Real>& h_avg_uth , Gpu::HostVector<Real>& h_avg_vth, Gpu::HostVector<Real>& h_avg_wth,
+                               Gpu::HostVector<Real>& h_avg_thth,
+                               Gpu::HostVector<Real>& h_avg_uiuiu  , Gpu::HostVector<Real>& h_avg_uiuiv,
+                               Gpu::HostVector<Real>& h_avg_uiuiw,
+                               Gpu::HostVector<Real>& h_avg_p,
+                               Gpu::HostVector<Real>& h_avg_pu  , Gpu::HostVector<Real>& h_avg_pv , Gpu::HostVector<Real>& h_avg_pw)
 
 {
     // We assume that this is always called at level 0

diff --git a/Source/IO/ERF_WriteScalarProfiles.cpp b/Source/IO/ERF_WriteScalarProfiles.cpp
@@ -203,7 +203,15 @@ ERF::cloud_fraction (Real time)
     AMREX_ASSERT(numpts < Long(std::numeric_limits<int>::max));
 
 #if 1
-    ParallelDescriptor::ReduceIntMax(p,static_cast<int>(numpts));
+    if (amrex::UseGpuAwareMpi()) {
+        ParallelDescriptor::ReduceIntMax(p,static_cast<int>(numpts));
+    } else {
+        Gpu::PinnedVector<int> hv(numpts);
+        Gpu::copyAsync(Gpu::deviceToHost, p, p+numpts, hv.data());
+        Gpu::streamSynchronize();
+        ParallelDescriptor::ReduceIntMax(hv.data(),static_cast<int>(numpts));
+        Gpu::copyAsync(Gpu::hostToDevice, hv,data(), hv.data()+numpts, p);
+    }
 
     // Sum over component 0
     Long num_cloudy = qc_2d.template sum<RunOn::Device>(0);