Skip to content

Commit

Permalink
Initial work to use MF registry to write checkpoints
Browse files Browse the repository at this point in the history 
Note that this does not yet compile, since it has a duplicate method API
  • Loading branch information
@dpgrote
dpgrote committed Oct 10, 2024
1 parent 0337ee1 commit 2f7a450
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 10 deletions.
34 changes: 34 additions & 0 deletions Source/ablastr/fields/MultiFabRegister.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,10 @@ namespace ablastr::fields
/** the MR level of this (i)MultiFab */
int m_level = 0;

/** Name of the MultiFab in a checkpoint file.
* If the name is an empty string, it is not written out. */
std::string m_checkpoint_name = "";

/** remake distribution map on load balance, @see amrex::AmrCore::RemakeLevel */
bool m_remake = true;

Expand Down Expand Up @@ -180,6 +184,7 @@ namespace ablastr::fields
* @param ncomp the number of components of the field (all with the same staggering)
* @param ngrow the number of guard (ghost, halo) cells
* @param initial_value the optional initial value
* @param checkpoint_name name of the MultiFab in a checkpoint file
* @param remake follow the default domain decomposition of the simulation
* @param redistribute_on_remake redistribute on @see amrex::AmrCore::RemakeLevel
* @return pointer to newly allocated MultiFab
Expand All @@ -194,6 +199,7 @@ namespace ablastr::fields
int ncomp,
amrex::IntVect const & ngrow,
std::optional<amrex::Real const> initial_value = std::nullopt,
std::string const& checkpoint_name = "",
bool remake = true,
bool redistribute_on_remake = true
)
Expand All @@ -206,6 +212,7 @@ namespace ablastr::fields
ncomp,
ngrow,
initial_value,
checkpoint_name,
remake,
redistribute_on_remake
);
Expand All @@ -224,6 +231,7 @@ namespace ablastr::fields
* @param ncomp the number of components of the field (all with the same staggering)
* @param ngrow the number of guard (ghost, halo) cells
* @param initial_value the optional initial value
* @param checkpoint_name name of the MultiFab in a checkpoint file
* @param remake follow the default domain decomposition of the simulation
* @param redistribute_on_remake redistribute on @see amrex::AmrCore::RemakeLevel
* @return pointer to newly allocated MultiFab
Expand All @@ -239,6 +247,7 @@ namespace ablastr::fields
int ncomp,
amrex::IntVect const & ngrow,
std::optional<amrex::Real const> initial_value = std::nullopt,
std::string const& checkpoint_name = "",
bool remake = true,
bool redistribute_on_remake = true
)
Expand All @@ -252,6 +261,7 @@ namespace ablastr::fields
ncomp,
ngrow,
initial_value,
checkpoint_name,
remake,
redistribute_on_remake
);
Expand Down Expand Up @@ -635,6 +645,28 @@ namespace ablastr::fields
amrex::DistributionMapping const & new_dm
);

/** Write data for MultiFabs to checkpoint file
*
* @param checkpoint_file_name checkpoint file name
* @param default_level_prefix prefix added to checkpoint names for the level
*/
void
write_to_checkpoint(
std::string const& checkpoint_file_name,
std::string const& default_level_prefix
);

/** Read data for MultiFabs from checkpoint file
*
* @param checkpoint_file_name checkpoint file name
* @param default_level_prefix prefix added to checkpoint names for the level
*/
void
read_from_checkpoint(
std::string const& checkpoint_file_name,
std::string const& default_level_prefix
);

/** Create the register name of scalar field and MR level
*
* @param name the name of the field
Expand Down Expand Up @@ -680,6 +712,7 @@ namespace ablastr::fields
int ncomp,
amrex::IntVect const & ngrow,
std::optional<amrex::Real const> initial_value = std::nullopt,
std::string const& checkpoint_name = "",
bool remake = true,
bool redistribute_on_remake = true
);
Expand All @@ -693,6 +726,7 @@ namespace ablastr::fields
int ncomp,
amrex::IntVect const & ngrow,
std::optional<amrex::Real const> initial_value = std::nullopt,
std::string const& checkpoint_name = "",
bool remake = true,
bool redistribute_on_remake = true
);
Expand Down
55 changes: 45 additions & 10 deletions Source/ablastr/fields/MultiFabRegister.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
#include "MultiFabRegister.H"

#include <AMReX_MakeType.H>
#include <AMReX_PlotFileUtil.H>
#include <AMReX_VisMF.H>

#include <array>
#include <memory>
Expand All @@ -29,6 +31,7 @@ namespace ablastr::fields
int ncomp,
amrex::IntVect const & ngrow,
std::optional<amrex::Real const> initial_value,
std::string const& checkpoint_name,
bool remake,
bool redistribute_on_remake
)
Expand All @@ -49,6 +52,7 @@ namespace ablastr::fields
{ba, dm, ncomp, ngrow, tag},
std::nullopt, // scalar: no direction
level,
checkpoint_name,
remake,
redistribute_on_remake,
"" // we own the memory
Expand Down Expand Up @@ -79,6 +83,7 @@ namespace ablastr::fields
int ncomp,
amrex::IntVect const & ngrow,
std::optional<amrex::Real const> initial_value,
std::string const& checkpoint_name,
bool remake,
bool redistribute_on_remake
)
Expand All @@ -103,6 +108,7 @@ namespace ablastr::fields
{ba, dm, ncomp, ngrow, tag},
dir,
level,
checkpoint_name,
remake,
redistribute_on_remake,
"" // we own the memory
Expand Down Expand Up @@ -161,6 +167,7 @@ namespace ablastr::fields
{mf_alias, amrex::make_alias, 0, mf_alias.nComp()},
std::nullopt, // scalar: no direction
level,
"",
alias.m_remake,
alias.m_redistribute_on_remake,
internal_alias_name
Expand Down Expand Up @@ -221,6 +228,7 @@ namespace ablastr::fields
{mf_alias, amrex::make_alias, 0, mf_alias.nComp()},
dir,
level,
"",
alias.m_remake,
alias.m_redistribute_on_remake,
internal_alias_name
Expand Down Expand Up @@ -274,25 +282,52 @@ namespace ablastr::fields
mf_owner.m_mf = std::move(new_mf);
}
}
}

// Aliases
void
MultiFabRegister::write_to_checkpoint (
std::string const& checkpoint_file_name,
std::string const& default_level_prefix
)
{
for (auto & element : m_mf_register )
{
MultiFabOwner & mf_owner = element.second;

// keep distribution map as it is?
if (!mf_owner.m_remake) {
continue;
// Only write MultiFabs with a defined checkpoint name
if (!mf_owner.m_checkpoint_name.empty()) {

amrex::VisMF::Write(mf_owner.m_mf,
amrex::MultiFabFileFullPrefix(mf_owner.m_level,
checkpoint_file_name,
default_level_prefix,
mf_owner.m_checkpoint_name));

}
}
}

if (mf_owner.m_level == level && mf_owner.is_alias()) {
const amrex::MultiFab & mf = m_mf_register[mf_owner.m_owner].m_mf;
amrex::MultiFab new_mf(mf, amrex::make_alias, 0, mf.nComp());
void
MultiFabRegister::read_from_checkpoint (
std::string const& checkpoint_file_name,
std::string const& default_level_prefix
)
{
for (auto & element : m_mf_register )
{
MultiFabOwner & mf_owner = element.second;

// no copy via Redistribute: the owner was already redistributed
// Only read MultiFabs with a defined checkpoint name
if (!mf_owner.m_checkpoint_name.empty()) {

mf_owner.m_mf.setVal(0.);

amrex::VisMF::Read(mf_owner.m_mf,
amrex::MultiFabFileFullPrefix(mf_owner.m_level,
checkpoint_file_name,
default_level_prefix,
mf_owner.m_checkpoint_name));

// replace old MultiFab with new one, deallocate old one
mf_owner.m_mf = std::move(new_mf);
}
}
}
Expand Down

0 comments on commit 2f7a450

Please sign in to comment.