apply solver builder to jit solver

include/micm/solver/backward_euler.hpp

@@ -46,18 +46,23 @@ namespace micm
     /// @brief Default constructor
     BackwardEuler(
         BackwardEulerSolverParameters parameters,
-        LinearSolverPolicy linear_solver,
-        ProcessSetPolicy process_set,
-        std::vector<std::size_t> jacobian_diagonal_elements,
+        LinearSolverPolicy&& linear_solver,
+        ProcessSetPolicy&& process_set,
+        auto& jacobian,
         std::vector<micm::Process>& processes)
         : parameters_(parameters),
-          linear_solver_(linear_solver),
-          process_set_(process_set),
-          jacobian_diagonal_elements_(jacobian_diagonal_elements),
+          linear_solver_(std::move(linear_solver)),
+          process_set_(std::move(process_set)),
+          jacobian_diagonal_elements_(jacobian.DiagonalIndices(0)),
           processes_(processes)
     {
     }
 
+    BackwardEuler(const BackwardEuler&) = delete;
+    BackwardEuler& operator=(const BackwardEuler&) = delete;
+    BackwardEuler(BackwardEuler&&) = default;
+    BackwardEuler& operator=(BackwardEuler&&) = default;
+
     virtual ~BackwardEuler() = default;
 
     /// @brief Advances the given step over the specified time step

include/micm/solver/jit_rosenbrock.hpp

@@ -19,6 +19,7 @@
 #include <micm/process/jit_process_set.hpp>
 #include <micm/solver/jit_linear_solver.hpp>
 #include <micm/solver/rosenbrock.hpp>
+#include <micm/solver/rosenbrock_solver_parameters.hpp>
 #include <micm/util/random_string.hpp>
 #include <micm/util/sparse_matrix_vector_ordering.hpp>
 #include <micm/util/vector_matrix.hpp>
@@ -30,22 +31,24 @@
 namespace micm
 {
 
-  template<
-      template<class> class MatrixPolicy = VectorMatrix,
-      class SparseMatrixPolicy = DefaultVectorSparseMatrix,
-      class LinearSolverPolicy = JitLinearSolver<MICM_DEFAULT_VECTOR_SIZE, SparseMatrixPolicy>,
-      class ProcessSetPolicy = JitProcessSet<MICM_DEFAULT_VECTOR_SIZE>>
-  class JitRosenbrockSolver : public RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy, ProcessSetPolicy>
+  /// @brief A Rosenbrock solver with JIT-compiled optimizations
+
+  template<class ProcessSetPolicy, class LinearSolverPolicy>
+  class JitRosenbrockSolver : public RosenbrockSolver<ProcessSetPolicy, LinearSolverPolicy>
   {
     llvm::orc::ResourceTrackerSP function_resource_tracker_;
     using FuncPtr = void (*)(double*, const double);
     FuncPtr alpha_minus_jacobian_ = nullptr;
 
    public:
+
+    /// @brief Solver parameters typename
+    using ParametersType = RosenbrockSolverParameters;
+
     JitRosenbrockSolver(const JitRosenbrockSolver&) = delete;
     JitRosenbrockSolver& operator=(const JitRosenbrockSolver&) = delete;
     JitRosenbrockSolver(JitRosenbrockSolver&& other)
-        : RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy, ProcessSetPolicy>(std::move(other)),
+        : RosenbrockSolver<ProcessSetPolicy, LinearSolverPolicy>(std::move(other)),
           function_resource_tracker_(std::move(other.function_resource_tracker_)),
           alpha_minus_jacobian_(std::move(other.alpha_minus_jacobian_))
     {
@@ -54,42 +57,33 @@ namespace micm
 
     JitRosenbrockSolver& operator=(JitRosenbrockSolver&& other)
     {
-      RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy, ProcessSetPolicy>::operator=(std::move(other));
+      RosenbrockSolver<ProcessSetPolicy, LinearSolverPolicy>::operator=(std::move(other));
       function_resource_tracker_ = std::move(other.function_resource_tracker_);
       alpha_minus_jacobian_ = std::move(other.alpha_minus_jacobian_);
       other.alpha_minus_jacobian_ = NULL;
       return *this;
     }
 
     /// @brief Builds a Rosenbrock solver for the given system, processes, and solver parameters
-    /// @param system The chemical system to create the solver for
-    /// @param processes The collection of chemical processes that will be applied during solving
+    /// @param parameters Solver parameters
+    /// @param linear_solver Linear solver
+    /// @param process_set Process set
+    /// @param jacobian Jacobian matrix
+    /// @param processes Vector of processes
     JitRosenbrockSolver(
-        const System& system,
-        const std::vector<Process>& processes,
-        const RosenbrockSolverParameters& parameters)
-        : RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy, ProcessSetPolicy>(
-              system,
-              processes,
+        RosenbrockSolverParameters parameters,
+        LinearSolverPolicy linear_solver,
+        ProcessSetPolicy process_set,
+        auto& jacobian,
+        std::vector<Process>& processes)
+        : RosenbrockSolver<ProcessSetPolicy, LinearSolverPolicy>(
               parameters,
-              [&](const SparseMatrixPolicy& matrix, double initial_value) -> LinearSolverPolicy {
-                return LinearSolverPolicy{ matrix, initial_value };
-              },
-              [&](const std::vector<Process>& processes,
-                  const std::map<std::string, std::size_t>& variable_map) -> ProcessSetPolicy {
-                return ProcessSetPolicy{ processes, variable_map };
-              })
+              std::move(linear_solver),
+              std::move(process_set),
+              jacobian,
+              processes)
     {
-      MatrixPolicy<double> temp{};
-      if (temp.GroupVectorSize() != parameters.number_of_grid_cells_)
-      {
-        std::string msg =
-            "JIT functions require the number of grid cells solved together to match the vector dimension template "
-            "parameter, currently: " +
-            std::to_string(temp.GroupVectorSize());
-        throw std::system_error(make_error_code(MicmJitErrc::InvalidMatrix), msg);
-      }
-      this->GenerateAlphaMinusJacobian();
+      this->GenerateAlphaMinusJacobian(jacobian);
     }
 
     ~JitRosenbrockSolver()
@@ -104,8 +98,17 @@ namespace micm
     /// @brief compute [alpha * I - dforce_dy]
     /// @param jacobian Jacobian matrix (dforce_dy)
     /// @param alpha
+    template<class SparseMatrixPolicy>
     void AlphaMinusJacobian(SparseMatrixPolicy& jacobian, const double& alpha) const
     {
+      if (jacobian.GroupVectorSize() != jacobian.NumberOfBlocks())
+      {
+        std::string msg =
+            "JIT functions require the number of grid cells solved together to match the vector dimension template "
+            "parameter, currently: " +
+            std::to_string(jacobian.GroupVectorSize());
+        throw std::system_error(make_error_code(MicmJitErrc::InvalidMatrix), msg);
+      }
       double a = alpha;
       if (alpha_minus_jacobian_)
       {
@@ -119,9 +122,9 @@ namespace micm
     }
 
    private:
-    void GenerateAlphaMinusJacobian()
+    void GenerateAlphaMinusJacobian(auto& jacobian)
     {
-      auto jacobian = this->GetState().jacobian_;
+      auto diagonal_elements = jacobian.DiagonalIndices(0);
       // save sizes needed throughout the function
       std::size_t n_cells = jacobian.GroupVectorSize();
       std::size_t number_of_nonzero_jacobian_elements = jacobian.AsVector().size();
@@ -143,7 +146,7 @@ namespace micm
 
       // iterative over the blocks of the jacobian and add the alpha value
       // jacobian_vector[i_elem + i_cell] += alpha;
-      for (const auto& i_elem : this->state_parameters_.jacobian_diagonal_elements_)
+      for (const auto& i_elem : diagonal_elements)
       {
         llvm::Value* ptr_index[1];
 

include/micm/solver/jit_solver_builder.hpp

@@ -11,10 +11,15 @@ namespace micm
 {
   /// @brief Builder of CPU-based solvers optimized for a particular chemical system using JIT compilation
   /// @tparam SolverParametersPolicy Policy for the ODE solver
-  /// @tparam T Primary data type
   /// @tparam L Vector size
   ///
   /// JIT-compiled solvers only work with vector-ordered matrices
-  template<class SolverParametersPolicy, class T = double, std::size_t L = MICM_DEFAULT_VECTOR_SIZE>
-  using JitSolverBuilder = SolverBuilder<SolverParametersPolicy, VectorMatrix<T, L>, SparseMatrix<T, SparseMatrixVectorOrdering<L>>, JitProcessSet, JitLinearSolver<SparseMatrixVectorOrdering<L>>, JitLuDecomposition>>;
+  template<class SolverParametersPolicy, std::size_t L = MICM_DEFAULT_VECTOR_SIZE>
+  using JitSolverBuilder = SolverBuilder<SolverParametersPolicy,
+                                         VectorMatrix<double, L>,
+                                         SparseMatrix<double, SparseMatrixVectorOrdering<L>>,
+                                         JitProcessSet<L>,
+                                         JitLinearSolver<L,
+                                                         SparseMatrix<double, SparseMatrixVectorOrdering<L>>,
+                                                         JitLuDecomposition<L>>>;
 } // namespace micm

include/micm/solver/jit_solver_parameters.hpp

@@ -0,0 +1,25 @@
+// Copyright (C) 2023-2024 National Center for Atmospheric Research
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include "rosenbrock_solver_parameters.hpp"
+#include <micm/solver/rosenbrock_solver_parameters.hpp>
+#include <micm/solver/jit_rosenbrock.hpp>
+
+namespace micm
+{
+    /// @brief Parameters for the JIT Rosenbrock solver
+    struct JitRosenbrockSolverParameters : public RosenbrockSolverParameters
+    {
+        template<class ProcessSetPolicy, class LinearSolverPolicy>
+        using SolverType = JitRosenbrockSolver<ProcessSetPolicy, LinearSolverPolicy>;
+
+        /// @brief Constructor from base class
+        /// @param base
+        JitRosenbrockSolverParameters(const RosenbrockSolverParameters& base)
+            : RosenbrockSolverParameters(base)
+        {
+        }
+
+    };
+}

include/micm/solver/linear_solver.hpp

@@ -57,6 +57,11 @@ namespace micm
     /// @brief default constructor
     LinearSolver(){};
 
+    LinearSolver(const LinearSolver&) = delete;
+    LinearSolver& operator=(const LinearSolver&) = delete;
+    LinearSolver(LinearSolver&&) = default;
+    LinearSolver& operator=(LinearSolver&&) = default;
+
     /// @brief Constructs a linear solver for the sparsity structure of the given matrix
     /// @param matrix Sparse matrix
     /// @param initial_value Initial value for matrix elements

include/micm/solver/lu_decomposition.hpp

@@ -7,6 +7,17 @@
 #include <micm/profiler/instrumentation.hpp>
 #include <micm/util/sparse_matrix.hpp>
 
+namespace
+{
+  template<typename T>
+  concept SparseMatrixConcept = requires(T t)
+  {
+    t.NumRows();
+    t.NumColumns();
+    t.NumberOfBlocks();
+  };
+}
+
 namespace micm
 {
 
@@ -76,20 +87,28 @@ namespace micm
     /// @brief default constructor
     LuDecomposition();
 
+    LuDecomposition(const LuDecomposition&) = delete;
+    LuDecomposition& operator=(const LuDecomposition&) = delete;
+    LuDecomposition(LuDecomposition&&) = default;
+    LuDecomposition& operator=(LuDecomposition&&) = default;
+
     /// @brief Construct an LU decomposition algorithm for a given sparse matrix
     /// @param matrix Sparse matrix
     template<class SparseMatrixPolicy>
+    requires(SparseMatrixConcept<SparseMatrixPolicy>)
     LuDecomposition(const SparseMatrixPolicy& matrix);
 
     /// @brief Create an LU decomposition algorithm for a given sparse matrix policy
     /// @param matrix Sparse matrix
     template<class SparseMatrixPolicy>
+    requires(SparseMatrixConcept<SparseMatrixPolicy>)
     static LuDecomposition Create(const SparseMatrixPolicy& matrix);
 
     /// @brief Create sparse L and U matrices for a given A matrix
     /// @param A Sparse matrix that will be decomposed
     /// @return L and U Sparse matrices
     template<class SparseMatrixPolicy>
+    requires(SparseMatrixConcept<SparseMatrixPolicy>)
     static std::pair<SparseMatrixPolicy, SparseMatrixPolicy> GetLUMatrices(
         const SparseMatrixPolicy& A,
         typename SparseMatrixPolicy::value_type initial_value);
@@ -99,6 +118,7 @@ namespace micm
     /// @param L The lower triangular matrix created by decomposition
     /// @param U The upper triangular matrix created by decomposition
     template<class SparseMatrixPolicy>
+    requires(SparseMatrixConcept<SparseMatrixPolicy>)
     void Decompose(const SparseMatrixPolicy& A, SparseMatrixPolicy& L, SparseMatrixPolicy& U) const;
 
     /// @brief Perform an LU decomposition on a given A matrix
@@ -123,6 +143,7 @@ namespace micm
     /// @brief Initialize arrays for the LU decomposition
     /// @param A Sparse matrix to decompose
     template<class SparseMatrixPolicy>
+    requires(SparseMatrixConcept<SparseMatrixPolicy>)
     void Initialize(const SparseMatrixPolicy& matrix, auto initial_value);
   };
 

include/micm/solver/lu_decomposition.inl

@@ -10,12 +10,14 @@ namespace micm
   }
 
   template<class SparseMatrixPolicy>
+  requires(SparseMatrixConcept<SparseMatrixPolicy>)
   inline LuDecomposition::LuDecomposition(const SparseMatrixPolicy& matrix)
   {
     Initialize<SparseMatrixPolicy>(matrix, typename SparseMatrixPolicy::value_type());
   }
 
   template<class SparseMatrixPolicy>
+  requires(SparseMatrixConcept<SparseMatrixPolicy>)
   inline LuDecomposition LuDecomposition::Create(const SparseMatrixPolicy& matrix)
   {
     LuDecomposition lu_decomp{};
@@ -24,6 +26,7 @@ namespace micm
   }
 
   template<class SparseMatrixPolicy>
+  requires(SparseMatrixConcept<SparseMatrixPolicy>)
   inline void LuDecomposition::Initialize(const SparseMatrixPolicy& matrix, auto initial_value)
   {
     MICM_PROFILE_FUNCTION();
@@ -100,6 +103,7 @@ namespace micm
   }
 
   template<class SparseMatrixPolicy>
+  requires(SparseMatrixConcept<SparseMatrixPolicy>)
   inline std::pair<SparseMatrixPolicy, SparseMatrixPolicy> LuDecomposition::GetLUMatrices(
       const SparseMatrixPolicy& A,
       typename SparseMatrixPolicy::value_type initial_value)
@@ -162,6 +166,7 @@ namespace micm
   }
 
   template<class SparseMatrixPolicy>
+  requires(SparseMatrixConcept<SparseMatrixPolicy>)
   inline void LuDecomposition::Decompose(const SparseMatrixPolicy& A, SparseMatrixPolicy& L, SparseMatrixPolicy& U) const
   {
     bool is_singular = false;

include/micm/solver/rosenbrock.hpp

@@ -68,16 +68,21 @@ namespace micm
         RosenbrockSolverParameters parameters,
         LinearSolverPolicy linear_solver,
         ProcessSetPolicy process_set,
-        std::vector<std::size_t> jacobian_diagonal_elements,
+        auto& jacobian,
         std::vector<Process>& processes)
         : parameters_(parameters),
-          linear_solver_(linear_solver),
-          process_set_(process_set),
-          jacobian_diagonal_elements_(jacobian_diagonal_elements),
+          linear_solver_(std::move(linear_solver)),
+          process_set_(std::move(process_set)),
+          jacobian_diagonal_elements_(jacobian.DiagonalIndices(0)),
           processes_(processes)
     {
     }
 
+    RosenbrockSolver(const RosenbrockSolver&) = delete;
+    RosenbrockSolver& operator=(const RosenbrockSolver&) = delete;
+    RosenbrockSolver(RosenbrockSolver&&) = default;
+    RosenbrockSolver& operator=(RosenbrockSolver&&) = default;
+
     virtual ~RosenbrockSolver() = default;
 
     /// @brief Advances the given step over the specified time step

include/micm/solver/solver.hpp

@@ -21,12 +21,12 @@ namespace micm
     SolverPolicy solver_;
 
     Solver(
-        SolverPolicy solver,
+        SolverPolicy&& solver,
         StateParameters state_parameters,
         std::size_t number_of_grid_cells,
         std::size_t number_of_species,
         std::size_t number_of_reactions)
-        : solver_(solver),
+        : solver_(std::move(solver)),
           number_of_grid_cells_(number_of_grid_cells),
           number_of_species_(number_of_species),
           number_of_reactions_(number_of_reactions),

include/micm/solver/solver_builder.hpp

@@ -102,10 +102,6 @@ namespace micm
     /// @return The labels of the custom parameters
     std::vector<std::string> GetCustomParameterLabels() const;
 
-    /// @brief Returns the diagonal elements of the jacobian matrix
-    /// @param jacobian
-    /// @return The diagonal elements of the jacobian matrix
-    std::vector<std::size_t> GetJacobianDiagonalElements(auto jacobian) const;
   };
 
   /// @brief Builder of CPU-based general solvers