Add preconditoner routines to TimeIntegrator with SUNDIALS support

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

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

Src/Base/AMReX_IntegratorBase.H

@@ -179,6 +179,18 @@ private:
     pp.query("fast_abs_tol", fast_abs_tol);
   }
 protected:
+    using PrecondSetupFunction = std::function<void(T& state, T& rhs,
+                                                    const amrex::Real time,
+                                                    bool jok, bool& jcur,
+                                                    const amrex::Real gamma)>;
+
+    using PrecondSolveFunction = std::function<void(T& soln, T& rhs,
+                                                    T& state, T& state_rhs,
+                                                    const amrex::Real time,
+                                                    const amrex::Real gamma,
+                                                    const amrex::Real delta,
+                                                    int lr)>;
+
     /**
      * \brief Rhs is the right-hand-side function the integrator will use.
      */
@@ -226,6 +238,30 @@ protected:
      */
     std::function<void (T&, amrex::Real)> 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<void (T&, amrex::Real)> A)
     {
         post_stage_action = A;

Src/Base/AMReX_TimeIntegrator.H

@@ -147,6 +147,22 @@ public:
         integrator_ptr->set_fast_rhs(F);
     }
 
+    void set_preconditioner (
+        std::function<void(T&, T&, const amrex::Real, bool, bool&, const amrex::Real)> Psetup,
+        std::function<void(T&, T&, T&, T&, const amrex::Real, const amrex::Real,
+                           const amrex::Real, int)> Psolve)
+    {
+        integrator_ptr->set_preconditioner(std::move(Psetup), std::move(Psolve));
+    }
+
+    void set_fast_preconditioner (
+        std::function<void(T&, T&, const amrex::Real, bool, bool&, const amrex::Real)> Psetup,
+        std::function<void(T&, T&, T&, T&, const amrex::Real, const amrex::Real,
+                           const amrex::Real, int)> Psolve)
+    {
+        integrator_ptr->set_fast_preconditioner(std::move(Psetup), std::move(Psolve));
+    }
+
     void set_post_stage_action (std::function<void (T&, amrex::Real)> A)
     {
         integrator_ptr->set_post_stage_action(A);