From 5f42f05340b50fb659da05a8532b34df3a6684e7 Mon Sep 17 00:00:00 2001
From: Matt Dawson <mattdawson@ucar.edu>
Date: Thu, 19 Dec 2024 12:18:51 -0800
Subject: [PATCH] update jit lu decomp

---
 .../solver/jit_lu_decomposition_doolittle.hpp |  7 ++++++-
 .../solver/jit_lu_decomposition_doolittle.inl | 20 +++++++++++++++----
 .../solver/lu_decomposition_doolittle.hpp     | 10 +++++-----
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/include/micm/jit/solver/jit_lu_decomposition_doolittle.hpp b/include/micm/jit/solver/jit_lu_decomposition_doolittle.hpp
index b6f445d99..3045b7789 100644
--- a/include/micm/jit/solver/jit_lu_decomposition_doolittle.hpp
+++ b/include/micm/jit/solver/jit_lu_decomposition_doolittle.hpp
@@ -35,12 +35,17 @@ namespace micm
 
     ~JitLuDecompositionDoolittle();
 
+    /// @brief Create an LU decomposition algorithm for a given sparse matrix policy
+    /// @param matrix Sparse matrix
+    template<class SparseMatrixPolicy, class LMatrixPolicy = SparseMatrixPolicy, class UMatrixPolicy = SparseMatrixPolicy>
+    requires(SparseMatrixConcept<SparseMatrixPolicy>) static JitLuDecompositionDoolittle<L> Create(const SparseMatrixPolicy& matrix);
+
     /// @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;
+    void Decompose(const SparseMatrixPolicy &A, auto& lower, auto& upper) const;
 
    private:
     /// @brief Generates a function to perform the LU decomposition for a specific matrix sparsity structure
diff --git a/include/micm/jit/solver/jit_lu_decomposition_doolittle.inl b/include/micm/jit/solver/jit_lu_decomposition_doolittle.inl
index 542d60848..95cc27590 100644
--- a/include/micm/jit/solver/jit_lu_decomposition_doolittle.inl
+++ b/include/micm/jit/solver/jit_lu_decomposition_doolittle.inl
@@ -26,9 +26,8 @@ namespace micm
   template<std::size_t L>
   inline JitLuDecompositionDoolittle<L>::JitLuDecompositionDoolittle(
       const SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix)
-      : LuDecompositionDoolittle(
-            LuDecompositionDoolittle::Create<SparseMatrix<double, SparseMatrixVectorOrdering<L>>>(matrix))
   {
+    using SparseMatrixPolicy = SparseMatrix<double, SparseMatrixVectorOrdering<L>>;
     decompose_function_ = NULL;
     if (matrix.NumberOfBlocks() > L)
     {
@@ -39,6 +38,7 @@ namespace micm
                         std::to_string(L);
       throw std::system_error(make_error_code(MicmJitErrc::InvalidMatrix), msg);
     }
+    Initialize<SparseMatrixPolicy, SparseMatrixPolicy, SparseMatrixPolicy>(matrix, typename SparseMatrixPolicy::value_type());
     GenerateDecomposeFunction();
   }
 
@@ -52,6 +52,18 @@ namespace micm
     }
   }
 
+  template<std::size_t L>
+  template<class SparseMatrixPolicy, class LMatrixPolicy, class UMatrixPolicy>
+  requires(SparseMatrixConcept<SparseMatrixPolicy>)
+  inline JitLuDecompositionDoolittle<L> JitLuDecompositionDoolittle<L>::Create(
+    const SparseMatrixPolicy& matrix)
+  {
+    static_assert(std::is_same_v<SparseMatrixPolicy, LMatrixPolicy>, "SparseMatrixPolicy must be the same as LMatrixPolicy for JIT LU decomposition");
+    static_assert(std::is_same_v<SparseMatrixPolicy, UMatrixPolicy>, "SparseMatrixPolicy must be the same as UMatrixPolicy for JIT LU decomposition");
+    JitLuDecompositionDoolittle<L> lu_decomp(matrix);
+    return lu_decomp;
+  }
+
   template<std::size_t L>
   void JitLuDecompositionDoolittle<L>::GenerateDecomposeFunction()
   {
@@ -210,8 +222,8 @@ namespace micm
   template<class SparseMatrixPolicy>
   void JitLuDecompositionDoolittle<L>::Decompose(
       const SparseMatrixPolicy &A,
-      SparseMatrixPolicy &lower,
-      SparseMatrixPolicy &upper) const
+      auto& lower,
+      auto& upper) const
   {
     decompose_function_(A.AsVector().data(), lower.AsVector().data(), upper.AsVector().data());
     for (size_t block = 0; block < A.NumberOfBlocks(); ++block)
diff --git a/include/micm/solver/lu_decomposition_doolittle.hpp b/include/micm/solver/lu_decomposition_doolittle.hpp
index d1fe311e3..ef4d98b45 100644
--- a/include/micm/solver/lu_decomposition_doolittle.hpp
+++ b/include/micm/solver/lu_decomposition_doolittle.hpp
@@ -89,20 +89,20 @@ namespace micm
 
     /// @brief Construct an LU decomposition algorithm for a given sparse matrix
     /// @param matrix Sparse matrix
-    template<class SparseMatrixPolicy, class LMatrixPolicy, class UMatrixPolicy>
+    template<class SparseMatrixPolicy, class LMatrixPolicy = SparseMatrixPolicy, class UMatrixPolicy = SparseMatrixPolicy>
     requires(SparseMatrixConcept<SparseMatrixPolicy>) LuDecompositionDoolittle(const SparseMatrixPolicy& matrix);
 
     ~LuDecompositionDoolittle() = default;
 
     /// @brief Create an LU decomposition algorithm for a given sparse matrix policy
     /// @param matrix Sparse matrix
-    template<class SparseMatrixPolicy, class LMatrixPolicy, class UMatrixPolicy>
+    template<class SparseMatrixPolicy, class LMatrixPolicy = SparseMatrixPolicy, class UMatrixPolicy = SparseMatrixPolicy>
     requires(SparseMatrixConcept<SparseMatrixPolicy>) static LuDecompositionDoolittle 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, class LMatrixPolicy, class UMatrixPolicy>
+    template<class SparseMatrixPolicy, class LMatrixPolicy = SparseMatrixPolicy, class UMatrixPolicy = SparseMatrixPolicy>
     requires(SparseMatrixConcept<SparseMatrixPolicy>) static std::pair<LMatrixPolicy, UMatrixPolicy> GetLUMatrices(
         const SparseMatrixPolicy& A,
         typename SparseMatrixPolicy::value_type initial_value);
@@ -122,10 +122,10 @@ namespace micm
         auto& L,
         auto& U) const;
 
-   private:
+   protected:
     /// @brief Initialize arrays for the LU decomposition
     /// @param A Sparse matrix to decompose
-    template<class SparseMatrixPolicy, class LMatrixPolicy, class UMatrixPolicy>
+    template<class SparseMatrixPolicy, class LMatrixPolicy = SparseMatrixPolicy, class UMatrixPolicy = SparseMatrixPolicy>
     requires(SparseMatrixConcept<SparseMatrixPolicy>) void Initialize(const SparseMatrixPolicy& matrix, auto initial_value);
   };