Merge pull request #287 from NCAR/develop-261-jit-solver-test-alt

`JitRosenbrockSolver` integration and regression tests
NCAR · Oct 4, 2023 · 1caf3ee · 1caf3ee
2 parents 3940a85 + c82e881
commit 1caf3ee
Show file tree

Hide file tree

Showing 35 changed files with 1,937 additions and 1,028 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -42,6 +42,7 @@ option(ENABLE_CUDA "Build with Cuda support" OFF)
 option(ENABLE_OPENACC "Build with OpenACC Support" OFF)
 option(ENABLE_LLVM "Build with LLVM support for JIT-compiling" OFF)
 option(ENABLE_TESTS "Build the tests" ON)
+set(DEFAULT_VECTOR_MATRIX_SIZE "4" CACHE STRING "Default size for vectorizable matrix types")
 
 include(CMakeDependentOption)
 # Option to collect custom OpenACC flags
@@ -52,6 +53,8 @@ if(ENABLE_GPU_TYPE)
   set(GPU_TYPE "" CACHE STRING "The GPU type being targeted")
 endif()
 
+add_compile_definitions(DEFAULT_VECTOR_SIZE=${DEFAULT_VECTOR_MATRIX_SIZE})
+
 ################################################################################
 # Dependencies
 

diff --git a/include/micm/process/jit_process_set.hpp b/include/micm/process/jit_process_set.hpp
@@ -14,7 +14,7 @@ namespace micm
 
   /// @brief JIT-compiled solver function calculators for a collection of processes
   ///        The template parameter is the number of grid cells to solve simultaneously
-  template<std::size_t L>
+  template<std::size_t L = DEFAULT_VECTOR_SIZE>
   class JitProcessSet : public ProcessSet
   {
     std::shared_ptr<JitCompiler> compiler_;
@@ -24,6 +24,11 @@ namespace micm
     void (*jacobian_function_)(const double *, const double *, double *);
 
    public:
+    JitProcessSet(const JitProcessSet &) = delete;
+    JitProcessSet &operator=(const JitProcessSet &) = delete;
+    JitProcessSet(JitProcessSet &&);
+    JitProcessSet &operator=(JitProcessSet &&);
+
     /// @brief Create a JITed process set calculator for a given set of processes
     /// @param compiler JIT compiler
     /// @param processes Processes to create calculator for
@@ -70,6 +75,33 @@ namespace micm
     void GenerateJacobianFunction(const SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix);
   };
 
+  template<std::size_t L>
+  inline JitProcessSet<L>::JitProcessSet(JitProcessSet &&other)
+      : ProcessSet(std::move(other)),
+        compiler_(std::move(other.compiler_)),
+        forcing_function_resource_tracker_(std::move(other.forcing_function_resource_tracker_)),
+        forcing_function_(std::move(other.forcing_function_)),
+        jacobian_function_resource_tracker_(std::move(other.jacobian_function_resource_tracker_)),
+        jacobian_function_(std::move(other.jacobian_function_))
+  {
+    other.forcing_function_ = NULL;
+    other.jacobian_function_ = NULL;
+  }
+
+  template<std::size_t L>
+  inline JitProcessSet<L> &JitProcessSet<L>::operator=(JitProcessSet &&other)
+  {
+    ProcessSet::operator=(std::move(other));
+    compiler_ = std::move(other.compiler_);
+    forcing_function_resource_tracker_ = std::move(other.forcing_function_resource_tracker_);
+    forcing_function_ = std::move(other.forcing_function_);
+    jacobian_function_resource_tracker_ = std::move(other.jacobian_function_resource_tracker_);
+    jacobian_function_ = std::move(other.jacobian_function_);
+    other.forcing_function_ = NULL;
+    other.jacobian_function_ = NULL;
+    return *this;
+  }
+
   template<std::size_t L>
   template<template<class> class MatrixPolicy>
   inline JitProcessSet<L>::JitProcessSet(
@@ -81,6 +113,10 @@ namespace micm
   {
     forcing_function_ = NULL;
     jacobian_function_ = NULL;
+    if (state.variables_.size() != L || state.variables_.GroupVectorSize() != L)
+    {
+      throw std::runtime_error("Invalid state for JitProcessSet. Check the the VectorMatrix template parameters.");
+    }
     this->GenerateForcingFunction();
   }
 

diff --git a/include/micm/solver/jit_linear_solver.hpp b/include/micm/solver/jit_linear_solver.hpp
@@ -5,6 +5,7 @@
 
 #include <micm/jit/jit_compiler.hpp>
 #include <micm/jit/jit_function.hpp>
+#include <micm/solver/jit_lu_decomposition.hpp>
 #include <micm/solver/linear_solver.hpp>
 #include <micm/util/sparse_matrix.hpp>
 #include <micm/util/sparse_matrix_vector_ordering.hpp>
@@ -16,14 +17,21 @@ namespace micm
   ///
   /// See LinearSolver class description for algorithm details
   /// The template parameter is the number of blocks (i.e. grid cells) in the block-diagonal matrix
-  template<std::size_t L, template<class> class SparseMatrixPolicy>
-  class JitLinearSolver : public LinearSolver<double, SparseMatrixPolicy>
+  template<std::size_t L, template<class> class SparseMatrixPolicy, class LuDecompositionPolicy = JitLuDecomposition<L>>
+  class JitLinearSolver : public LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>
   {
     std::shared_ptr<JitCompiler> compiler_;
     llvm::orc::ResourceTrackerSP solve_function_resource_tracker_;
     void (*solve_function_)(const double*, double*, const double*, const double*);
 
    public:
+    JitLinearSolver(){};
+
+    JitLinearSolver(const JitLinearSolver&) = delete;
+    JitLinearSolver& operator=(const JitLinearSolver&) = delete;
+    JitLinearSolver(JitLinearSolver&&);
+    JitLinearSolver& operator=(JitLinearSolver&&);
+
     /// @brief Create a JITed linear solver for a given sparse matrix structure
     /// @param compiler JIT compiler
     /// @param matrix Block-diagonal sparse matrix to create solver for

diff --git a/include/micm/solver/jit_linear_solver.inl b/include/micm/solver/jit_linear_solver.inl
@@ -4,19 +4,50 @@
 namespace micm
 {
 
-  template<std::size_t L, template<class> class SparseMatrixPolicy>
-  inline JitLinearSolver<L, SparseMatrixPolicy>::JitLinearSolver(
+  template<std::size_t L, template<class> class SparseMatrixPolicy, class LuDecompositionPolicy>
+  inline JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::JitLinearSolver(JitLinearSolver &&other)
+      : LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>(std::move(other)),
+        compiler_(std::move(other.compiler_)),
+        solve_function_resource_tracker_(std::move(other.solve_function_resource_tracker_)),
+        solve_function_(std::move(other.solve_function_))
+  {
+    other.solve_function_ = NULL;
+  }
+
+  template<std::size_t L, template<class> class SparseMatrixPolicy, class LuDecompositionPolicy>
+  inline JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy> &
+  JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::operator=(JitLinearSolver &&other)
+  {
+    LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::operator=(std::move(other));
+    compiler_ = std::move(other.compiler_);
+    solve_function_resource_tracker_ = std::move(other.solve_function_resource_tracker_);
+    solve_function_ = std::move(other.solve_function_);
+    other.solve_function_ = NULL;
+    return *this;
+  }
+
+  template<std::size_t L, template<class> class SparseMatrixPolicy, class LuDecompositionPolicy>
+  inline JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::JitLinearSolver(
       std::shared_ptr<JitCompiler> compiler,
       const SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix,
       double initial_value)
-      : LinearSolver<double, SparseMatrixPolicy>(matrix, initial_value),
+      : LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>(
+            matrix,
+            initial_value,
+            [&](const SparseMatrixPolicy<double> &m) -> LuDecompositionPolicy
+            { return LuDecompositionPolicy(compiler, m); }),
         compiler_(compiler)
   {
     solve_function_ = NULL;
+    if (matrix.size() != L || matrix.GroupVectorSize() != L)
+    {
+      throw std::runtime_error("Invalid matrix for JitLinearSolver. Check the the VectorMatrix template parameters.");
+    }
+    GenerateSolveFunction();
   }
 
-  template<std::size_t L, template<class> class SparseMatrixPolicy>
-  inline JitLinearSolver<L, SparseMatrixPolicy>::~JitLinearSolver()
+  template<std::size_t L, template<class> class SparseMatrixPolicy, class LuDecompositionPolicy>
+  inline JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::~JitLinearSolver()
   {
     if (solve_function_ != NULL)
     {
@@ -25,26 +56,28 @@ namespace micm
     }
   }
 
-  template<std::size_t L, template<class> class SparseMatrixPolicy>
-  inline void JitLinearSolver<L, SparseMatrixPolicy>::Factor(SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix)
+  template<std::size_t L, template<class> class SparseMatrixPolicy, class LuDecompositionPolicy>
+  inline void JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::Factor(
+      SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix)
   {
-    LinearSolver<double, SparseMatrixPolicy>::Factor(matrix);
-    GenerateSolveFunction();
+    LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::Factor(matrix);
   }
 
-  template<std::size_t L, template<class> class SparseMatrixPolicy>
+  template<std::size_t L, template<class> class SparseMatrixPolicy, class LuDecompositionPolicy>
   template<template<class> class MatrixPolicy>
-  inline void JitLinearSolver<L, SparseMatrixPolicy>::Solve(const MatrixPolicy<double> &b, MatrixPolicy<double> &x)
+  inline void JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::Solve(
+      const MatrixPolicy<double> &b,
+      MatrixPolicy<double> &x)
   {
     solve_function_(
         b.AsVector().data(),
         x.AsVector().data(),
-        LinearSolver<double, SparseMatrixPolicy>::lower_matrix_.AsVector().data(),
-        LinearSolver<double, SparseMatrixPolicy>::upper_matrix_.AsVector().data());
+        LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::lower_matrix_.AsVector().data(),
+        LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::upper_matrix_.AsVector().data());
   }
 
-  template<std::size_t L, template<class> class SparseMatrixPolicy>
-  inline void JitLinearSolver<L, SparseMatrixPolicy>::GenerateSolveFunction()
+  template<std::size_t L, template<class> class SparseMatrixPolicy, class LuDecompositionPolicy>
+  inline void JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::GenerateSolveFunction()
   {
     JitFunction func = JitFunction::create(compiler_)
                            .name("linear_solve")
@@ -54,10 +87,10 @@ namespace micm
                                         { "U", JitType::DoublePtr } })
                            .return_type(JitType::Void);
     llvm::Type *double_type = func.GetType(JitType::Double);
-    auto Lij_yj = LinearSolver<double, SparseMatrixPolicy>::Lij_yj_.begin();
-    auto Uij_xj = LinearSolver<double, SparseMatrixPolicy>::Uij_xj_.begin();
+    auto Lij_yj = LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::Lij_yj_.begin();
+    auto Uij_xj = LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::Uij_xj_.begin();
     std::size_t offset = 0;
-    for (auto &nLij_Lii : LinearSolver<double, SparseMatrixPolicy>::nLij_Lii_)
+    for (auto &nLij_Lii : LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::nLij_Lii_)
     {
       // the x vector is used for y values to conserve memory
       {
@@ -107,8 +140,8 @@ namespace micm
       }
       offset += L;
     }
-    offset = L * LinearSolver<double, SparseMatrixPolicy>::nUij_Uii_.size();
-    for (auto &nUij_Uii : LinearSolver<double, SparseMatrixPolicy>::nUij_Uii_)
+    offset = L * LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::nUij_Uii_.size();
+    for (auto &nUij_Uii : LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::nUij_Uii_)
     {
       offset -= L;
       for (std::size_t i = 0; i < nUij_Uii.first; ++i)

diff --git a/include/micm/solver/jit_lu_decomposition.hpp b/include/micm/solver/jit_lu_decomposition.hpp
@@ -15,14 +15,21 @@ namespace micm
   ///
   /// See LuDecomposition class description for algorithm details
   /// The template parameter is the number of blocks (i.e. grid cells) in the block-diagonal matrix
-  template<std::size_t L>
+  template<std::size_t L = DEFAULT_VECTOR_SIZE>
   class JitLuDecomposition : public LuDecomposition
   {
     std::shared_ptr<JitCompiler> compiler_;
     llvm::orc::ResourceTrackerSP decompose_function_resource_tracker_;
     void (*decompose_function_)(const double *, double *, double *);
 
    public:
+    JitLuDecomposition(){};
+
+    JitLuDecomposition(const JitLuDecomposition &) = delete;
+    JitLuDecomposition &operator=(const JitLuDecomposition &) = delete;
+    JitLuDecomposition(JitLuDecomposition &&);
+    JitLuDecomposition &operator=(JitLuDecomposition &&);
+
     /// @brief Create a JITed LU decomposer for a given sparse matrix structure
     /// @param compiler JIT compiler
     /// @param matrix Sparse matrix to create LU decomposer for

diff --git a/include/micm/solver/jit_lu_decomposition.inl b/include/micm/solver/jit_lu_decomposition.inl
@@ -3,6 +3,28 @@
 
 namespace micm
 {
+
+  template<std::size_t L>
+  inline JitLuDecomposition<L>::JitLuDecomposition(JitLuDecomposition &&other)
+      : LuDecomposition(std::move(other)),
+        compiler_(std::move(other.compiler_)),
+        decompose_function_resource_tracker_(std::move(other.decompose_function_resource_tracker_)),
+        decompose_function_(std::move(other.decompose_function_))
+  {
+    other.decompose_function_ = NULL;
+  }
+
+  template<std::size_t L>
+  inline JitLuDecomposition<L> &JitLuDecomposition<L>::operator=(JitLuDecomposition &&other)
+  {
+    LuDecomposition::operator=(std::move(other));
+    compiler_ = std::move(other.compiler_);
+    decompose_function_resource_tracker_ = std::move(other.decompose_function_resource_tracker_);
+    decompose_function_ = std::move(other.decompose_function_);
+    other.decompose_function_ = NULL;
+    return *this;
+  }
+
   template<std::size_t L>
   inline JitLuDecomposition<L>::JitLuDecomposition(
       std::shared_ptr<JitCompiler> compiler,
@@ -11,7 +33,10 @@ namespace micm
         compiler_(compiler)
   {
     decompose_function_ = NULL;
-    assert(matrix.size() <= L && "Jit LU Decomposition matrix size mismatch");
+    if (matrix.size() > L)
+    {
+      throw std::runtime_error("Invalid matrix for JitLuDecomposition. Check the the VectorMatrix template parameters.");
+    }
     GenerateDecomposeFunction();
   }
 

diff --git a/include/micm/solver/jit_rosenbrock.hpp b/include/micm/solver/jit_rosenbrock.hpp
@@ -19,20 +19,47 @@
 #include <ctime>
 #include <micm/jit/jit_compiler.hpp>
 #include <micm/jit/jit_function.hpp>
+#include <micm/solver/jit_linear_solver.hpp>
 #include <micm/solver/rosenbrock.hpp>
+#include <micm/util/sparse_matrix_vector_ordering.hpp>
+#include <micm/util/vector_matrix.hpp>
 
 namespace micm
 {
 
-  template<template<class> class MatrixPolicy = Matrix, template<class> class SparseMatrixPolicy = SparseMatrix>
-  class JitRosenbrockSolver : public RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy>
+  template<
+      template<class> class MatrixPolicy = VectorMatrix,
+      template<class> class SparseMatrixPolicy = VectorSparseMatrix,
+      class LinearSolverPolicy = JitLinearSolver<DEFAULT_VECTOR_SIZE, SparseMatrixPolicy>>
+  class JitRosenbrockSolver : public RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy>
   {
     std::shared_ptr<JitCompiler> compiler_;
     llvm::orc::ResourceTrackerSP function_resource_tracker_;
     using FuncPtr = void (*)(double*, const double);
     FuncPtr alpha_minus_jacobian_ = nullptr;
 
    public:
+    JitRosenbrockSolver(const JitRosenbrockSolver&) = delete;
+    JitRosenbrockSolver& operator=(const JitRosenbrockSolver&) = delete;
+    JitRosenbrockSolver(JitRosenbrockSolver&& other)
+        : RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy>(std::move(other)),
+          compiler_(std::move(other.compiler_)),
+          function_resource_tracker_(std::move(other.function_resource_tracker_)),
+          alpha_minus_jacobian_(std::move(other.alpha_minus_jacobian_))
+    {
+      other.alpha_minus_jacobian_ = NULL;
+    }
+
+    JitRosenbrockSolver& operator=(JitRosenbrockSolver&& other)
+    {
+      RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy>::operator=(std::move(other));
+      compiler_ = std::move(other.compiler_);
+      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
@@ -41,9 +68,20 @@ namespace micm
         const System& system,
         const std::vector<Process>& processes,
         const RosenbrockSolverParameters& parameters)
-        : RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy>(system, processes, parameters),
+        : RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy>(
+              system,
+              processes,
+              parameters,
+              [&](const SparseMatrixPolicy<double>& matrix, double initial_value) -> LinearSolverPolicy {
+                return LinearSolverPolicy{ compiler, matrix, initial_value };
+              }),
           compiler_(compiler)
     {
+      MatrixPolicy<double> temp{};
+      if (temp.GroupVectorSize() != parameters.number_of_grid_cells_)
+      {
+        throw std::runtime_error("Number of grid cells for JitRosenbrockSolver must match template parameter.");
+      }
       this->GenerateAlphaMinusJacobian();
     }