diff --git a/Docs/sphinx_documentation/source/RuntimeParameters.rst b/Docs/sphinx_documentation/source/RuntimeParameters.rst index 9642de257b2..72839c2a2ec 100644 --- a/Docs/sphinx_documentation/source/RuntimeParameters.rst +++ b/Docs/sphinx_documentation/source/RuntimeParameters.rst @@ -1502,6 +1502,21 @@ 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.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 @@ -1533,7 +1548,22 @@ 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_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 diff --git a/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst b/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst index 722fdbe973a..d0fe5b014e8 100644 --- a/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst +++ b/Docs/sphinx_documentation/source/TimeIntegration_Chapter.rst @@ -173,6 +173,15 @@ 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. 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 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..2df06da8782 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..d634bf8a481 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..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 @@ -37,6 +38,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 +73,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); @@ -117,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 @@ -139,12 +182,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 +198,122 @@ 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" && + type != "IM-MRI" && type != "IMEX-MRI") || + nonlinear_solver != "Newton") + { + return; + } + + 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; + } + + SUNLinSolFree(LS); + LS = SUNLinSol_SPGMR(y_data, + preconditioning, + 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 (use_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 (use_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; + } + + 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; + } + + SUNLinSolFree(fast_LS); + fast_LS = SUNLinSol_SPGMR(y_data, + 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 (use_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"); @@ -184,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); @@ -262,12 +430,12 @@ 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"; } - 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); } } @@ -329,18 +497,18 @@ 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 = 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") { @@ -427,12 +595,11 @@ 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"; } - 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 +619,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 +836,88 @@ 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 = 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, + 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."); + 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, + 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 = false; + BaseT::FastPrecondSetup(S_data, S_rhs, rhs_time, jok, jcur, gamma); + *jcurPtr = jcur ? SUNTRUE : SUNFALSE; + + return 0; + }; + + 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."); + 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, void * /* user_data */) -> int { @@ -762,6 +1012,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 +1042,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 +1082,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) {