Skip to content

Scale corrected ir smoother #2036

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
greole wants to merge 7 commits into
base: develop
Choose a base branch
from
Draft

Scale corrected ir smoother #2036

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e94d3a2
wip scale corrected ir smoother
greole Jun 10, 2026
2bdf624
add missing header change
greole Jun 23, 2026
cb192b9
address review comments
greole Jun 26, 2026
1f306a6
update example
greole Jun 26, 2026
2b0b609
multigrid based scale correction
Jul 22, 2026
c9ff246
Remove abandoned IR scale-correction smoother
Jul 22, 2026
4251932
Fix multigrid scale correction for develop copy_from API
Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions benchmark/utils/preconditioners.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,18 @@ DEFINE_double(mg_tolerance, false, "The tolerance for the coarse solver");
DEFINE_uint32(mg_max_iters, false,
"The max number of iterations for the coarse solver");

DEFINE_string(mg_scale_correction, "none",
"OpenFOAM-style Rayleigh scale correction mode for the Multigrid "
"preconditioner: 'none' (off), 'post' (coarse-correction scaling "
"only), or 'both' (pre-smooth + coarse-correction scaling)");

DEFINE_uint32(mg_smoother_iters, 2,
"Number of block-Jacobi Richardson sweeps used for the Multigrid "
"pre/post smoother (the coarsest solver uses 4x this)");

DEFINE_double(mg_smoother_relax, 0.8,
"Relaxation factor of the Multigrid block-Jacobi smoother");


// parses the Jacobi storage optimization command line argument
gko::precision_reduction parse_storage_optimization(const std::string& flag)
Expand Down Expand Up @@ -315,6 +327,7 @@ const std::map<std::string, std::function<std::unique_ptr<gko::LinOpFactory>(
{"mg",
[](std::shared_ptr<const gko::Executor> exec) {
using ir = gko::solver::Ir<etype>;
using jacobi = gko::preconditioner::Jacobi<etype>;
auto iter_stop = gko::share(gko::stop::Iteration::build()
.with_max_iters(FLAGS_mg_max_iters)
.on(exec));
Expand All @@ -323,12 +336,45 @@ const std::map<std::string, std::function<std::unique_ptr<gko::LinOpFactory>(
.with_baseline(gko::stop::mode::absolute)
.with_reduction_factor(FLAGS_mg_tolerance)
.on(exec));
// damped block-Jacobi Richardson smoother (and a heavier variant
// as the coarsest solver) -- required for the multigrid V-cycle to
// be an effective preconditioner, and hence for scale correction
// to have a measurable effect.
auto smoother = gko::share(
ir::build()
.with_solver(jacobi::build().with_max_block_size(1u))
.with_relaxation_factor(
static_cast<etype>(FLAGS_mg_smoother_relax))
.with_criteria(gko::stop::Iteration::build().with_max_iters(
FLAGS_mg_smoother_iters))
.on(exec));
auto coarsest = gko::share(
ir::build()
.with_solver(jacobi::build().with_max_block_size(1u))
.with_relaxation_factor(
static_cast<etype>(FLAGS_mg_smoother_relax))
.with_criteria(gko::stop::Iteration::build().with_max_iters(
4u * FLAGS_mg_smoother_iters))
.on(exec));
const auto& sc_mode = FLAGS_mg_scale_correction;
if (sc_mode != "none" && sc_mode != "post" && sc_mode != "both") {
throw std::runtime_error(
"Unknown -mg_scale_correction mode '" + sc_mode +
"', expected 'none', 'post' or 'both'");
}
const bool scale_correction = sc_mode != "none";
const bool scale_correction_pre = sc_mode == "both";
return gko::solver::Multigrid::build()
.with_mg_level(
gko::multigrid::Pgm<etype, itype>::build()
.with_deterministic(FLAGS_pgm_deterministic))
.with_pre_smoother(smoother)
.with_post_smoother(smoother)
.with_coarsest_solver(coarsest)
.with_criteria(iter_stop, tol_stop)
.with_max_levels(FLAGS_mg_max_num_levels)
.with_scale_correction(scale_correction)
.with_scale_correction_pre_smooth(scale_correction_pre)
.on(exec);
}}
#if GINKGO_BUILD_MPI
Expand Down
170 changes: 153 additions & 17 deletions core/solver/multigrid.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include "ginkgo/core/solver/multigrid.hpp"

#include <complex>
#include <limits>
#include <string>

#include <ginkgo/core/base/exception.hpp>
Expand Down Expand Up @@ -291,6 +292,13 @@
std::vector<std::shared_ptr<const LinOp>> one_list;
std::vector<std::shared_ptr<const LinOp>> next_one_list;
std::vector<std::shared_ptr<const LinOp>> neg_one_list;
// scale correction workspace (only allocated when scale_correction=true)
std::vector<std::shared_ptr<LinOp>> acf_list; // A*δ scratch, nrows×nrhs
std::vector<std::shared_ptr<LinOp>>
delta_pre_list; // A*δ_c / smoother scratch
std::vector<std::shared_ptr<LinOp>> alpha_list; // Rayleigh scalar, 1×nrhs
std::vector<std::shared_ptr<LinOp>>
denom_list; // denominator scalar, 1×nrhs
const LinOp* system_matrix;
const Multigrid* multigrid;
size_type nrhs;
Expand All @@ -314,6 +322,12 @@
clear_and_reserve(one_list, list_size);
clear_and_reserve(next_one_list, list_size);
clear_and_reserve(neg_one_list, list_size);
if (multigrid_in->get_parameters().scale_correction) {
clear_and_reserve(acf_list, list_size);
clear_and_reserve(delta_pre_list, list_size);
clear_and_reserve(alpha_list, list_size);
clear_and_reserve(denom_list, list_size);
}
// Allocate memory first such that reusing allocation in each iter.
for (int i = 0; i < mg_level_list.size(); i++) {
auto next_nrows = mg_level_list.at(i)->get_coarse_op()->get_size()[0];
Expand Down Expand Up @@ -401,6 +415,13 @@
}
one_list.emplace_back(initialize<vec>({one<value_type>()}, exec));
neg_one_list.emplace_back(initialize<vec>({-one<value_type>()}, exec));
if (multigrid->get_parameters().scale_correction) {
acf_list.emplace_back(vec::create(exec, dim<2>{current_nrows, nrhs}));
delta_pre_list.emplace_back(
vec::create(exec, dim<2>{current_nrows, nrhs}));
alpha_list.emplace_back(vec::create(exec, dim<2>{1, nrhs}));
denom_list.emplace_back(vec::create(exec, dim<2>{1, nrhs}));
}
}


Expand Down Expand Up @@ -449,6 +470,16 @@
one_list.emplace_back(initialize<dense_vec>({one<value_type>()}, exec));
neg_one_list.emplace_back(
initialize<dense_vec>({-one<value_type>()}, exec));
if (multigrid->get_parameters().scale_correction) {
acf_list.emplace_back(vec::create(exec, current_comm,
dim<2>{current_nrows, nrhs},
dim<2>{current_local_nrows, nrhs}));
delta_pre_list.emplace_back(
vec::create(exec, current_comm, dim<2>{current_nrows, nrhs},
dim<2>{current_local_nrows, nrhs}));
alpha_list.emplace_back(dense_vec::create(exec, dim<2>{1, nrhs}));
denom_list.emplace_back(dense_vec::create(exec, dim<2>{1, nrhs}));
}
}


Expand Down Expand Up @@ -503,22 +534,47 @@
auto r = r_list.at(level);
auto g = g_list.at(level);
auto e = e_list.at(level);
// get mg_level
auto mg_level = multigrid->get_mg_level_list().at(level);
// get the pre_smoother
auto pre_smoother = multigrid->get_pre_smoother_list().at(level);
// get the mid_smoother
std::shared_ptr<const LinOp> mid_smoother{nullptr};
auto mid_case = multigrid->get_parameters().mid_case;
if (mid_case == multigrid::mid_smooth_type::standalone) {
mid_smoother = multigrid->get_mid_smoother_list().at(level);
}
// get the post_smoother
auto post_smoother = multigrid->get_post_smoother_list().at(level);
auto one = one_list.at(level).get();
auto next_one = next_one_list.at(level).get();
auto neg_one = neg_one_list.at(level).get();
// origin or next or first

// scale correction applies at all levels except immediately above coarsest
// (at that level the coarse solver provides a near-exact result, sf ≈ 1)
bool do_scale =
multigrid->get_parameters().scale_correction && level < total_level - 1;
// the pre-smooth (downward) scaling can be disabled independently, leaving
// only the post-smooth (coarse-correction) scaling ("post-only" mode)
bool do_pre_scale =
do_scale && multigrid->get_parameters().scale_correction_pre_smooth;

// [NeoN patch] device-side guarded Rayleigh reciprocal: sf = num / (denom + eps),
// computed entirely on-device to avoid the per-correction-point copy_val_to_host(denom)
// D2H sync (that guard fires a synchronizing device->host copy at every level every
// V-cycle -- ~90/solve on the occDrivAer pressure solve, a leading cost when the solve is

Check warning on line 561 in core/solver/multigrid.cpp

View workflow job for this annotation

GitHub Actions / Spell Check with Typos 

"Aer" should be "Are".
// synchronization-bound). Safe because num = delta.b and denom = delta.A.delta are BOTH
// exactly zero iff delta = 0, so at delta=0 sf = 0/eps = 0 (the correction reduces to
// smoother(b) -- benign, and does not occur after a pre-smoother sweep on nonzero b);
// eps = smallest positive normal is negligible vs any nonzero denom. Replaces the old
// `if (copy_val_to_host(denom) != 0) { inv_scale(...); <body> }`.
auto safe_inv_scale = [&](matrix::Dense<value_type>* alpha_d,
matrix::Dense<value_type>* denom_d) {
using real_type = gko::remove_complex<value_type>;
auto exec_l = multigrid->get_executor();
auto eps = matrix::Dense<value_type>::create(
exec_l, dim<2>{1, alpha_d->get_size()[1]});
eps->fill(value_type{std::numeric_limits<real_type>::min()});
denom_d->add_scaled(one, eps.get()); // denom += eps (device, no sync)
alpha_d->inv_scale(denom_d); // sf = num / (denom + eps)
};

bool use_pre = has_property(mode, cycle_mode::first_of_cycle) ||
mid_case == multigrid::mid_smooth_type::both ||
mid_case == multigrid::mid_smooth_type::pre_smoother;
Expand All @@ -540,15 +596,49 @@
pre_smoother->apply(b, x);
}
}

// Pre-smooth scale correction (OpenFOAM GAMGSolverSolve.C downward pass):
// Rayleigh-scale δ_pre = x, deflating r before restriction.
// Aδ = A * δ_pre
// sf = (δ_pre · b) / (δ_pre · Aδ)
// δ_pre = sf * δ_pre + smoother(b − sf * Aδ) [reuses acf, r scratch]
if (do_pre_scale && use_pre && pre_smoother) {
auto acf = acf_list.at(level);
auto dp = delta_pre_list.at(level);
auto alpha_dense = as<matrix::Dense<value_type>>(alpha_list.at(level));
auto denom_dense = as<matrix::Dense<value_type>>(denom_list.at(level));

matrix->apply(x, acf); // acf = A * δ_pre
as<VectorType>(x)->compute_dot(b, alpha_list.at(level));
as<VectorType>(x)->compute_dot(acf, denom_list.at(level));
// [NeoN patch] device-side guarded reciprocal (no copy_val_to_host D2H sync).
safe_inv_scale(alpha_dense.get(), denom_dense.get()); // sf = (δ·b)/(δ·Aδ)
{
// r temporarily holds r_scaled = b − sf * Aδ
as<VectorType>(acf)->scale(alpha_dense); // acf = sf * Aδ
as<Cloneable>(r)->copy_from(as<Cloneable>(b));
as<VectorType>(r)->add_scaled(neg_one, acf); // r = b − sf*Aδ

// dp = smoother(r_scaled) starting from zero
as<VectorType>(dp)->fill(zero<value_type>());
pre_smoother->apply(r, dp);

// x = sf * δ_pre + smoother(b − sf*Aδ)
as<VectorType>(x)->scale(alpha_dense);
as<VectorType>(x)->add_scaled(one, dp);
}
// r is overwritten with the actual (deflated) residual below
}

// The common smoother is wrapped by IR and IR already split the iter and
// residual check. Thus, when the IR only contains iter limit, there's no
// additional residual computation
// additional residual computation.
// TODO: if already computes the residual outside, the first level may not
// need this residual computation when no presmoother in the first level.
as<Cloneable>(r)->copy_from(as<Cloneable>(b)); // n * b
matrix->apply(neg_one, x, one, r);
as<Cloneable>(r)->copy_from(as<Cloneable>(b));
matrix->apply(neg_one, x, one, r); // r = b − A*x (deflated if scaled)

// first cycle
// restrict
mg_level->get_restrict_op()->apply(r, g);
// next level
if (level + 1 == total_level) {
Expand All @@ -568,9 +658,8 @@
next_mode);
if (level < multigrid->get_mg_level_list().size() - 1) {
// additional work for non-v_cycle
// next level
if (cycle == multigrid::cycle::f) {
// f_cycle call v_cycle in the second cycle
// f_cycle calls v_cycle in the second cycle
this->run_mg_cycle(multigrid::cycle::v, level + 1,
next_level_matrix, g.get(), e.get(),
cycle_mode::end_of_cycle);
Expand All @@ -579,20 +668,61 @@
e.get(), cycle_mode::end_of_cycle);
}
}
// prolong
mg_level->get_prolong_op()->apply(next_one, e, next_one, x);

// end or origin previous
// Post-smooth scale correction (OpenFOAM GAMGSolverSolve.C upward pass):
// Prolong coarse correction δ_c into acf, Rayleigh-scale it w.r.t. the
// deflated residual r, then merge with δ_pre (= current x).
// δ_c = prolong(e)
// Aδ = A * δ_c [stored in delta_pre scratch]
// sf = (δ_c · r) / (δ_c · Aδ)
// δ_c = sf * δ_c + smoother(r − sf * Aδ)
// x += δ_c [x = δ_pre + scale-corrected δ_c]
if (do_scale) {
auto acf = acf_list.at(level);
auto dp = delta_pre_list.at(level);
auto alpha_dense = as<matrix::Dense<value_type>>(alpha_list.at(level));
auto denom_dense = as<matrix::Dense<value_type>>(denom_list.at(level));

// prolong e into acf (δ_c = prolong(e))
as<VectorType>(acf)->fill(zero<value_type>());
mg_level->get_prolong_op()->apply(next_one, e, next_one, acf);

matrix->apply(acf, dp); // dp = A * δ_c
as<VectorType>(acf)->compute_dot(r, alpha_list.at(level));
as<VectorType>(acf)->compute_dot(dp, denom_list.at(level));
// [NeoN patch] device-side guarded reciprocal (no copy_val_to_host D2H sync).
safe_inv_scale(alpha_dense.get(),
denom_dense.get()); // sf = (δ_c·r)/(δ_c·Aδ_c)
{
// r temporarily holds r_scaled = r − sf * Aδ_c
as<VectorType>(dp)->scale(alpha_dense); // dp = sf * Aδ_c
as<VectorType>(r)->add_scaled(neg_one, dp); // r = r − sf*Aδ_c

// dp = smoother(r_scaled) starting from zero
as<VectorType>(dp)->fill(zero<value_type>());
if (pre_smoother) {
pre_smoother->apply(r, dp);
}

// acf = sf * δ_c + smoother(r − sf*Aδ_c)
as<VectorType>(acf)->scale(alpha_dense);
as<VectorType>(acf)->add_scaled(one, dp);
}
// x = δ_pre + scale-corrected δ_c
as<VectorType>(x)->add_scaled(one, acf);
} else {
// standard prolongation: x += prolong(e)
mg_level->get_prolong_op()->apply(next_one, e, next_one, x);
}

bool use_post = has_property(mode, cycle_mode::end_of_cycle) ||
mid_case == multigrid::mid_smooth_type::both ||
mid_case == multigrid::mid_smooth_type::post_smoother;
// post-smooth
if (use_post && post_smoother) {
post_smoother->apply(b, x);
}

// put the mid smoother into the end of previous cycle
// only W/F cycle
// put the mid smoother into the end of previous cycle (W/F cycle only)
bool use_mid =
(cycle == multigrid::cycle::w || cycle == multigrid::cycle::f) &&
!has_property(mode, cycle_mode::end_of_cycle) &&
Expand Down Expand Up @@ -695,6 +825,12 @@
params.with_default_initial_guess(
config::get_value<solver::initial_guess_mode>(obj));
}
if (auto& obj = config_check.get("scale_correction")) {
params.with_scale_correction(config::get_value<bool>(obj));
}
if (auto& obj = config_check.get("scale_correction_pre_smooth")) {
params.with_scale_correction_pre_smooth(config::get_value<bool>(obj));
}

return params;
}
Expand Down
Loading
Loading