NCAR · K20shores · Jul 17, 2024 · Jul 11, 2024 · Jul 11, 2024 · Jul 12, 2024
diff --git a/include/micm/jit/solver/jit_linear_solver.hpp b/include/micm/jit/solver/jit_linear_solver.hpp
@@ -49,10 +49,6 @@ namespace micm
         SparseMatrixPolicy& upper_matrix,
         bool& is_singular);
 
-    /// @brief Decompose the matrix into upper and lower triangular matrices and general JIT functions
-    /// @param matrix Matrix that will be factored into lower and upper triangular matrices
-    void Factor(SparseMatrixPolicy& matrix, SparseMatrixPolicy& lower_matrix, SparseMatrixPolicy& upper_matrix);
-
     /// @brief Solve for x in Ax = b
     template<class MatrixPolicy>
     void Solve(MatrixPolicy& x, SparseMatrixPolicy& lower_matrix, SparseMatrixPolicy& upper_matrix);

diff --git a/include/micm/jit/solver/jit_linear_solver.inl b/include/micm/jit/solver/jit_linear_solver.inl
@@ -66,15 +66,6 @@ namespace micm
     LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::Factor(matrix, lower_matrix, upper_matrix, is_singular);
   }
 
-  template<std::size_t L, class SparseMatrixPolicy, class LuDecompositionPolicy>
-  inline void JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::Factor(
-      SparseMatrixPolicy &matrix,
-      SparseMatrixPolicy &lower_matrix,
-      SparseMatrixPolicy &upper_matrix)
-  {
-    LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::Factor(matrix, lower_matrix, upper_matrix);
-  }
-
   template<std::size_t L, class SparseMatrixPolicy, class LuDecompositionPolicy>
   template<class MatrixPolicy>
   inline void JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::Solve(

diff --git a/include/micm/jit/solver/jit_lu_decomposition.hpp b/include/micm/jit/solver/jit_lu_decomposition.hpp
@@ -18,7 +18,8 @@ namespace micm
   class JitLuDecomposition : public LuDecomposition
   {
     llvm::orc::ResourceTrackerSP decompose_function_resource_tracker_;
-    void (*decompose_function_)(const double *, double *, double *);
+    using FuncPtr = void (*)(const double *, double *, double *);
+    FuncPtr decompose_function_ = nullptr;
 
    public:
     JitLuDecomposition(){};
@@ -43,13 +44,6 @@ namespace micm
     void Decompose(const SparseMatrixPolicy &A, SparseMatrixPolicy &lower, SparseMatrixPolicy &upper, bool &is_singular)
         const;
 
-    /// @brief Create sparse L and U matrices for a given A matrix
-    /// @param A Sparse matrix that will be decomposed
-    /// @param lower The lower triangular matrix created by decomposition
-    /// @param upper The upper triangular matrix created by decomposition
-    template<class SparseMatrixPolicy>
-    void Decompose(const SparseMatrixPolicy &A, SparseMatrixPolicy &lower, SparseMatrixPolicy &upper) const;
-
    private:
     /// @brief Generates a function to perform the LU decomposition for a specific matrix sparsity structure
     void GenerateDecomposeFunction();

diff --git a/include/micm/jit/solver/jit_lu_decomposition.inl b/include/micm/jit/solver/jit_lu_decomposition.inl
@@ -58,7 +58,7 @@ namespace micm
                            .SetName(function_name)
                            .SetArguments({ { "A matrix", JitType::DoublePtr },
                                            { "lower matrix", JitType::DoublePtr },
-                                           { "upper matrix", JitType::DoublePtr } })
+                                           { "upper matrix", JitType::DoublePtr }})
                            .SetReturnType(JitType::Void);
     llvm::Type *double_type = func.GetType(JitType::Double);
     llvm::Value *zero = llvm::ConstantInt::get(*(func.context_), llvm::APInt(64, 0));
@@ -180,26 +180,27 @@ namespace micm
     func.builder_->CreateRetVoid();
 
     auto target = func.Generate();
-    decompose_function_ = (void (*)(const double *, double *, double *))(intptr_t)target.second;
+    decompose_function_ = (FuncPtr)(intptr_t)target.second;
     decompose_function_resource_tracker_ = target.first;
   }
 
   template<std::size_t L>
   template<class SparseMatrixPolicy>
-  void JitLuDecomposition<L>::Decompose(
-      const SparseMatrixPolicy &A,
-      SparseMatrixPolicy &lower,
-      SparseMatrixPolicy &upper,
-      bool &is_singular) const
-  {
-    LuDecomposition::Decompose<SparseMatrixPolicy>(A, lower, upper, is_singular);
-  }
-
-  template<std::size_t L>
-  template<class SparseMatrixPolicy>
-  void JitLuDecomposition<L>::Decompose(const SparseMatrixPolicy &A, SparseMatrixPolicy &lower, SparseMatrixPolicy &upper)
-      const
+  void JitLuDecomposition<L>::Decompose(const SparseMatrixPolicy &A, SparseMatrixPolicy &lower, SparseMatrixPolicy &upper, bool& is_singular) const
   {
+    is_singular = false;
     decompose_function_(A.AsVector().data(), lower.AsVector().data(), upper.AsVector().data());
+    for(size_t block = 0; block < A.NumberOfBlocks(); ++block)
+    {
+      auto diagonals = upper.DiagonalIndices(block);
+      for(const auto& diagonal : diagonals)
+      {
+        if(upper.AsVector()[diagonal] == 0)
+        {
+          is_singular = true;
+          return;
+        }
+      }
+    }
   }
 }  // namespace micm
diff --git a/include/micm/solver/linear_solver.hpp b/include/micm/solver/linear_solver.hpp
@@ -76,9 +76,6 @@ namespace micm
 
     virtual ~LinearSolver() = default;
 
-    /// @brief Decompose the matrix into upper and lower triangular matrices
-    void Factor(const SparseMatrixPolicy& matrix, SparseMatrixPolicy& lower_matrix, SparseMatrixPolicy& upper_matrix) const;
-
     /// @brief Decompose the matrix into upper and lower triangular matrices
     /// @param matrix Matrix to decompose into lower and upper triangular matrices
     /// @param is_singular Flag that is set to true if matrix is singular; false otherwise

diff --git a/include/micm/solver/linear_solver.inl b/include/micm/solver/linear_solver.inl
@@ -107,17 +107,6 @@ namespace micm
     }
   };
 
-  template<class SparseMatrixPolicy, class LuDecompositionPolicy>
-  inline void LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::Factor(
-      const SparseMatrixPolicy& matrix,
-      SparseMatrixPolicy& lower_matrix,
-      SparseMatrixPolicy& upper_matrix) const
-  {
-    MICM_PROFILE_FUNCTION();
-
-    lu_decomp_.template Decompose<SparseMatrixPolicy>(matrix, lower_matrix, upper_matrix);
-  }
-
   template<class SparseMatrixPolicy, class LuDecompositionPolicy>
   inline void LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::Factor(
       const SparseMatrixPolicy& matrix,

diff --git a/include/micm/solver/lu_decomposition.hpp b/include/micm/solver/lu_decomposition.hpp
@@ -111,16 +111,6 @@ namespace micm
         const SparseMatrixPolicy& A,
         typename SparseMatrixPolicy::value_type initial_value);
 
-    /// @brief Perform an LU decomposition on a given A matrix
-    /// @param A Sparse matrix to decompose
-    /// @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
     /// @param A Sparse matrix to decompose
     /// @param L The lower triangular matrix created by decomposition

diff --git a/include/micm/solver/lu_decomposition.inl b/include/micm/solver/lu_decomposition.inl
@@ -99,6 +99,7 @@ namespace micm
       }
       niLU_.push_back(iLU);
     }
+    uii_.push_back(LU.second.VectorIndex(0, n - 1, n - 1));
   }
 
   template<class SparseMatrixPolicy>
@@ -163,16 +164,6 @@ namespace micm
     return LU;
   }
 
-  template<class SparseMatrixPolicy>
-  requires(SparseMatrixConcept<SparseMatrixPolicy>) inline void LuDecomposition::Decompose(
-      const SparseMatrixPolicy& A,
-      SparseMatrixPolicy& L,
-      SparseMatrixPolicy& U) const
-  {
-    bool is_singular = false;
-    Decompose<SparseMatrixPolicy>(A, L, U, is_singular);
-  }
-
   template<class SparseMatrixPolicy>
   requires(!VectorizableSparse<SparseMatrixPolicy>) inline void LuDecomposition::Decompose(
       const SparseMatrixPolicy& A,
@@ -181,6 +172,7 @@ namespace micm
       bool& is_singular) const
   {
     MICM_PROFILE_FUNCTION();
+    is_singular = false;
 
     // Loop over blocks
     for (std::size_t i_block = 0; i_block < A.NumberOfBlocks(); ++i_block)
@@ -197,7 +189,6 @@ namespace micm
       auto lki_nkj = lki_nkj_.begin();
       auto lkj_uji = lkj_uji_.begin();
       auto uii = uii_.begin();
-      is_singular = false;
       for (auto& inLU : niLU_)
       {
         // Upper trianglur matrix
@@ -223,16 +214,21 @@ namespace micm
             L_vector[lki_nkj->first] -= L_vector[lkj_uji->first] * U_vector[lkj_uji->second];
             ++lkj_uji;
           }
+
           if (U_vector[*uii] == 0.0)
           {
             is_singular = true;
-            return;
           }
           L_vector[lki_nkj->first] /= U_vector[*uii];
           ++lki_nkj;
           ++uii;
         }
       }
+      // check the bottom right corner of the matrix
+      if (U_vector[*uii] == 0.0)
+      {
+        is_singular = true;
+      }
     }
   }
 
@@ -250,6 +246,7 @@ namespace micm
     std::size_t A_GroupSizeOfFlatBlockSize = A.GroupSize(A.FlatBlockSize());
     std::size_t L_GroupSizeOfFlatBlockSize = L.GroupSize(L.FlatBlockSize());
     std::size_t U_GroupSizeOfFlatBlockSize = U.GroupSize(U.FlatBlockSize());
+    is_singular = false;
 
     // Loop over groups of blocks
     for (std::size_t i_group = 0; i_group < A.NumberOfGroups(A_BlockSize); ++i_group)
@@ -266,7 +263,6 @@ namespace micm
       auto lki_nkj = lki_nkj_.begin();
       auto lkj_uji = lkj_uji_.begin();
       auto uii = uii_.begin();
-      is_singular = false;
       const std::size_t n_cells = std::min(A_GroupVectorSize, A_BlockSize - i_group * A_GroupVectorSize);
       for (auto& inLU : niLU_)
       {
@@ -316,15 +312,26 @@ namespace micm
             if (U_vector[uii_deref + i_cell] == 0.0)
             {
               is_singular = true;
-              return;
             }
             L_vector[lki_nkj_first + i_cell] /= U_vector[uii_deref + i_cell];
           }
           ++lki_nkj;
           ++uii;
         }
       }
+      std::size_t uii_deref = *uii;
+      if (n_cells != A_GroupVectorSize) {
+        // check the bottom right corner of the matrix
+        for (std::size_t i_cell = 0; i_cell < n_cells; ++i_cell)
+        {
+          if (U_vector[uii_deref + i_cell] == 0.0)
+          {
+            is_singular = true;
+            break;
+          }
+        }
+      }
     }
   }
 
-}  // namespace micm
+}  // namespace micm
diff --git a/include/micm/solver/rosenbrock.inl b/include/micm/solver/rosenbrock.inl
@@ -240,26 +240,21 @@ namespace micm
   {
     MICM_PROFILE_FUNCTION();
 
-    auto jacobian = state.jacobian_;
     uint64_t n_consecutive = 0;
     singular = false;
     while (true)
     {
+      auto jacobian = state.jacobian_;
       double alpha = 1 / (H * gamma);
       static_cast<Derived*>(this)->AlphaMinusJacobian(jacobian, alpha);
-      if (parameters_.check_singularity_)
-      {
-        linear_solver_.Factor(jacobian, state.lower_matrix_, state.upper_matrix_, singular);
-      }
-      else
-      {
-        singular = false;
-        linear_solver_.Factor(jacobian, state.lower_matrix_, state.upper_matrix_);
-      }
-      singular = false;  // TODO This should be evaluated in some way
+      linear_solver_.Factor(jacobian, state.lower_matrix_, state.upper_matrix_, singular);
       stats.decompositions_ += 1;
-      if (!singular)
+
+      // if we are checking for singularity and the matrix is not singular, we can break the loop
+      // if we are not checking for singularity, we always break the loop
+      if (!singular || !parameters_.check_singularity_)
         break;
+
       stats.singular_ += 1;
       if (++n_consecutive > 5)
         break;

diff --git a/include/micm/solver/state.inl b/include/micm/solver/state.inl
@@ -92,7 +92,7 @@ namespace micm
     jacobian_ = BuildJacobian<SparseMatrixPolicy>(
         parameters.nonzero_jacobian_elements_, parameters.number_of_grid_cells_, state_size_);
 
-    auto lu = LuDecomposition::GetLUMatrices<SparseMatrixPolicy>(jacobian_, 1.0e-30);
+    auto lu = LuDecomposition::GetLUMatrices<SparseMatrixPolicy>(jacobian_, 0);
     auto lower_matrix = std::move(lu.first);
     auto upper_matrix = std::move(lu.second);
     lower_matrix_ = lower_matrix;

diff --git a/src/solver/lu_decomposition.cu b/src/solver/lu_decomposition.cu
@@ -102,6 +102,11 @@ namespace micm
             ++uii_offset;
           }
         }
+        // check the bottom right corner of the matrix
+        if (d_U[d_uii[uii_offset] + tid] == 0.0)
+        {
+          *d_is_singular = true;
+        }
       }
     }  // end of CUDA kernel
 

diff --git a/test/integration/test_analytical_rosenbrock.cpp b/test/integration/test_analytical_rosenbrock.cpp
@@ -285,7 +285,7 @@ TEST(AnalyticalExamples, SurfaceRxn)
   test_analytical_surface_rxn(rosenbrock_6stage_da, 1e-7);
 }
 
-TEST(AnalyticalExamples, HIRESConfig)
+TEST(AnalyticalExamples, HIRES)
 {
   auto rosenbrock_solver = [](auto params)
   {

diff --git a/test/unit/cuda/solver/test_cuda_linear_solver.cpp b/test/unit/cuda/solver/test_cuda_linear_solver.cpp
@@ -75,7 +75,8 @@ std::vector<double> linearSolverGenerator(std::size_t number_of_blocks)
   CopyToDeviceSparse<SparseMatrixPolicy>(lower_matrix);
   CopyToDeviceSparse<SparseMatrixPolicy>(upper_matrix);
 
-  solver.Factor(A, lower_matrix, upper_matrix);
+  bool singular{ false };
+  solver.Factor(A, lower_matrix, upper_matrix, singular);
   solver.template Solve<MatrixPolicy>(x, lower_matrix, upper_matrix);
 
   // Only copy the data to the host when it is a CudaMatrix

diff --git a/test/unit/cuda/solver/test_cuda_lu_decomposition.cpp b/test/unit/cuda/solver/test_cuda_lu_decomposition.cpp
@@ -50,7 +50,8 @@ void testCudaRandomMatrix(size_t n_grids)
 
   micm::LuDecomposition cpu_lud = micm::LuDecomposition::Create<CPUSparseMatrixPolicy>(cpu_A);
   auto cpu_LU = micm::LuDecomposition::GetLUMatrices<CPUSparseMatrixPolicy>(cpu_A, 1.0e-30);
-  cpu_lud.Decompose<CPUSparseMatrixPolicy>(cpu_A, cpu_LU.first, cpu_LU.second);
+  bool singular{ false };
+  cpu_lud.Decompose<CPUSparseMatrixPolicy>(cpu_A, cpu_LU.first, cpu_LU.second, singular);
 
   // checking GPU result again CPU
   size_t L_size = cpu_LU.first.AsVector().size();