Fix SuperLU_DIST memory leak (#4258)

* Fix memory bug in options.

* Add check on incoming matrix.

* More defensive check for passing new operator A.

* Add some state after reviewing against PETSc - unavoidable.

* And add to Python interface

* Improve doctoring

Author: jhale

cpp/dolfinx/la/superlu_dist.cpp

@@ -103,6 +103,7 @@ std::vector<int_t> row_indices(const auto& A)
   // Write the per-scalar-row entry counts into `flattened_rowptr[1:]`, with
   // each block-row contributing `bs[0]` copies.
   std::vector<int_t> flattened_rowptr(m_loc_block * bs[0] + 1);
+  flattened_rowptr[0] = A_rowptr[0] * bs[1];
   for (std::int32_t i = 0; i < m_loc_block; ++i)
   {
     int_t delta = (A_rowptr[i + 1] - A_rowptr[i]) * bs[1];
@@ -462,8 +463,7 @@ template <typename T>
 void SuperLUDistSolver<T>::set_options(
     SuperLUDistStructs::superlu_dist_options_t options)
 {
-  _options = std::make_unique<SuperLUDistStructs::superlu_dist_options_t>(
-      std::move(options));
+  *_options = options;
 }
 //----------------------------------------------------------------------------
 template <typename T>
@@ -544,13 +544,69 @@ void SuperLUDistSolver<T>::set_option(std::string name, std::string value)
 }
 //----------------------------------------------------------------------------
 template <typename T>
-void SuperLUDistSolver<T>::set_A(std::shared_ptr<const SuperLUDistMatrix<T>> A)
+SuperLUDistSolver<T>::~SuperLUDistSolver()
 {
+  if (_factored)
+  {
+    int_t n = _superlu_matA->supermatrix()->ncol;
+    if constexpr (std::is_same_v<T, double>)
+      dDestroy_LU(n, _gridinfo.get(), _lustruct.get());
+    else if constexpr (std::is_same_v<T, float>)
+      sDestroy_LU(n, _gridinfo.get(), _lustruct.get());
+    else if constexpr (std::is_same_v<T, std::complex<double>>)
+      zDestroy_LU(n, _gridinfo.get(), _lustruct.get());
+    else
+      static_assert(always_false_v<T>, "Invalid scalar type");
+  }
+}
+//----------------------------------------------------------------------------
+template <typename T>
+void SuperLUDistSolver<T>::set_A(std::shared_ptr<const SuperLUDistMatrix<T>> A,
+                                 std::string fact)
+{
+  if (A->supermatrix()->nrow != _superlu_matA->supermatrix()->nrow
+      or A->supermatrix()->ncol != _superlu_matA->supermatrix()->ncol)
+  {
+    throw std::runtime_error(
+        "New matrix A has different size to the matrix used to construct the "
+        "solver.");
+  }
   _superlu_matA = A;
+
+  // See pddistribute in SuperLU_DIST: on Fact=DOFACT or Fact=SamePattern
+  // pdgssvx overwrites LUstruct's internal pointers without freeing the
+  // previous arrays, so a prior factorisation must be released first.
+  // Fact=SamePattern_SameRowPerm reuses the existing arrays in place.
+  if (fact == "DOFACT" or fact == "SamePattern")
+  {
+    if (_factored)
+    {
+      int_t n = _superlu_matA->supermatrix()->ncol;
+      if constexpr (std::is_same_v<T, double>)
+        dDestroy_LU(n, _gridinfo.get(), _lustruct.get());
+      else if constexpr (std::is_same_v<T, float>)
+        sDestroy_LU(n, _gridinfo.get(), _lustruct.get());
+      else if constexpr (std::is_same_v<T, std::complex<double>>)
+        zDestroy_LU(n, _gridinfo.get(), _lustruct.get());
+      else
+        static_assert(always_false_v<T>, "Invalid scalar type");
+      _factored = false;
+    }
+    _options->Fact = (fact == "DOFACT") ? DOFACT : SamePattern;
+  }
+  else if (fact == "SamePattern_SameRowPerm")
+  {
+    _options->Fact = SamePattern_SameRowPerm;
+  }
+  else
+  {
+    throw std::runtime_error("set_A fact must be one of 'DOFACT', "
+                             "'SamePattern', 'SamePattern_SameRowPerm'");
+  }
 }
 //----------------------------------------------------------------------------
 template <typename T>
-int SuperLUDistSolver<T>::solve(const la::Vector<T>& b, la::Vector<T>& u) const
+int SuperLUDistSolver<T>::solve(const la::Vector<T>& b, la::Vector<T>& u)
 {
   common::Timer tsolve("SuperLU_DIST solve");
 
@@ -614,6 +670,10 @@ int SuperLUDistSolver<T>::solve(const la::Vector<T>& b, la::Vector<T>& u) const
   if (info != 0)
     spdlog::info("SuperLU_DIST p*gssvx() error: {}", info);
 
+  // pdgssvx allocates LUstruct internals during factorisation. Record this
+  // so the destructor / set_A can call Destroy_LU to release them.
+  _factored = true;
+
   PStatPrint(_options.get(), &stat, _gridinfo.get());
   PStatFree(&stat);
 

cpp/dolfinx/la/superlu_dist.h

@@ -204,6 +204,9 @@ class SuperLUDistSolver
   /// Copy assignment
   SuperLUDistSolver& operator=(const SuperLUDistSolver&) = delete;
 
+  /// @brief Destructor. Frees internal LU arrays before LUstructFree.
+  ~SuperLUDistSolver();
+
   /// @brief Set solver option name to value
   ///
   /// See SuperLU_DIST User's Guide for option names and values.
@@ -237,11 +240,14 @@ class SuperLUDistSolver
 
   /// @brief Set assembled left-hand side matrix A.
   ///
-  /// For advanced use with SuperLU_DIST option `Factor` allowing use of
-  /// previously computed permutations when solving with new matrix A.
+  /// New matrix must have the same size/parallel layout as the matrix
+  /// used to construct the solver.
   ///
   /// @param A Assembled left-hand side matrix.
-  void set_A(std::shared_ptr<const SuperLUDistMatrix<T>> A);
+  /// @param fact One of `"DOFACT"`, `"SamePattern"`,
+  ///   `"SamePattern_SameRowPerm"`. See the SuperLU_DIST documentation
+  ///   for the meaning of these values.
+  void set_A(std::shared_ptr<const SuperLUDistMatrix<T>> A, std::string fact);
 
   /// @brief Solve linear system Au = b.
   ///
@@ -255,7 +261,7 @@ class SuperLUDistSolver
   /// @note The values of `A` are modified in-place during the solve.
   /// @note To solve with successive right-hand sides the caller must
   ///   `solver.set_options("Factor", "FACTORED")` after the first solve.
-  int solve(const Vector<T>& b, Vector<T>& u) const;
+  int solve(const Vector<T>& b, Vector<T>& u);
 
 private:
   // Assembled left-hand side matrix
@@ -275,6 +281,10 @@ class SuperLUDistSolver
   // Pointer to 'typed' struct *SOLVEstruct
   std::unique_ptr<typename map_t<T>::SOLVEstruct_t, SolveStructDeleter>
       _solvestruct;
+
+  // True once pdgssvx has populated LUstruct with per-block-column arrays
+  // that must be released via Destroy_LU before LUstructFree.
+  bool _factored = false;
 };
 } // namespace dolfinx::la
 #endif

python/dolfinx/la/superlu_dist.py

@@ -119,16 +119,16 @@ def set_option(self, name: str, value: str):
         """
         self._cpp_object.set_option(name, value)
 
-    def set_A(self, A: SuperLUDistMatrix[_T]):
+    def set_A(self, A: SuperLUDistMatrix[_T], fact: str):
         """Set assembled left-hand side matrix.
 
-        For advanced use with SuperLU_DIST option `Factor` allowing use of
-        previously computed permutations when solving with new matrix A.
-
         Args:
             A: Assembled left-hand side matrix :math:`A`.
+            fact: One of ``"DOFACT"``, ``"SamePattern"``,
+                ``"SamePattern_SameRowPerm"``. See the SuperLU_DIST
+                documentation for the meaning of these values.
         """
-        self._cpp_object.set_A(A._cpp_object)
+        self._cpp_object.set_A(A._cpp_object, fact)
 
     def solve(self, b: dolfinx.la.Vector[_T], u: dolfinx.la.Vector[_T]) -> int:
         """Solve linear system :math:`Au = b`.

python/test/unit/la/test_superlu_dist.py

@@ -128,8 +128,7 @@ def solve_and_check(solver, b):
     A_3.scatter_reverse()
 
     A_superlu_3 = superlu_dist_matrix(A_3)
-    solver_2.set_A(A_superlu_3)
-    solver_2.set_option("Fact", "SamePattern")
+    solver_2.set_A(A_superlu_3, "SamePattern")
     uh_2 = solve_and_check(solver_2, b_2)
     check_error(u_ex, uh_2)