tboudreaux/MeanField
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(allocations): reduced overall allocations by 95%, increaseed jacobian applicatin by 2x

Browse Source 
This commit uses global pre allocated work space to dramatically reduce memory usage and allocation time
This commit is contained in:
Emily Boudreaux
2026-09-10 06:50:56 -04:00
parent b3c04d507a
commit 75cc638739
66 changed files with 207183 additions and 99552 deletions
Download Patch File Download Diff File Expand all files Collapse all files

218
libmeanfield/interface/solver/stellar_context.cppm
View File

@@ -1,5 +1,6 @@
module;
#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
@@ -8,6 +9,7 @@ module;
#include <cstdint>
#include <exception>
#include <expected>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
@@ -17,6 +19,7 @@ module;
#include <string_view>
#include <type_traits>
#include <utility>
#include <vector>
#include <mfem.hpp>
#include <mpi.h>
@@ -37,6 +40,11 @@ export namespace mean_field::solver {
template <typename Context, typename NewtonConfiguration, typename Observer> class StellarEquilibriumSolver;
} // namespace mean_field::solver
export namespace mean_field::solver::detail {
// Internal, synchronous experiment access; not a stable solver API.
struct StellarEquilibriumContextDiagnostics;
}
namespace mean_field::solver::detail {
template <typename Model, typename Discretization>
using StellarContextProblem =
@@ -256,6 +264,7 @@ export namespace mean_field::solver {
private:
template <typename, typename, typename> friend class StellarEquilibriumSolver;
friend struct detail::StellarEquilibriumContextAssembly;
friend struct detail::StellarEquilibriumContextDiagnostics;
using NormalizedOperatorType = detail::StellarContextNormalizedOperator<ProblemType>;
using PhysicalInverseType = detail::StellarContextPhysicalInverse<PreconditionerPrescriptionType, ProblemType>;
@@ -323,11 +332,16 @@ export namespace mean_field::solver {
trialNormalizedResidual(RequireProblem(storage).EquationSize()),
linearRightHandSide(RequireProblem(storage).EquationSize()),
normalizedCorrection(RequireProblem(storage).StateSize()),
physicalCorrection(RequireProblem(storage).StateSize()),
volumeDisplacementDirection(
RequireProblem(storage).GetPhysicalOperator().GetDomainDeformation().volumeDisplacementSize()
),
candidatePhysicalState(RequireProblem(storage).StateSize()),
normalizedOperator(std::make_unique<NormalizedOperatorType>(RequireProblem(storage))) {
ValidateInitialState();
InitializeWorkspaces();
PrepareInitialOperator();
InitializeGeometryPreflightRules();
physicalInverse = std::unique_ptr<PhysicalInverseType>{new PhysicalInverseType(
PreparePhysicalInverse(std::move(preconditionerPrescription), RequireProblem(storage))
@@ -375,6 +389,9 @@ export namespace mean_field::solver {
trialNormalizedResidual.Size() == problem->EquationSize() &&
linearRightHandSide.Size() == problem->EquationSize() &&
normalizedCorrection.Size() == problem->StateSize() &&
physicalCorrection.Size() == problem->StateSize() &&
volumeDisplacementDirection.Size() ==
problem->GetPhysicalOperator().GetDomainDeformation().volumeDisplacementSize() &&
candidatePhysicalState.Size() == problem->StateSize() &&
storage->physicalState->Size() == problem->StateSize();
}
@@ -464,6 +481,26 @@ export namespace mean_field::solver {
}
}
[[nodiscard]] deformation::LargestSafeNewtonStepSizeEstimate EstimateLargestSafeStepSize(
const double maximumStepSize,
const double fractionToBoundarySafety
) {
normalizedOperator->DenormalizeState(normalizedCorrection, physicalCorrection);
const auto &physicalOperator = Problem().GetPhysicalOperator();
Problem().BuildVolumeDisplacementDirection(physicalCorrection, volumeDisplacementDirection);
const fem::FEM &finiteElements =
equilibrium::detail::StellarEquilibriumProblemFactory::FiniteElementModel(Problem());
return deformation::estimate_largest_safe_newton_step_size(
Problem().GetDiscretization().domainMapper(), *finiteElements.displacementFes,
*finiteElements.compactificationCoordinate, physicalOperator.GetGeneratedVolumeDisplacement(),
volumeDisplacementDirection, geometryPreflightRules,
{.maximumStepSize = maximumStepSize,
.determinantFloor = 0.0,
.fractionToBoundarySafety = fractionToBoundarySafety}
);
}
std::shared_ptr<Storage> storage;
DependencyLedger dependencyLedger;
mfem::Vector acceptedNormalizedState;
@@ -472,7 +509,10 @@ export namespace mean_field::solver {
mfem::Vector trialNormalizedResidual;
mfem::Vector linearRightHandSide;
mfem::Vector normalizedCorrection;
mfem::Vector physicalCorrection;
mfem::Vector volumeDisplacementDirection;
mfem::Vector candidatePhysicalState;
std::vector<deformation::NewtonStepGeometryRule> geometryPreflightRules;
double acceptedMinimumJacobianDeterminant{std::numeric_limits<double>::quiet_NaN()};
std::unique_ptr<NormalizedOperatorType> normalizedOperator;
std::unique_ptr<PhysicalInverseType> physicalInverse;
@@ -565,12 +605,14 @@ export namespace mean_field::solver {
void InitializeWorkspaces() {
normalizedOperator->NormalizeState(AcceptedPhysicalState(), acceptedNormalizedState);
trialNormalizedState = acceptedNormalizedState;
acceptedNormalizedResidual = 0.0;
trialNormalizedResidual = 0.0;
linearRightHandSide = 0.0;
normalizedCorrection = 0.0;
candidatePhysicalState = AcceptedPhysicalState();
trialNormalizedState = acceptedNormalizedState;
acceptedNormalizedResidual = 0.0;
trialNormalizedResidual = 0.0;
linearRightHandSide = 0.0;
normalizedCorrection = 0.0;
physicalCorrection = 0.0;
volumeDisplacementDirection = 0.0;
candidatePhysicalState = AcceptedPhysicalState();
}
void PrepareInitialOperator() {
@@ -582,6 +624,97 @@ export namespace mean_field::solver {
trialNormalizedResidual = acceptedNormalizedResidual;
}
void AppendGeometryPreflightRule(
const int element,
const mfem::IntegrationRule &integrationRule
) {
geometryPreflightRules.push_back({.element = element, .integrationRule = &integrationRule});
}
void InitializeGeometryPreflightRules() {
const ProblemType &problem = Problem();
const fem::FEM &finiteElements =
equilibrium::detail::StellarEquilibriumProblemFactory::FiniteElementModel(problem);
if (finiteElements.mesh == nullptr || finiteElements.displacementFes == nullptr ||
finiteElements.compactificationCoordinate == nullptr) {
throw std::logic_error(
"The Newton geometry preflight requires complete displacement geometry data."
);
}
if (!problem.GetPhysicalOperator().GetDomainDeformation().descriptor().linearOnReferenceGeometry) {
throw std::invalid_argument(
"The Newton geometry preflight requires a domain deformation that is linear on the "
"reference geometry."
);
}
geometryPreflightRules.clear();
geometryPreflightRules.reserve(
static_cast<std::size_t>(finiteElements.mesh->GetNE()) * static_cast<std::size_t>(10)
);
const int dimension = problem.GetDiscretization().domainMapper().GetDimension();
for (int element = 0; element < finiteElements.mesh->GetNE(); ++element) {
const mfem::FiniteElement *finiteElement = finiteElements.displacementFes->GetFE(element);
mfem::ElementTransformation *transformation =
finiteElements.mesh->GetElementTransformation(element);
if (finiteElement == nullptr || transformation == nullptr) {
throw std::logic_error(
"The Newton geometry preflight encountered incomplete element geometry data."
);
}
const int geometryInspectionOrder =
std::max(finiteElement->GetOrder() + 2, 2 * dimension * finiteElement->GetOrder());
AppendGeometryPreflightRule(
element, mfem::IntRules.Get(transformation->GetGeometryType(), geometryInspectionOrder)
);
}
const auto appendPreparedRules = [this](const auto &preparedOperator) {
preparedOperator.VisitMappedGeometryRules(
[this](const int element, const mfem::IntegrationRule &integrationRule) {
AppendGeometryPreflightRule(element, integrationRule);
}
);
};
const auto &physicalOperator = problem.GetPhysicalOperator();
const auto &gravityGeometry = physicalOperator.GetGravityContext().GetGeometryContext();
appendPreparedRules(gravityGeometry.GetMassOperator());
appendPreparedRules(gravityGeometry.GetSourceOperator());
appendPreparedRules(physicalOperator.GetBarotropicClosureOperator());
appendPreparedRules(physicalOperator.GetHydrostaticOperator());
const auto &displacementOperator = physicalOperator.GetDisplacementOperator();
appendPreparedRules(displacementOperator.GetPressureOperator());
appendPreparedRules(displacementOperator.GetGravityOperator());
appendPreparedRules(displacementOperator.GetRotationalOperator());
appendPreparedRules(physicalOperator.GetMassNormalizationOperator());
if constexpr (ProblemType::hasFixedAngularMomentum) {
appendPreparedRules(problem.GetPreparedOperator().GetAngularMomentumConstraint());
}
const auto ruleLess = [](const deformation::NewtonStepGeometryRule &left,
const deformation::NewtonStepGeometryRule &right) {
if (left.element != right.element) {
return left.element < right.element;
}
return std::less<const mfem::IntegrationRule *>{}(left.integrationRule, right.integrationRule);
};
std::sort(geometryPreflightRules.begin(), geometryPreflightRules.end(), ruleLess);
geometryPreflightRules.erase(
std::unique(
geometryPreflightRules.begin(), geometryPreflightRules.end(),
[](const deformation::NewtonStepGeometryRule &left,
const deformation::NewtonStepGeometryRule &right) {
return left.element == right.element && left.integrationRule == right.integrationRule;
}
),
geometryPreflightRules.end()
);
}
[[nodiscard]] auto PrepareOperator(const mfem::Vector &normalizedState) {
if constexpr (ProblemType::generatedRotationProviderCount == 0) {
return normalizedOperator->Prepare(
@@ -695,6 +828,42 @@ export namespace mean_field::solver {
concept StellarEquilibriumContextType = IsStellarEquilibriumContext<std::remove_cvref_t<Candidate>>::value;
} // namespace mean_field::solver
export namespace mean_field::solver::detail {
struct StellarEquilibriumContextDiagnostics final {
// The callback must not retain references to runtime storage. It may
// prepare trial states, but accepted vectors must remain unchanged.
// Restore the production preparation and correction on every exit.
template <typename Context, typename Callback>
static void WithState(Context &context, Callback &&callback) {
if (context.hasActiveSolver() || !context.isReady()) {
throw std::logic_error("Diagnostics require a ready context with no active solver.");
}
auto &state = *context.m_state;
mfem::Vector savedCorrection(state.normalizedCorrection);
context.AcquireSolver();
try {
state.BeginEvaluation();
std::invoke(std::forward<Callback>(callback), state,
equilibrium::detail::StellarEquilibriumProblemFactory::FiniteElementModel(state.Problem()));
state.RestoreAccepted();
state.normalizedCorrection = savedCorrection;
context.ReleaseSolver();
} catch (...) {
const auto original = std::current_exception();
try {
state.RestoreAccepted();
state.normalizedCorrection = savedCorrection;
} catch (...) {
context.ReleaseSolver();
throw;
}
context.ReleaseSolver();
std::rethrow_exception(original);
}
}
};
}
namespace mean_field::solver::detail {
[[nodiscard]] inline physics::RigidRotation ZeroRigidRotation() {
mfem::Vector angularVelocity(3);
@@ -973,12 +1142,14 @@ export namespace mean_field::solver {
struct IterationTimings final {
Clock::time_point start{};
double geometryPreflightSeconds{0.0};
double lineSearchSeconds{0.0};
double trialPreparationSeconds{0.0};
double metricEvaluationSeconds{0.0};
double preconditionerRefreshSeconds{0.0};
double rollbackSeconds{0.0};
double observerSeconds{0.0};
std::optional<deformation::LargestSafeNewtonStepSizeEstimate> geometryPreflight;
};
public:
@@ -1100,14 +1271,39 @@ export namespace mean_field::solver {
}
const nonlinear::MetricEvaluation previousMetric = acceptedMetric;
bool accepted = false;
double acceptedStepLength = 0.0;
int lineSearchTrials = 0;
const Clock::time_point geometryPreflightStart = Clock::now();
timings.geometryPreflight = state.EstimateLargestSafeStepSize(
nextLineSearchStepLength, options.backtracking.fractionToBoundarySafety
);
timings.geometryPreflightSeconds =
std::chrono::duration<double>(Clock::now() - geometryPreflightStart).count();
diagnostics.totalGeometryPreflightSeconds += timings.geometryPreflightSeconds;
diagnostics.lastGeometryPreflight = timings.geometryPreflight;
if (timings.geometryPreflight->limitedByGeometry) {
++diagnostics.geometryLimitedIterations;
}
if (timings.geometryPreflight->stepSize < options.backtracking.minimumStepLength) {
diagnostics.finalResidualNorm = acceptedMetric.residualNorm;
NotifyAfter(
iteration, nonlinear::IterationDisposition::globalization_failure, false, 0.0, 0,
diagnostics.initialResidualNorm, previousMetric, acceptedMetric, linearReport, timings, state
);
return Failure(
state, std::move(diagnostics), StellarEquilibriumFailureReason::globalization_failure,
"The geometry preflight found no orientation-preserving Newton step at or above the "
"configured minimum step length."
);
}
bool accepted = false;
double acceptedStepLength = 0.0;
int lineSearchTrials = 0;
nonlinear::MetricEvaluation trialMetric{};
StellarEquilibriumFailureReason rejectionReason =
StellarEquilibriumFailureReason::globalization_failure;
std::string rejectionMessage = "The backtracking line search found no acceptable Newton step.";
double stepLength = nextLineSearchStepLength;
double stepLength = timings.geometryPreflight->stepSize;
const Clock::time_point lineSearchStart = Clock::now();
try {
@@ -1553,11 +1749,13 @@ export namespace mean_field::solver {
.relativeResidualNorm = RelativeResidual(metric.residualNorm, initialResidualNorm),
.merit = metric.merit,
.iterationSeconds = DurationExcludingObserver(timings.start, timings.observerSeconds),
.geometryPreflightSeconds = timings.geometryPreflightSeconds,
.lineSearchSeconds = timings.lineSearchSeconds,
.trialPreparationSeconds = timings.trialPreparationSeconds,
.metricEvaluationSeconds = timings.metricEvaluationSeconds,
.preconditionerRefreshSeconds = timings.preconditionerRefreshSeconds,
.rollbackSeconds = timings.rollbackSeconds,
.geometryPreflight = timings.geometryPreflight,
.linearSolve = linearReport,
.communicator = state.Problem().GetCommunicator(),
.physicalState = detail::ReadOnlySpan(state.AcceptedPhysicalState()),