Skip to content

Commit

Permalink
dotMask now owned by WarpX. Defined only when needed.
Browse files Browse the repository at this point in the history
  • Loading branch information
JustinRayAngus committed Aug 24, 2024
1 parent 95e19e6 commit 6ba8132
Show file tree
Hide file tree
Showing 4 changed files with 80 additions and 47 deletions.
10 changes: 0 additions & 10 deletions Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.H
Original file line number Diff line number Diff line change
Expand Up @@ -216,12 +216,6 @@ public:
// solver vec types are for now limited to warpx::fields::FieldType
warpx::fields::FieldType getVecType () const { return m_field_type; }

// clearDotMask() must be called by the highest class that owns WarpXSolverVec()
// after it is done being used ( typically in the destructor ) to avoid the
// following error message after the simulation finishes:
// malloc_consolidate(): unaligned fastbin chunk detected
static void clearDotMask() { m_dotMask.clear(); }

private:

bool m_is_defined = false;
Expand All @@ -235,10 +229,6 @@ private:
inline static WarpX* m_WarpX = nullptr;
void SetWarpXPointer( WarpX* a_WarpX );

inline static bool m_dot_mask_defined = false;
inline static amrex::Vector<std::array<std::unique_ptr<amrex::iMultiFab>,3>> m_dotMask;
void SetDotMask();

};

#endif
45 changes: 8 additions & 37 deletions Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ void WarpXSolverVec::Define ( WarpX* a_WarpX,
m_field_type = a_field_type;
m_is_defined = true;
SetWarpXPointer(a_WarpX);
SetDotMask();
}

void WarpXSolverVec::Copy ( const warpx::fields::FieldType a_field_type )
Expand Down Expand Up @@ -58,50 +57,22 @@ void WarpXSolverVec::SetWarpXPointer( WarpX* a_WarpX )
m_warpx_ptr_defined = true;
}

void WarpXSolverVec::SetDotMask()
{
if (m_dot_mask_defined) { return; }
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
IsDefined(),
"WarpXSolverVec::SetDotMask() called from undefined instance ");
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
m_warpx_ptr_defined,
"WarpXSolverVec::SetDotMask() called before SetWarpXPointer() ");

const amrex::Vector<amrex::Geometry>& Geom = m_WarpX->Geom();
m_dotMask.resize(m_num_amr_levels);
for ( int n = 0; n < 3; n++) {
const amrex::BoxArray& grids = m_field_vec[0][n]->boxArray();
const amrex::MultiFab tmp( grids, m_field_vec[0][n]->DistributionMap(),
1, 0, amrex::MFInfo().SetAlloc(false) );
const amrex::Periodicity& period = Geom[0].periodicity();
m_dotMask[0][n] = tmp.OwnerMask(period);
}
m_dot_mask_defined = true;

// If the function below is not called, then the following
// error message occurs after the simulation finishes:
// malloc_consolidate(): unaligned fastbin chunk detected
amrex::ExecOnFinalize(WarpXSolverVec::clearDotMask);
}

[[nodiscard]] amrex::Real WarpXSolverVec::dotProduct ( const WarpXSolverVec& a_X ) const
{
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
m_dot_mask_defined,
"WarpXSolverVec::dotProduct called with m_dotMask not yet defined");
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
a_X.m_field_type==m_field_type,
"WarpXSolverVec::dotProduct(X) called with solver vecs of different types");

amrex::Real result = 0.0;
const int lev = 0;
const bool local = true;
for (int n = 0; n < 3; ++n) {
auto rtmp = amrex::MultiFab::Dot( *m_dotMask[lev][n],
*m_field_vec[lev][n], 0,
*a_X.getVec()[lev][n], 0, 1, 0, local);
result += rtmp;
for (int lev = 0; lev < m_num_amr_levels; ++lev) {
for (int n = 0; n < 3; ++n) {
const amrex::iMultiFab* dotMask = m_WarpX->getFieldDotMaskPointer(m_field_type,lev,n);
auto rtmp = amrex::MultiFab::Dot( *dotMask,
*m_field_vec[lev][n], 0,
*a_X.getVec()[lev][n], 0, 1, 0, local);
result += rtmp;
}
}
amrex::ParallelAllReduce::Sum(result, amrex::ParallelContext::CommunicatorSub());
return result;
Expand Down
21 changes: 21 additions & 0 deletions Source/WarpX.H
Original file line number Diff line number Diff line change
Expand Up @@ -535,6 +535,21 @@ public:
[[nodiscard]] const amrex::Vector<std::array< std::unique_ptr<amrex::MultiFab>,3>>&
getMultiLevelField(warpx::fields::FieldType field_type) const;

/**
* \brief
* Get pointer to the amrex::MultiFab containing the dotMask for the specified field
*/
[[nodiscard]] const amrex::iMultiFab*
getFieldDotMaskPointer (warpx::fields::FieldType field_type, const int lev, const int dir) const;

/**
* \brief
* Set the dotMask container
*/
void SetDotMask( std::unique_ptr<amrex::iMultiFab>& field_dotMask,
warpx::fields::FieldType field_type,
const int lev, const int dir ) const;

[[nodiscard]] bool DoPML () const {return do_pml;}
[[nodiscard]] bool DoFluidSpecies () const {return do_fluid_species;}

Expand Down Expand Up @@ -1502,6 +1517,12 @@ private:
amrex::Vector<std::array< std::unique_ptr<amrex::MultiFab>, 3 > > Efield_avg_fp;
amrex::Vector<std::array< std::unique_ptr<amrex::MultiFab>, 3 > > Bfield_avg_fp;

// Masks for computing global moments of fields when using grids that have
// shared locations across different ranks (e.g., a Yee grid)
mutable amrex::Vector<std::array< std::unique_ptr<amrex::iMultiFab>,3 > > Efield_dotMask;
mutable amrex::Vector<std::array< std::unique_ptr<amrex::iMultiFab>,3 > > Bfield_dotMask;
mutable amrex::Vector< std::unique_ptr<amrex::iMultiFab> > phi_dotMask;

// Memory buffers for computing magnetostatic fields
// Vector Potential A and previous step. Time buffer needed for computing dA/dt to first order
amrex::Vector<std::array< std::unique_ptr<amrex::MultiFab>, 3 > > vector_potential_fp_nodal;
Expand Down
51 changes: 51 additions & 0 deletions Source/WarpX.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -326,6 +326,10 @@ WarpX::WarpX ()
Efield_fp.resize(nlevs_max);
Bfield_fp.resize(nlevs_max);

Efield_dotMask.resize(nlevs_max);
Bfield_dotMask.resize(nlevs_max);
phi_dotMask.resize(nlevs_max);

// Only allocate vector potential arrays when using the Magnetostatic Solver
if (electrostatic_solver_id == ElectrostaticSolverAlgo::LabFrameElectroMagnetostatic)
{
Expand Down Expand Up @@ -2077,6 +2081,9 @@ WarpX::ClearLevel (int lev)
Efield_fp [lev][i].reset();
Bfield_fp [lev][i].reset();

Efield_dotMask [lev][i].reset();
Bfield_dotMask [lev][i].reset();

current_store[lev][i].reset();

if (do_current_centering)
Expand Down Expand Up @@ -2123,6 +2130,8 @@ WarpX::ClearLevel (int lev)
G_cp [lev].reset();
rho_cp[lev].reset();

phi_dotMask[lev].reset();

#ifdef WARPX_USE_FFT
if (WarpX::electromagnetic_solver_id == ElectromagneticSolverAlgo::PSATD) {
spectral_solver_fp[lev].reset();
Expand Down Expand Up @@ -3615,3 +3624,45 @@ WarpX::getMultiLevelField(warpx::fields::FieldType field_type) const
return Efield_fp;
}
}

const amrex::iMultiFab*
WarpX::getFieldDotMaskPointer (warpx::fields::FieldType field_type, const int lev, const int dir) const
{
switch(field_type)
{
case FieldType::Efield_fp :
if (Efield_dotMask[lev][dir] == nullptr) {
SetDotMask( Efield_dotMask[lev][dir], field_type, lev, dir );
}
return Efield_dotMask[lev][dir].get();
case FieldType::Bfield_fp :
if (Bfield_dotMask[lev][dir] == nullptr) {
SetDotMask( Bfield_dotMask[lev][dir], field_type, lev, dir );
}
return Bfield_dotMask[lev][dir].get();
case FieldType::phi_fp :
if (phi_dotMask[lev] == nullptr) {
SetDotMask( phi_dotMask[lev], field_type, lev, 0 );
}
return phi_dotMask[lev].get();
default:
WARPX_ABORT_WITH_MESSAGE("Invalid field type for dotMask");
return Efield_dotMask[lev][dir].get();
}
}

void WarpX::SetDotMask( std::unique_ptr<amrex::iMultiFab>& field_dotMask,
warpx::fields::FieldType field_type,
const int lev, const int dir ) const
{
// Define the dot mask for this field_type needed to properly compute dotProduct()
// for field values that have shared locations on different MPI ranks

const amrex::MultiFab* this_field = getFieldPointer(field_type,lev,dir);
const amrex::BoxArray& this_ba = this_field->boxArray();
const amrex::MultiFab tmp( this_ba, this_field->DistributionMap(),
1, 0, amrex::MFInfo().SetAlloc(false) );
const amrex::Periodicity& period = Geom(lev).periodicity();
field_dotMask = tmp.OwnerMask(period);

}

0 comments on commit 6ba8132

Please sign in to comment.