Probe FEM gravity strategies and make reduced contracts collective

This commit is contained in:
2026-09-17 09:30:12 -04:00
parent 9e0e1bc42a
commit a7f080f265
2 changed files with 361 additions and 6 deletions

View File

@@ -128,6 +128,7 @@ export namespace mean_field::solver {
mfem::Vector &output
) const {
RequireCurrent();
RequireInputSize(input, Width());
const auto &partition = GetPartition();
partition.InjectRetainedValues(input, m_fullDirection);
FullMult(m_fullDirection, m_fullAction);
@@ -165,6 +166,7 @@ export namespace mean_field::solver {
) const {
RequireCurrent();
const auto &partition = GetPartition();
RequireInputSize(fullRightHandSide, partition.RootResidualSize());
partition.GatherGravityResiduals(fullRightHandSide, m_gravityRightHandSide);
if (!SolveGravity()) {
return m_lastInnerSolve;
@@ -185,12 +187,11 @@ export namespace mean_field::solver {
) const {
RequireCurrent();
const auto &partition = GetPartition();
RequireInputSize(retainedCorrection, partition.RetainedValueSize());
RequireInputSize(fullRightHandSide, partition.RootResidualSize());
partition.InjectRetainedValues(retainedCorrection, m_fullDirection);
FullMult(m_fullDirection, m_fullAction);
m_fullAction *= -1.0;
if (fullRightHandSide.Size() != m_fullAction.Size()) {
throw std::invalid_argument("A lifted gravity correction requires a full residual right-hand side.");
}
m_fullAction += fullRightHandSide;
partition.GatherGravityResiduals(m_fullAction, m_gravityRightHandSide);
return FinishLift(fullCorrection);
@@ -202,6 +203,7 @@ export namespace mean_field::solver {
) const {
RequireCurrent();
const auto &partition = GetPartition();
RequireInputSize(retainedDirection, partition.RetainedValueSize());
partition.InjectRetainedValues(retainedDirection, m_fullDirection);
FullMult(m_fullDirection, m_fullAction);
partition.GatherGravityResiduals(m_fullAction, m_gravityRightHandSide);
@@ -210,12 +212,29 @@ export namespace mean_field::solver {
}
private:
[[nodiscard]] bool AllRanksSatisfy(const bool valid) const {
const int local = valid ? 1 : 0;
int global = 0;
if (MPI_Allreduce(&local, &global, 1, MPI_INT, MPI_MIN, GetCommunicator()) != MPI_SUCCESS) {
throw std::runtime_error("Reduced gravity could not verify its distributed application contract.");
}
return global != 0;
}
void RequireCurrent() const {
if (!IsCurrent()) {
// An early rank-local exception would strand the remaining ranks
// inside the next nested gravity solve.
if (!AllRanksSatisfy(IsCurrent())) {
throw std::logic_error("A reduced-gravity action requires the current prepared gravity subsystem.");
}
}
void RequireInputSize(const mfem::Vector &input, const int size) const {
if (!AllRanksSatisfy(input.Size() == size)) {
throw std::invalid_argument("A reduced-gravity action requires compatible input vector sizes.");
}
}
void FullMult(
const mfem::Vector &input,
mfem::Vector &output
@@ -302,13 +321,16 @@ export namespace mean_field::solver {
const mfem::Vector &input,
mfem::Vector &output
) const override {
if (!m_operation->IsCurrent()) {
if (!AllRanksSatisfy(m_operation->IsCurrent())) {
throw std::logic_error("The reduced-gravity preconditioner requires the current prepared subsystem.");
}
if (!AllRanksSatisfy(input.Size() == Width() && output.Size() == Height())) {
throw std::invalid_argument("The reduced-gravity inverse requires preallocated compatible vectors.");
}
const auto &partition = m_operation->GetPartition();
partition.InjectRetainedResiduals(input, m_fullRightHandSide);
m_fullInverse->Mult(m_fullRightHandSide, m_fullCorrection);
if (!m_operation->IsCurrent()) {
if (!AllRanksSatisfy(m_operation->IsCurrent())) {
throw std::logic_error("The full inverse changed the reduced-gravity preparation.");
}
partition.GatherRetainedValues(m_fullCorrection, output);
@@ -319,6 +341,15 @@ export namespace mean_field::solver {
}
private:
[[nodiscard]] bool AllRanksSatisfy(const bool valid) const {
const int local = valid ? 1 : 0;
int global = 0;
if (MPI_Allreduce(&local, &global, 1, MPI_INT, MPI_MIN, m_operation->GetCommunicator()) != MPI_SUCCESS) {
throw std::runtime_error("The reduced-gravity inverse could not verify distributed validity.");
}
return global != 0;
}
const GravityReducedOperator<Subsystem> *m_operation;
FullInverse *m_fullInverse;
mutable mfem::Vector m_fullRightHandSide;