diff options
Diffstat (limited to 'Source/FieldSolver/SpectralSolver')
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp | 56 |
1 files changed, 56 insertions, 0 deletions
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp index 645d605c8..f62eaaa16 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp @@ -398,6 +398,8 @@ SpectralFieldDataRZ::ForwardTransform (const int lev, amrex::MultiFab const & field_mf, int const field_index, int const i_comp) { + amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(lev); + // Check field index type, in order to apply proper shift in spectral space. // Only cell centered in r is supported. bool const is_nodal_z = field_mf.is_nodal(1); @@ -419,6 +421,12 @@ SpectralFieldDataRZ::ForwardTransform (const int lev, // Loop over boxes. for (amrex::MFIter mfi(field_mf); mfi.isValid(); ++mfi){ + if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers) + { + amrex::Gpu::synchronize(); + } + amrex::Real wt = amrex::second(); + // Perform the Hankel transform first. // tempHTransformedSplit includes the imaginary component of mode 0. // field_mf does not. @@ -434,6 +442,12 @@ SpectralFieldDataRZ::ForwardTransform (const int lev, FABZForwardTransform(mfi, realspace_bx, tempHTransformedSplit, field_index, is_nodal_z); + if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers) + { + amrex::Gpu::synchronize(); + wt = amrex::second() - wt; + amrex::HostDevice::Atomic::Add( &(*cost)[mfi.index()], wt); + } } } @@ -445,6 +459,8 @@ SpectralFieldDataRZ::ForwardTransform (const int lev, amrex::MultiFab const & field_mf_r, int const field_index_r, amrex::MultiFab const & field_mf_t, int const field_index_t) { + amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(lev); + // Check field index type, in order to apply proper shift in spectral space. // Only cell centered in r is supported. bool const is_nodal_z = field_mf_r.is_nodal(1); @@ -461,6 +477,12 @@ SpectralFieldDataRZ::ForwardTransform (const int lev, // Loop over boxes. for (amrex::MFIter mfi(field_mf_r); mfi.isValid(); ++mfi){ + if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers) + { + amrex::Gpu::synchronize(); + } + amrex::Real wt = amrex::second(); + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); if ( !(field_mf_r[mfi].box().contains(field_mf_r_copy[mfi].box())) ) { @@ -484,6 +506,12 @@ SpectralFieldDataRZ::ForwardTransform (const int lev, FABZForwardTransform(mfi, realspace_bx, tempHTransformedSplit_p, field_index_r, is_nodal_z); FABZForwardTransform(mfi, realspace_bx, tempHTransformedSplit_m, field_index_t, is_nodal_z); + if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers) + { + amrex::Gpu::synchronize(); + wt = amrex::second() - wt; + amrex::HostDevice::Atomic::Add( &(*cost)[mfi.index()], wt); + } } } @@ -494,6 +522,8 @@ SpectralFieldDataRZ::BackwardTransform (const int lev, amrex::MultiFab& field_mf, int const field_index, int const i_comp) { + amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(lev); + // Check field index type, in order to apply proper shift in spectral space. bool const is_nodal_z = field_mf.is_nodal(1); @@ -509,6 +539,12 @@ SpectralFieldDataRZ::BackwardTransform (const int lev, // Loop over boxes. for (amrex::MFIter mfi(field_mf); mfi.isValid(); ++mfi){ + if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers) + { + amrex::Gpu::synchronize(); + } + amrex::Real wt = amrex::second(); + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); FABZBackwardTransform(mfi, realspace_bx, field_index, tempHTransformedSplit, is_nodal_z); @@ -519,6 +555,12 @@ SpectralFieldDataRZ::BackwardTransform (const int lev, multi_spectral_hankel_transformer[mfi].SpectralToPhysical_Scalar(tempHTransformedSplit[mfi], field_mf_copy[mfi]); field_mf[mfi].copy<amrex::RunOn::Device>(field_mf_copy[mfi], 0, i_comp*ncomp, ncomp); + if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers) + { + amrex::Gpu::synchronize(); + wt = amrex::second() - wt; + amrex::HostDevice::Atomic::Add( &(*cost)[mfi.index()], wt); + } } } @@ -529,6 +571,8 @@ SpectralFieldDataRZ::BackwardTransform (const int lev, amrex::MultiFab& field_mf_r, int const field_index_r, amrex::MultiFab& field_mf_t, int const field_index_t) { + amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(lev); + // Check field index type, in order to apply proper shift in spectral space. bool const is_nodal_z = field_mf_r.is_nodal(1); @@ -543,6 +587,12 @@ SpectralFieldDataRZ::BackwardTransform (const int lev, // Loop over boxes. for (amrex::MFIter mfi(field_mf_r); mfi.isValid(); ++mfi){ + if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers) + { + amrex::Gpu::synchronize(); + } + amrex::Real wt = amrex::second(); + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); FABZBackwardTransform(mfi, realspace_bx, field_index_r, tempHTransformedSplit_p, is_nodal_z); @@ -572,6 +622,12 @@ SpectralFieldDataRZ::BackwardTransform (const int lev, } }); + if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers) + { + amrex::Gpu::synchronize(); + wt = amrex::second() - wt; + amrex::HostDevice::Atomic::Add( &(*cost)[mfi.index()], wt); + } } } |