perf(jacobian-action): major updates to jacobian action application by removing redudant quadrature work. ~5x increase in speed

This commit is contained in:
2026-09-02 17:01:50 -04:00
parent 85500fef3b
commit 25510008dd
74 changed files with 8967 additions and 814 deletions

View File

@@ -0,0 +1,259 @@
module;
#include <cstdint>
#include <span>
#include <string>
#include <string_view>
#include <vector>
#include <mfem.hpp>
#include <mpi.h>
export module mean_field:solver.preconditioning_diagnostics;
export import :operators.root_manifest;
export namespace mean_field::solver {
struct OperatorApplicationStatistics final {
std::uint64_t applications{0};
double totalSeconds{0.0};
double maximumSeconds{0.0};
};
struct PreconditionerLifecycleStatistics final {
std::uint64_t setups{0};
std::uint64_t refreshes{0};
double setupSeconds{0.0};
double refreshSeconds{0.0};
};
/*
* A non-owning measurement wrapper. Statistics are local to an MPI rank;
* cross-rank wall-clock reductions are performed when a solve report is
* assembled. Krylov application is sequential, so counters intentionally
* do not impose atomic overhead.
*/
class InstrumentedOperator final : public mfem::Operator {
public:
explicit InstrumentedOperator(const mfem::Operator &operation);
void Mult(
const mfem::Vector &input,
mfem::Vector &output
) const override;
void ResetStatistics() const noexcept;
[[nodiscard]] const OperatorApplicationStatistics &GetStatistics() const noexcept;
[[nodiscard]] const mfem::Operator &GetOperation() const noexcept;
private:
const mfem::Operator *m_operation;
mutable OperatorApplicationStatistics m_statistics;
};
class InstrumentedPreconditioner final : public mfem::Solver {
public:
explicit InstrumentedPreconditioner(mfem::Solver &preconditioner);
void SetOperator(const mfem::Operator &operation) override;
void Mult(
const mfem::Vector &input,
mfem::Vector &output
) const override;
void ResetStatistics() const noexcept;
[[nodiscard]] const OperatorApplicationStatistics &GetStatistics() const noexcept;
[[nodiscard]] const PreconditionerLifecycleStatistics &GetLifecycleStatistics() const noexcept;
[[nodiscard]] const mfem::Solver &GetPreconditioner() const noexcept;
private:
mfem::Solver *m_preconditioner;
mutable OperatorApplicationStatistics m_statistics;
PreconditionerLifecycleStatistics m_lifecycleStatistics;
};
class IdentityPreconditioner final : public mfem::Solver {
public:
explicit IdentityPreconditioner(int size);
void SetOperator(const mfem::Operator &operation) override;
void Mult(
const mfem::Vector &input,
mfem::Vector &output
) const override;
};
/*
* If the supplied solver applies M^{-1}, this operator represents the
* fixed right-preconditioned product J M^{-1}. It is deliberately
* independent of the Krylov implementation used in production.
*/
class FixedRightPreconditionedOperator final : public mfem::Operator {
public:
FixedRightPreconditionedOperator(
const mfem::Operator &jacobian,
const mfem::Solver &inversePreconditioner
);
void Mult(
const mfem::Vector &input,
mfem::Vector &output
) const override;
[[nodiscard]] const mfem::Operator &GetJacobian() const noexcept;
[[nodiscard]] const mfem::Solver &GetInversePreconditioner() const noexcept;
private:
const mfem::Operator *m_jacobian;
const mfem::Solver *m_inversePreconditioner;
mutable mfem::Vector m_preconditionedDirection;
};
struct IterationResidualMeasurement final {
int iteration;
double reportedNorm;
bool final;
};
class ResidualHistoryMonitor final : public mfem::IterativeSolverMonitor {
public:
void Reset() override;
void MonitorResidual(
int iteration,
double norm,
const mfem::Vector &residual,
bool final
) override;
[[nodiscard]] const std::vector<IterationResidualMeasurement> &GetHistory() const noexcept;
private:
std::vector<IterationResidualMeasurement> m_history;
};
struct ResidualBlockMeasurement final {
std::string stableId;
int size{0};
double descriptorScale{1.0};
double rightHandSideNorm{0.0};
double trueResidualNorm{0.0};
double blockRelativeResidual{0.0};
double scaledRightHandSideNorm{0.0};
double scaledTrueResidualNorm{0.0};
double contributionToGlobalRelativeResidual{0.0};
double fractionOfGlobalSquaredResidualNorm{0.0};
};
struct DirectResidualMeasurement final {
double rightHandSideNorm{0.0};
double trueResidualNorm{0.0};
double relativeResidual{0.0};
std::vector<ResidualBlockMeasurement> blocks;
};
[[nodiscard]] DirectResidualMeasurement measureDirectResidual(
const mfem::Operator &jacobian,
const mfem::Vector &rightHandSide,
const mfem::Vector &solution,
std::span<const operators::RootBlockDescriptor> residualBlocks,
MPI_Comm communicator,
double denominatorFloor = 1.0e-300
);
struct LinearSolveMeasurement final {
bool solverConverged{false};
int outerIterations{0};
double solverReportedInitialNorm{0.0};
double solverReportedFinalNorm{0.0};
double solverReportedResidualReduction{0.0};
double trueResidualDigitsReducedPerJacobianApplication{0.0};
double solveSecondsMaximumRank{0.0};
OperatorApplicationStatistics jacobian;
OperatorApplicationStatistics inversePreconditioner;
PreconditionerLifecycleStatistics inversePreconditionerLifecycle;
DirectResidualMeasurement directResidual;
std::vector<IterationResidualMeasurement> reportedResidualHistory;
};
[[nodiscard]] LinearSolveMeasurement measureLinearSolve(
const mfem::IterativeSolver &iterativeSolver,
const mfem::Operator &jacobian,
const mfem::Vector &rightHandSide,
const mfem::Vector &solution,
std::span<const operators::RootBlockDescriptor> residualBlocks,
const OperatorApplicationStatistics &jacobianStatistics,
const OperatorApplicationStatistics &inversePreconditionerStatistics,
const PreconditionerLifecycleStatistics &inversePreconditionerLifecycle,
const ResidualHistoryMonitor &monitor,
double localSolveSeconds,
MPI_Comm communicator,
double denominatorFloor = 1.0e-300
);
struct ArnoldiOptions final {
int krylovDimension{40};
double breakdownRelativeTolerance{1.0e-13};
double ritzConvergenceRelativeTolerance{1.0e-8};
bool reorthogonalize{true};
};
struct RitzValueMeasurement final {
double realPart{0.0};
double imaginaryPart{0.0};
double magnitude{0.0};
double distanceFromOne{0.0};
double residualEstimate{0.0};
double relativeResidualEstimate{0.0};
bool converged{false};
};
enum class RitzValueOrdering { closest_to_zero, farthest_from_one, smallest_real_part, largest_magnitude };
struct ArnoldiSpectralMeasurement final {
int requestedDimension{0};
int achievedDimension{0};
bool invariantSubspaceFound{false};
std::uint64_t operatorApplications{0};
double operatorApplicationSecondsMaximumRank{0.0};
double operatorMaximumApplicationSecondsMaximumRank{0.0};
double measurementSecondsMaximumRank{0.0};
double nonApplicationSecondsMaximumRank{0.0};
int convergedRitzValueCount{0};
int negativeRealPartCount{0};
double projectedLargestSingularValue{0.0};
double projectedSmallestSingularValue{0.0};
double projectedConditionProxy{0.0};
double centroidRealPart{0.0};
double centroidImaginaryPart{0.0};
double rmsDistanceFromOne{0.0};
double rmsClusterRadius{0.0};
double minimumMagnitude{0.0};
double maximumMagnitude{0.0};
double minimumRealPart{0.0};
double maximumRealPart{0.0};
double maximumAbsoluteImaginaryPart{0.0};
double conjugatePairDefect{0.0};
double projectedDepartureFromNormality{0.0};
double projectedFieldOfValuesMinimumRealPart{0.0};
double projectedFieldOfValuesMaximumRealPart{0.0};
std::vector<RitzValueMeasurement> ritzValues;
};
[[nodiscard]] ArnoldiSpectralMeasurement measureArnoldiSpectrum(
const mfem::Operator &operation,
const mfem::Vector &initialDirection,
MPI_Comm communicator,
const ArnoldiOptions &options = {}
);
[[nodiscard]] std::vector<RitzValueMeasurement> selectRitzValues(
const ArnoldiSpectralMeasurement &measurement,
RitzValueOrdering ordering,
int count
);
} // namespace mean_field::solver