From a162d448f02af5e8816bd8c6a71de6a126ef04a5 Mon Sep 17 00:00:00 2001 From: Cody Balos Date: Wed, 24 Jun 2026 15:46:10 -0700 Subject: [PATCH 1/3] add set_preconditioner and set_fast_preconditioner --- .../source/RuntimeParameters.rst | 6 +- .../source/TimeIntegration_Chapter.rst | 7 + Src/Base/AMReX_IntegratorBase.H | 50 ++++ Src/Base/AMReX_TimeIntegrator.H | 16 ++ .../SUNDIALS/AMReX_SundialsIntegrator.H | 266 +++++++++++++++++- 5 files changed, 332 insertions(+), 13 deletions(-) diff --git a/Docs/sphinx_documentation/source/RuntimeParameters.rst b/Docs/sphinx_documentation/source/RuntimeParameters.rst index 9642de257b2..5c7c67e0dce 100644 --- a/Docs/sphinx_documentation/source/RuntimeParameters.rst +++ b/Docs/sphinx_documentation/source/RuntimeParameters.rst @@ -1502,6 +1502,8 @@ with implicit methods (e.g., DIRK). The linear solver used with Newton's method for single rate methods (e.g., DIRK) or at the slow time scale with multirate methods (e.g., IM-MRI). + AMReX currently provides GMRES here, with optional preconditioning supplied + through :cpp:`TimeIntegrator::set_preconditioner()`. .. py:data:: integration.sundials.max_linear_iters :type: int @@ -1533,7 +1535,9 @@ with implicit methods (e.g., DIRK). :value: GMRES The linear solver used with Newton's method at the fast time scale with - multirate methods (e.g., when the fast method is DIRK). + multirate methods (e.g., when the fast method is DIRK). Optional fast + preconditioning can be attached through + :cpp:`TimeIntegrator::set_fast_preconditioner()`. .. py:data:: integration.sundials.fast_max_linear_iters :type: int diff --git a/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst b/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst index 722fdbe973a..8ba19cafc96 100644 --- a/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst +++ b/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst @@ -173,6 +173,13 @@ methods, one also needs to select the fast time scale method type using the input option :py:data:`integration.sundials.fast_type`, which may be set to ``ERK`` or ``DIRK``. +Implicit SUNDIALS solves can also attach preconditioner callbacks through +``TimeIntegrator::set_preconditioner()`` and, for fast implicit MRI solves, +``TimeIntegrator::set_fast_preconditioner()``. These callbacks receive AMReX +data structures unpacked from the SUNDIALS vectors, which makes it possible to +reuse AMReX linear solver infrastructure such as hypre-backed solves inside the +SUNDIALS Newton iteration. + To select a specific SUNDIALS method, use the input option :py:data:`integration.sundials.method` for ERK and DIRK methods as well as the slow time scale method with an MRI integrator. Use diff --git a/Src/Base/AMReX_IntegratorBase.H b/Src/Base/AMReX_IntegratorBase.H index a6ba598aade..d9f35573234 100644 --- a/Src/Base/AMReX_IntegratorBase.H +++ b/Src/Base/AMReX_IntegratorBase.H @@ -179,6 +179,18 @@ private: pp.query("fast_abs_tol", fast_abs_tol); } protected: + using PrecondSetupFunction = std::function; + + using PrecondSolveFunction = std::function; + /** * \brief Rhs is the right-hand-side function the integrator will use. */ @@ -226,6 +238,30 @@ protected: */ std::function post_fast_step_action; + /** + * \brief Optional SUNDIALS-compatible setup callback for the slow or + * single-rate implicit preconditioner. + */ + PrecondSetupFunction PrecondSetup; + + /** + * \brief Optional SUNDIALS-compatible solve callback for the slow or + * single-rate implicit preconditioner. + */ + PrecondSolveFunction PrecondSolve; + + /** + * \brief Optional SUNDIALS-compatible setup callback for the fast implicit + * preconditioner in multirate methods. + */ + PrecondSetupFunction FastPrecondSetup; + + /** + * \brief Optional SUNDIALS-compatible solve callback for the fast implicit + * preconditioner in multirate methods. + */ + PrecondSolveFunction FastPrecondSolve; + /** * \brief Flag to enable/disable adaptive time stepping in single rate * methods or at the slow time scale in multirate methods (bool) @@ -315,6 +351,20 @@ public: RhsFast = F; } + virtual void set_preconditioner (PrecondSetupFunction Psetup, + PrecondSolveFunction Psolve) + { + PrecondSetup = std::move(Psetup); + PrecondSolve = std::move(Psolve); + } + + virtual void set_fast_preconditioner (PrecondSetupFunction Psetup, + PrecondSolveFunction Psolve) + { + FastPrecondSetup = std::move(Psetup); + FastPrecondSolve = std::move(Psolve); + } + void set_post_stage_action (std::function A) { post_stage_action = A; diff --git a/Src/Base/AMReX_TimeIntegrator.H b/Src/Base/AMReX_TimeIntegrator.H index fff178bae3c..fe67ecdcedf 100644 --- a/Src/Base/AMReX_TimeIntegrator.H +++ b/Src/Base/AMReX_TimeIntegrator.H @@ -147,6 +147,22 @@ public: integrator_ptr->set_fast_rhs(F); } + void set_preconditioner ( + std::function Psetup, + std::function Psolve) + { + integrator_ptr->set_preconditioner(std::move(Psetup), std::move(Psolve)); + } + + void set_fast_preconditioner ( + std::function Psetup, + std::function Psolve) + { + integrator_ptr->set_fast_preconditioner(std::move(Psetup), std::move(Psolve)); + } + void set_post_stage_action (std::function A) { integrator_ptr->set_post_stage_action(A); diff --git a/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H b/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H index aae204ed4cb..f3d5b07afde 100644 --- a/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H +++ b/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H @@ -37,6 +37,14 @@ struct SundialsUserData { std::function fi; //!< Implicit RHS for ImEx schemes. std::function fe; //!< Explicit RHS for ImEx schemes. std::function ff; //!< MRI fast-scale RHS. + std::function psetup; //!< Slow/single-rate preconditioner setup. + std::function psolve; //!< Slow/single-rate preconditioner solve. + std::function fast_psetup; //!< Fast preconditioner setup. + std::function fast_psolve; //!< Fast preconditioner solve. std::function post_stage; //!< Hook invoked after each stage. std::function post_step; //!< Hook invoked after each time step. std::function post_fast_stage; //!< Hook for MRI fast stages. @@ -64,6 +72,36 @@ namespace SundialsUserFun { return udata->ff(t, y_data, y_rhs, user_data); } + static int psetup (amrex::Real t, N_Vector y_data, N_Vector y_rhs, + sunbooleantype jok, sunbooleantype* jcurPtr, + amrex::Real gamma, void *user_data) { + SundialsUserData* udata = static_cast(user_data); + return udata->psetup(t, y_data, y_rhs, jok, jcurPtr, gamma, user_data); + } + + static int psolve (amrex::Real t, N_Vector y_data, N_Vector y_rhs, + N_Vector r_data, N_Vector z_data, amrex::Real gamma, + amrex::Real delta, int lr, void *user_data) { + SundialsUserData* udata = static_cast(user_data); + return udata->psolve(t, y_data, y_rhs, r_data, z_data, gamma, + delta, lr, user_data); + } + + static int fast_psetup (amrex::Real t, N_Vector y_data, N_Vector y_rhs, + sunbooleantype jok, sunbooleantype* jcurPtr, + amrex::Real gamma, void *user_data) { + SundialsUserData* udata = static_cast(user_data); + return udata->fast_psetup(t, y_data, y_rhs, jok, jcurPtr, gamma, user_data); + } + + static int fast_psolve (amrex::Real t, N_Vector y_data, N_Vector y_rhs, + N_Vector r_data, N_Vector z_data, amrex::Real gamma, + amrex::Real delta, int lr, void *user_data) { + SundialsUserData* udata = static_cast(user_data); + return udata->fast_psolve(t, y_data, y_rhs, r_data, z_data, gamma, + delta, lr, user_data); + } + static int post_stage (amrex::Real t, N_Vector y_data, void *user_data) { SundialsUserData* udata = static_cast(user_data); return udata->post_stage(t, y_data, user_data); @@ -139,12 +177,14 @@ private: void *arkode_mem = nullptr; SUNLinearSolver LS = nullptr; SUNNonlinearSolver NLS = nullptr; + bool slow_ls_needs_update = false; // Fast time scale void *arkode_fast_mem = nullptr; MRIStepInnerStepper fast_stepper = nullptr; SUNLinearSolver fast_LS = nullptr; SUNNonlinearSolver fast_NLS = nullptr; + bool fast_ls_needs_update = false; // Integrator stop time bool set_stop_time = false; @@ -153,6 +193,90 @@ private: // Max steps between returns amrex::Long max_num_steps = 0; + void ConfigureSlowLinearSolver (N_Vector y_data) + { + if ((type != "DIRK" && type != "IMEX-RK" && + type != "IM-MRI" && type != "IMEX-MRI") || + nonlinear_solver != "Newton") + { + return; + } + + const bool has_preconditioner = static_cast(BaseT::PrecondSolve); + if (BaseT::PrecondSetup && !has_preconditioner) { + amrex::Error("SUNDIALS preconditioner setup callback requires a solve callback."); + } + + if (!slow_ls_needs_update && LS != nullptr) { + return; + } + + SUNLinSolFree(LS); + LS = SUNLinSol_SPGMR(y_data, + has_preconditioner ? SUN_PREC_LEFT : SUN_PREC_NONE, + max_linear_iters, sunctx); + AMREX_ALWAYS_ASSERT(LS != nullptr); + + int flag = 0; + if (use_ark) { + flag = ARKStepSetLinearSolver(arkode_mem, LS, nullptr); + AMREX_ALWAYS_ASSERT(flag == 0); + if (has_preconditioner) { + flag = ARKStepSetPreconditioner( + arkode_mem, + BaseT::PrecondSetup ? SundialsUserFun::psetup : nullptr, + SundialsUserFun::psolve); + AMREX_ALWAYS_ASSERT(flag == 0); + } + } else if (use_mri) { + flag = MRIStepSetLinearSolver(arkode_mem, LS, nullptr); + AMREX_ALWAYS_ASSERT(flag == 0); + if (has_preconditioner) { + flag = MRIStepSetPreconditioner( + arkode_mem, + BaseT::PrecondSetup ? SundialsUserFun::psetup : nullptr, + SundialsUserFun::psolve); + AMREX_ALWAYS_ASSERT(flag == 0); + } + } + + slow_ls_needs_update = false; + } + + void ConfigureFastLinearSolver (N_Vector y_data) + { + if (fast_type != "DIRK" || fast_nonlinear_solver != "Newton") { + return; + } + + const bool has_preconditioner = static_cast(BaseT::FastPrecondSolve); + if (BaseT::FastPrecondSetup && !has_preconditioner) { + amrex::Error("SUNDIALS fast preconditioner setup callback requires a solve callback."); + } + + if (!fast_ls_needs_update && fast_LS != nullptr) { + return; + } + + SUNLinSolFree(fast_LS); + fast_LS = SUNLinSol_SPGMR(y_data, + has_preconditioner ? SUN_PREC_LEFT : SUN_PREC_NONE, + fast_max_linear_iters, sunctx); + AMREX_ALWAYS_ASSERT(fast_LS != nullptr); + + int flag = ARKStepSetLinearSolver(arkode_fast_mem, fast_LS, nullptr); + AMREX_ALWAYS_ASSERT(flag == 0); + if (has_preconditioner) { + flag = ARKStepSetPreconditioner( + arkode_fast_mem, + BaseT::FastPrecondSetup ? SundialsUserFun::fast_psetup : nullptr, + SundialsUserFun::fast_psolve); + AMREX_ALWAYS_ASSERT(flag == 0); + } + + fast_ls_needs_update = false; + } + void initialize_parameters () { amrex::ParmParse pp("integration.sundials"); @@ -264,10 +388,8 @@ private: amrex::Print() << "Linear solver: " << linear_solver << "\n"; amrex::Print() << "Max linear iters: " << max_linear_iters << "\n"; } - LS = SUNLinSol_SPGMR(y_data, SUN_PREC_NONE, max_linear_iters, sunctx); - AMREX_ALWAYS_ASSERT(LS != nullptr); - flag = ARKStepSetLinearSolver(arkode_mem, LS, nullptr); - AMREX_ALWAYS_ASSERT(flag == 0); + slow_ls_needs_update = true; + ConfigureSlowLinearSolver(y_data); } } @@ -331,16 +453,14 @@ private: amrex::Print() << "Linear solver: " << fast_linear_solver << "\n"; amrex::Print() << "Max linear iters: " << fast_max_linear_iters << "\n"; } - fast_LS = SUNLinSol_SPGMR(y_data, SUN_PREC_NONE, fast_max_linear_iters, sunctx); - AMREX_ALWAYS_ASSERT(fast_LS != nullptr); - flag = ARKStepSetLinearSolver(arkode_fast_mem, fast_LS, nullptr); - AMREX_ALWAYS_ASSERT(flag == 0); + fast_ls_needs_update = true; } } // Attach structure with user-supplied function wrappers flag = ARKStepSetUserData(arkode_fast_mem, &udata); AMREX_ALWAYS_ASSERT(flag == 0); + ConfigureFastLinearSolver(y_data); // Set integrator tolerances if (BaseT::use_adaptive_fast_time_step || fast_type == "DIRK" || fast_type == "IMEX-RK") { @@ -429,10 +549,7 @@ private: amrex::Print() << "Linear solver: " << linear_solver << "\n"; amrex::Print() << "Max linear iters: " << max_linear_iters << "\n"; } - LS = SUNLinSol_SPGMR(y_data, SUN_PREC_NONE, max_linear_iters, sunctx); - AMREX_ALWAYS_ASSERT(LS != nullptr); - flag = MRIStepSetLinearSolver(arkode_mem, LS, nullptr); - AMREX_ALWAYS_ASSERT(flag == 0); + slow_ls_needs_update = true; } } @@ -452,6 +569,7 @@ private: // Set max number of steps between returns flag = MRIStepSetMaxNumSteps(arkode_mem, max_num_steps); AMREX_ALWAYS_ASSERT(flag == 0); + ConfigureSlowLinearSolver(y_data); } // ------------------------------------- @@ -668,6 +786,104 @@ public: return 0; }; + udata.psetup = [&](amrex::Real rhs_time, N_Vector y_data, N_Vector y_rhs, + sunbooleantype jok, sunbooleantype* jcurPtr, + amrex::Real gamma, void * /* user_data */) -> int { + + T S_data; + unpack_vector(y_data, S_data); + + T S_rhs; + unpack_vector(y_rhs, S_rhs); + + bool jcur = (jcurPtr != nullptr) && (*jcurPtr != SUNFALSE); + if (!BaseT::PrecondSetup) { + if (jcurPtr != nullptr) { + *jcurPtr = jcur ? SUNTRUE : SUNFALSE; + } + return 0; + } + + int ierr = BaseT::PrecondSetup(S_data, S_rhs, rhs_time, + jok != SUNFALSE, jcur, gamma); + if (jcurPtr != nullptr) { + *jcurPtr = jcur ? SUNTRUE : SUNFALSE; + } + return ierr; + }; + + udata.psolve = [&](amrex::Real rhs_time, N_Vector y_data, N_Vector y_rhs, + N_Vector r_data, N_Vector z_data, amrex::Real gamma, + amrex::Real delta, int lr, + void * /* user_data */) -> int { + + T S_data; + unpack_vector(y_data, S_data); + + T S_rhs; + unpack_vector(y_rhs, S_rhs); + + T S_r; + unpack_vector(r_data, S_r); + + T S_z; + unpack_vector(z_data, S_z); + + AMREX_ALWAYS_ASSERT_WITH_MESSAGE(static_cast(BaseT::PrecondSolve), + "SUNDIALS preconditioner solve callback not set."); + return BaseT::PrecondSolve(S_z, S_r, S_data, S_rhs, + rhs_time, gamma, delta, lr); + }; + + udata.fast_psetup = [&](amrex::Real rhs_time, N_Vector y_data, N_Vector y_rhs, + sunbooleantype jok, sunbooleantype* jcurPtr, + amrex::Real gamma, void * /* user_data */) -> int { + + T S_data; + unpack_vector(y_data, S_data); + + T S_rhs; + unpack_vector(y_rhs, S_rhs); + + bool jcur = (jcurPtr != nullptr) && (*jcurPtr != SUNFALSE); + if (!BaseT::FastPrecondSetup) { + if (jcurPtr != nullptr) { + *jcurPtr = jcur ? SUNTRUE : SUNFALSE; + } + return 0; + } + + int ierr = BaseT::FastPrecondSetup(S_data, S_rhs, rhs_time, + jok != SUNFALSE, jcur, gamma); + if (jcurPtr != nullptr) { + *jcurPtr = jcur ? SUNTRUE : SUNFALSE; + } + return ierr; + }; + + udata.fast_psolve = [&](amrex::Real rhs_time, N_Vector y_data, N_Vector y_rhs, + N_Vector r_data, N_Vector z_data, amrex::Real gamma, + amrex::Real delta, int lr, + void * /* user_data */) -> int { + + T S_data; + unpack_vector(y_data, S_data); + + T S_rhs; + unpack_vector(y_rhs, S_rhs); + + T S_r; + unpack_vector(r_data, S_r); + + T S_z; + unpack_vector(z_data, S_z); + + AMREX_ALWAYS_ASSERT_WITH_MESSAGE(static_cast(BaseT::FastPrecondSolve), + "SUNDIALS fast preconditioner solve callback not set."); + return BaseT::FastPrecondSolve(S_z, S_r, S_data, S_rhs, + rhs_time, gamma, delta, lr); + }; + udata.post_stage = [&](amrex::Real time, N_Vector y_data, void * /* user_data */) -> int { @@ -762,6 +978,20 @@ public: SUNNonlinSolFree(fast_NLS); } + void set_preconditioner (typename BaseT::PrecondSetupFunction Psetup, + typename BaseT::PrecondSolveFunction Psolve) override + { + BaseT::set_preconditioner(std::move(Psetup), std::move(Psolve)); + slow_ls_needs_update = true; + } + + void set_fast_preconditioner (typename BaseT::PrecondSetupFunction Psetup, + typename BaseT::PrecondSolveFunction Psolve) override + { + BaseT::set_fast_preconditioner(std::move(Psetup), std::move(Psolve)); + fast_ls_needs_update = true; + } + /** * \brief Take a single time step of size \p dt starting from \p S_old. * @@ -778,6 +1008,12 @@ public: N_Vector y_old = wrap_data(S_old); N_Vector y_new = wrap_data(S_new); + if (use_ark) { + ConfigureSlowLinearSolver(y_old); + } else if (use_mri) { + ConfigureFastLinearSolver(y_old); + ConfigureSlowLinearSolver(y_old); + } if (use_ark) { ARKStepReset(arkode_mem, time, y_old); // should probably resize @@ -812,6 +1048,12 @@ public: amrex::Real time_ret; // SUNDIALS return time N_Vector y_out = wrap_data(S_out); + if (use_ark) { + ConfigureSlowLinearSolver(y_out); + } else if (use_mri) { + ConfigureFastLinearSolver(y_out); + ConfigureSlowLinearSolver(y_out); + } if (use_ark) { if (!BaseT::use_adaptive_time_step) { From 20bb072fe1079af9c6f3452568b2700d56f6f7b0 Mon Sep 17 00:00:00 2001 From: Cody Balos Date: Wed, 24 Jun 2026 16:12:14 -0700 Subject: [PATCH 2/3] make precond routines return void --- Src/Base/AMReX_IntegratorBase.H | 22 ++++++------ .../SUNDIALS/AMReX_SundialsIntegrator.H | 36 +++++-------------- 2 files changed, 20 insertions(+), 38 deletions(-) diff --git a/Src/Base/AMReX_IntegratorBase.H b/Src/Base/AMReX_IntegratorBase.H index d9f35573234..2df06da8782 100644 --- a/Src/Base/AMReX_IntegratorBase.H +++ b/Src/Base/AMReX_IntegratorBase.H @@ -179,17 +179,17 @@ private: pp.query("fast_abs_tol", fast_abs_tol); } protected: - using PrecondSetupFunction = std::function; - - using PrecondSolveFunction = std::function; + using PrecondSetupFunction = std::function; + + using PrecondSolveFunction = std::function; /** * \brief Rhs is the right-hand-side function the integrator will use. diff --git a/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H b/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H index f3d5b07afde..08aa4d78f64 100644 --- a/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H +++ b/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H @@ -796,20 +796,11 @@ public: T S_rhs; unpack_vector(y_rhs, S_rhs); - bool jcur = (jcurPtr != nullptr) && (*jcurPtr != SUNFALSE); - if (!BaseT::PrecondSetup) { - if (jcurPtr != nullptr) { - *jcurPtr = jcur ? SUNTRUE : SUNFALSE; - } - return 0; - } - - int ierr = BaseT::PrecondSetup(S_data, S_rhs, rhs_time, - jok != SUNFALSE, jcur, gamma); - if (jcurPtr != nullptr) { - *jcurPtr = jcur ? SUNTRUE : SUNFALSE; - } - return ierr; + bool jcur = false; + BaseT::PrecondSetup(S_data, S_rhs, rhs_time, jok, jcur, gamma); + *jcurPtr = jcur ? SUNTRUE : SUNFALSE; + + return 0; }; udata.psolve = [&](amrex::Real rhs_time, N_Vector y_data, N_Vector y_rhs, @@ -845,19 +836,10 @@ public: T S_rhs; unpack_vector(y_rhs, S_rhs); - bool jcur = (jcurPtr != nullptr) && (*jcurPtr != SUNFALSE); - if (!BaseT::FastPrecondSetup) { - if (jcurPtr != nullptr) { - *jcurPtr = jcur ? SUNTRUE : SUNFALSE; - } - return 0; - } - - int ierr = BaseT::FastPrecondSetup(S_data, S_rhs, rhs_time, - jok != SUNFALSE, jcur, gamma); - if (jcurPtr != nullptr) { - *jcurPtr = jcur ? SUNTRUE : SUNFALSE; - } + bool jcur = false; + BaseT::FastPrecondSetup(S_data, S_rhs, rhs_time, jcur, gamma); + *jcurPtr = jcur ? SUNTRUE : SUNFALSE; + return ierr; }; From 7823a8b1ee0797ef40b74e2144772892b15286e4 Mon Sep 17 00:00:00 2001 From: Cody Balos Date: Wed, 24 Jun 2026 16:37:37 -0700 Subject: [PATCH 3/3] make preconditioner an input option --- .../source/RuntimeParameters.rst | 26 ++++++ .../source/TimeIntegration_Chapter.rst | 4 +- Src/Base/AMReX_TimeIntegrator.H | 12 +-- .../SUNDIALS/AMReX_SundialsIntegrator.H | 86 +++++++++++++++---- 4 files changed, 104 insertions(+), 24 deletions(-) diff --git a/Docs/sphinx_documentation/source/RuntimeParameters.rst b/Docs/sphinx_documentation/source/RuntimeParameters.rst index 5c7c67e0dce..72839c2a2ec 100644 --- a/Docs/sphinx_documentation/source/RuntimeParameters.rst +++ b/Docs/sphinx_documentation/source/RuntimeParameters.rst @@ -1505,6 +1505,19 @@ with implicit methods (e.g., DIRK). AMReX currently provides GMRES here, with optional preconditioning supplied through :cpp:`TimeIntegrator::set_preconditioner()`. +.. py:data:: integration.sundials.linear_solver_preconditioning + :type: string + :value: LEFT + + The SPGMR preconditioning type used with the slow or single-rate linear + solver when a preconditioner is supplied through + :cpp:`TimeIntegrator::set_preconditioner()`. The supported values are: + + * LEFT + * RIGHT + * BOTH + * NONE + .. py:data:: integration.sundials.max_linear_iters :type: int :value: 5 @@ -1539,6 +1552,19 @@ with implicit methods (e.g., DIRK). preconditioning can be attached through :cpp:`TimeIntegrator::set_fast_preconditioner()`. +.. py:data:: integration.sundials.fast_linear_solver_preconditioning + :type: string + :value: LEFT + + The SPGMR preconditioning type used with the fast linear solver when a fast + preconditioner is supplied through + :cpp:`TimeIntegrator::set_fast_preconditioner()`. The supported values are: + + * LEFT + * RIGHT + * BOTH + * NONE + .. py:data:: integration.sundials.fast_max_linear_iters :type: int :value: 5 diff --git a/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst b/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst index 8ba19cafc96..d0fe5b014e8 100644 --- a/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst +++ b/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst @@ -178,7 +178,9 @@ Implicit SUNDIALS solves can also attach preconditioner callbacks through ``TimeIntegrator::set_fast_preconditioner()``. These callbacks receive AMReX data structures unpacked from the SUNDIALS vectors, which makes it possible to reuse AMReX linear solver infrastructure such as hypre-backed solves inside the -SUNDIALS Newton iteration. +SUNDIALS Newton iteration. The SPGMR preconditioning side can be selected with +:py:data:`integration.sundials.linear_solver_preconditioning` and +:py:data:`integration.sundials.fast_linear_solver_preconditioning`. To select a specific SUNDIALS method, use the input option :py:data:`integration.sundials.method` for ERK and DIRK methods as well as the diff --git a/Src/Base/AMReX_TimeIntegrator.H b/Src/Base/AMReX_TimeIntegrator.H index fe67ecdcedf..d634bf8a481 100644 --- a/Src/Base/AMReX_TimeIntegrator.H +++ b/Src/Base/AMReX_TimeIntegrator.H @@ -148,17 +148,17 @@ public: } void set_preconditioner ( - std::function Psetup, - std::function Psolve) + std::function Psetup, + std::function Psolve) { integrator_ptr->set_preconditioner(std::move(Psetup), std::move(Psolve)); } void set_fast_preconditioner ( - std::function Psetup, - std::function Psolve) + std::function Psetup, + std::function Psolve) { integrator_ptr->set_fast_preconditioner(std::move(Psetup), std::move(Psolve)); } diff --git a/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H b/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H index 08aa4d78f64..2f8962639b6 100644 --- a/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H +++ b/Src/Extern/SUNDIALS/AMReX_SundialsIntegrator.H @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -155,9 +156,13 @@ private: // Linear solver std::string linear_solver = "GMRES"; + std::string linear_solver_preconditioning = "LEFT"; + int linear_solver_preconditioning_type = SUN_PREC_LEFT; int max_linear_iters = 0; std::string fast_linear_solver = "GMRES"; + std::string fast_linear_solver_preconditioning = "LEFT"; + int fast_linear_solver_preconditioning_type = SUN_PREC_LEFT; int fast_max_linear_iters = 0; // SUNDIALS package flags, set based on type @@ -193,6 +198,32 @@ private: // Max steps between returns amrex::Long max_num_steps = 0; + static int PreconditioningType (const std::string& preconditioning, const char* parameter_name) + { + const std::string value = amrex::toUpper(preconditioning); + + if (value == "LEFT") { + return SUN_PREC_LEFT; + } + if (value == "RIGHT") { + return SUN_PREC_RIGHT; + } + if (value == "BOTH") { + return SUN_PREC_BOTH; + } + if (value == "NONE") { + return SUN_PREC_NONE; + } + + std::string msg("Unknown integration.sundials."); + msg += parameter_name; + msg += ": "; + msg += preconditioning; + msg += ". Supported values are LEFT, RIGHT, BOTH, and NONE."; + amrex::Error(msg.c_str()); + return SUN_PREC_NONE; + } + void ConfigureSlowLinearSolver (N_Vector y_data) { if ((type != "DIRK" && type != "IMEX-RK" && @@ -202,10 +233,13 @@ private: return; } - const bool has_preconditioner = static_cast(BaseT::PrecondSolve); - if (BaseT::PrecondSetup && !has_preconditioner) { + if (BaseT::PrecondSetup && !BaseT::PrecondSolve) { amrex::Error("SUNDIALS preconditioner setup callback requires a solve callback."); } + const int preconditioning = BaseT::PrecondSolve + ? linear_solver_preconditioning_type + : SUN_PREC_NONE; + const bool use_preconditioner = preconditioning != SUN_PREC_NONE; if (!slow_ls_needs_update && LS != nullptr) { return; @@ -213,7 +247,7 @@ private: SUNLinSolFree(LS); LS = SUNLinSol_SPGMR(y_data, - has_preconditioner ? SUN_PREC_LEFT : SUN_PREC_NONE, + preconditioning, max_linear_iters, sunctx); AMREX_ALWAYS_ASSERT(LS != nullptr); @@ -221,7 +255,7 @@ private: if (use_ark) { flag = ARKStepSetLinearSolver(arkode_mem, LS, nullptr); AMREX_ALWAYS_ASSERT(flag == 0); - if (has_preconditioner) { + if (use_preconditioner) { flag = ARKStepSetPreconditioner( arkode_mem, BaseT::PrecondSetup ? SundialsUserFun::psetup : nullptr, @@ -231,7 +265,7 @@ private: } else if (use_mri) { flag = MRIStepSetLinearSolver(arkode_mem, LS, nullptr); AMREX_ALWAYS_ASSERT(flag == 0); - if (has_preconditioner) { + if (use_preconditioner) { flag = MRIStepSetPreconditioner( arkode_mem, BaseT::PrecondSetup ? SundialsUserFun::psetup : nullptr, @@ -249,10 +283,13 @@ private: return; } - const bool has_preconditioner = static_cast(BaseT::FastPrecondSolve); - if (BaseT::FastPrecondSetup && !has_preconditioner) { + if (BaseT::FastPrecondSetup && !BaseT::FastPrecondSolve) { amrex::Error("SUNDIALS fast preconditioner setup callback requires a solve callback."); } + const int preconditioning = BaseT::FastPrecondSolve + ? fast_linear_solver_preconditioning_type + : SUN_PREC_NONE; + const bool use_preconditioner = preconditioning != SUN_PREC_NONE; if (!fast_ls_needs_update && fast_LS != nullptr) { return; @@ -260,13 +297,13 @@ private: SUNLinSolFree(fast_LS); fast_LS = SUNLinSol_SPGMR(y_data, - has_preconditioner ? SUN_PREC_LEFT : SUN_PREC_NONE, + preconditioning, fast_max_linear_iters, sunctx); AMREX_ALWAYS_ASSERT(fast_LS != nullptr); int flag = ARKStepSetLinearSolver(arkode_fast_mem, fast_LS, nullptr); AMREX_ALWAYS_ASSERT(flag == 0); - if (has_preconditioner) { + if (use_preconditioner) { flag = ARKStepSetPreconditioner( arkode_fast_mem, BaseT::FastPrecondSetup ? SundialsUserFun::fast_psetup : nullptr, @@ -308,9 +345,16 @@ private: pp.query("fast_max_nonlinear_iters", fast_max_nonlinear_iters); pp.query("linear_solver", linear_solver); + pp.query("linear_solver_preconditioning", linear_solver_preconditioning); + linear_solver_preconditioning_type = + PreconditioningType(linear_solver_preconditioning, "linear_solver_preconditioning"); pp.query("max_linear_iters", max_linear_iters); pp.query("fast_linear_solver", fast_linear_solver); + pp.query("fast_linear_solver_preconditioning", fast_linear_solver_preconditioning); + fast_linear_solver_preconditioning_type = + PreconditioningType(fast_linear_solver_preconditioning, + "fast_linear_solver_preconditioning"); pp.query("fast_max_linear_iters", fast_max_linear_iters); set_stop_time = pp.query("stop_time", stop_time); @@ -386,6 +430,8 @@ private: if (nonlinear_solver == "Newton") { if (amrex::Verbose()) { amrex::Print() << "Linear solver: " << linear_solver << "\n"; + amrex::Print() << "Linear solver preconditioning: " + << linear_solver_preconditioning << "\n"; amrex::Print() << "Max linear iters: " << max_linear_iters << "\n"; } slow_ls_needs_update = true; @@ -451,6 +497,8 @@ private: if (fast_nonlinear_solver == "Newton") { if (amrex::Verbose()) { amrex::Print() << "Linear solver: " << fast_linear_solver << "\n"; + amrex::Print() << "Linear solver preconditioning: " + << fast_linear_solver_preconditioning << "\n"; amrex::Print() << "Max linear iters: " << fast_max_linear_iters << "\n"; } fast_ls_needs_update = true; @@ -547,6 +595,8 @@ private: if (nonlinear_solver == "Newton") { if (amrex::Verbose()) { amrex::Print() << "Linear solver: " << linear_solver << "\n"; + amrex::Print() << "Linear solver preconditioning: " + << linear_solver_preconditioning << "\n"; amrex::Print() << "Max linear iters: " << max_linear_iters << "\n"; } slow_ls_needs_update = true; @@ -799,7 +849,7 @@ public: bool jcur = false; BaseT::PrecondSetup(S_data, S_rhs, rhs_time, jok, jcur, gamma); *jcurPtr = jcur ? SUNTRUE : SUNFALSE; - + return 0; }; @@ -822,8 +872,9 @@ public: AMREX_ALWAYS_ASSERT_WITH_MESSAGE(static_cast(BaseT::PrecondSolve), "SUNDIALS preconditioner solve callback not set."); - return BaseT::PrecondSolve(S_z, S_r, S_data, S_rhs, - rhs_time, gamma, delta, lr); + BaseT::PrecondSolve(S_z, S_r, S_data, S_rhs, + rhs_time, gamma, delta, lr); + return 0; }; udata.fast_psetup = [&](amrex::Real rhs_time, N_Vector y_data, N_Vector y_rhs, @@ -837,10 +888,10 @@ public: unpack_vector(y_rhs, S_rhs); bool jcur = false; - BaseT::FastPrecondSetup(S_data, S_rhs, rhs_time, jcur, gamma); + BaseT::FastPrecondSetup(S_data, S_rhs, rhs_time, jok, jcur, gamma); *jcurPtr = jcur ? SUNTRUE : SUNFALSE; - - return ierr; + + return 0; }; udata.fast_psolve = [&](amrex::Real rhs_time, N_Vector y_data, N_Vector y_rhs, @@ -862,8 +913,9 @@ public: AMREX_ALWAYS_ASSERT_WITH_MESSAGE(static_cast(BaseT::FastPrecondSolve), "SUNDIALS fast preconditioner solve callback not set."); - return BaseT::FastPrecondSolve(S_z, S_r, S_data, S_rhs, - rhs_time, gamma, delta, lr); + BaseT::FastPrecondSolve(S_z, S_r, S_data, S_rhs, + rhs_time, gamma, delta, lr); + return 0; }; udata.post_stage = [&](amrex::Real time, N_Vector y_data,