Wrap SuperLU_DIST solver options (#4072)

* Moving towards options class + wrap. Looks like a bit of a pain.

* Simplify

* Add safe casting to int_t on boundary of DOLFINx/SuperLU_DIST

* Switch to comparing on max limit, tidy up whitespace.

* Revert

* Work in progress.

* Add option setter for some options

* Add a basic setter so that C++ users can just use native SuperLU interface.

* Switch to YES and NO, following SuperLU_DIST

* Small tweaks.

* Tweaks.

* Remove verbose flag

* Tests run

* Final tweaks and comments

* Turn off output

* Fix.

* Fix

* Fix

* clang-format

* Add comment on options.

---------

Co-authored-by: Chris Richardson <chris@bpi.cam.ac.uk>

Author: jhale

cpp/dolfinx/la/superlu_dist.cpp

@@ -14,8 +14,10 @@ extern "C"
 #include <superlu_zdefs.h>
 }
 #include <algorithm>
+#include <dolfinx/common/Timer.h>
 #include <dolfinx/la/MatrixCSR.h>
 #include <dolfinx/la/Vector.h>
+#include <ranges>
 #include <stdexcept>
 #include <vector>
 
@@ -46,6 +48,25 @@ namespace
 template <typename...>
 constexpr bool dependent_false_v = false;
 
+template <typename V, typename W>
+void option_setter(std::string_view option_name, W& option,
+                   const std::initializer_list<V> values,
+                   std::initializer_list<std::string_view> value_names,
+                   std::string_view value_in)
+{
+  // TODO: Can be done nicely with std::views::zip in C++23.
+  for (auto i : std::views::iota(std::size_t{0}, value_names.size()))
+  {
+    if (value_in == *(value_names.begin() + i))
+    {
+      option = *(values.begin() + i);
+      spdlog::info("Set {} to {}", option_name, value_in);
+      return;
+    }
+  }
+  throw std::runtime_error("Unsupported value");
+}
+
 std::vector<int_t> col_indices(const auto& A)
 {
   // Local number of non-zeros
@@ -135,15 +156,13 @@ create_supermatrix(const auto& A, auto& rowptr, auto& cols)
 } // namespace
 //----------------------------------------------------------------------------
 template <typename T>
-SuperLUDistMatrix<T>::SuperLUDistMatrix(std::shared_ptr<const MatrixCSR<T>> A,
-                                        bool verbose)
+SuperLUDistMatrix<T>::SuperLUDistMatrix(std::shared_ptr<const MatrixCSR<T>> A)
     : _matA(std::move(A)),
       _cols(
           std::make_unique<SuperLUDistStructs::vec_int_t>(col_indices(*_matA))),
       _rowptr(
           std::make_unique<SuperLUDistStructs::vec_int_t>(row_indices(*_matA))),
-      _supermatrix(create_supermatrix<T>(*_matA, *_rowptr, *_cols)),
-      _verbose(verbose)
+      _supermatrix(create_supermatrix<T>(*_matA, *_rowptr, *_cols))
 {
 }
 //----------------------------------------------------------------------------
@@ -170,18 +189,32 @@ struct dolfinx::la::SuperLUDistStructs::gridinfo_t : public ::gridinfo_t
 {
 };
 //----------------------------------------------------------------------------
+struct dolfinx::la::SuperLUDistStructs::superlu_dist_options_t
+    : public ::superlu_dist_options_t
+{
+};
+//----------------------------------------------------------------------------
 void GridInfoDeleter::operator()(
     SuperLUDistStructs::gridinfo_t* gridinfo) const noexcept
 {
   superlu_gridexit(gridinfo);
   delete gridinfo;
-}
+};
 
 //----------------------------------------------------------------------------
 template <typename T>
 SuperLUDistSolver<T>::SuperLUDistSolver(
-    std::shared_ptr<const SuperLUDistMatrix<T>> A, bool verbose)
+    std::shared_ptr<const SuperLUDistMatrix<T>> A)
     : _superlu_matA(std::move(A)),
+      _options(
+          []
+          {
+            auto o = std::make_unique<
+                SuperLUDistStructs::superlu_dist_options_t>();
+            set_default_options_dist(o.get());
+            o->PrintStat = NO;
+            return o;
+          }()),
       _gridinfo(
           [comm = _superlu_matA->matA().comm()]
           {
@@ -191,28 +224,105 @@ SuperLUDistSolver<T>::SuperLUDistSolver(
                 new SuperLUDistStructs::gridinfo_t, GridInfoDeleter{});
             superlu_gridinit(comm, nprow, npcol, p.get());
             return p;
-          }()),
-      _verbose(verbose)
+          }())
+{
+}
+
+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));
+}
+
+template <typename T>
+void SuperLUDistSolver<T>::set_option(std::string name, std::string value)
 {
+  spdlog::info("Attempting to set option {} to {}", name, value);
+  const std::map<std::string, std::reference_wrapper<yes_no_t>> map_bool
+      = {{"Equil", _options->Equil},
+         {"DiagInv", _options->DiagInv},
+         {"SymmetricMode", _options->SymmetricMode},
+         {"PivotGrowth", _options->PivotGrowth},
+         {"ConditionNumber", _options->ConditionNumber},
+         {"ReplaceTinyPivot", _options->ReplaceTinyPivot},
+         {"SolveInitialized", _options->SolveInitialized},
+         {"RefineInitialized", _options->RefineInitialized},
+         {"PrintStat", _options->PrintStat},
+         {"lookahead_etree", _options->lookahead_etree},
+         {"SymPattern", _options->SymPattern},
+         {"Use_TensorCore", _options->Use_TensorCore},
+         {"Algo3d", _options->Algo3d}};
+
+  // Search in map_bool first
+  auto it = map_bool.find(name);
+  if (it != map_bool.end())
+  {
+    if (value == "YES")
+    {
+      spdlog::info("Set {} to YES", name);
+      it->second.get() = YES;
+    }
+    else if (value == "NO")
+    {
+      spdlog::info("Set {} to NO", name);
+      it->second.get() = NO;
+    }
+    else
+    {
+      throw std::runtime_error("Boolean values must be string 'YES' or 'NO'");
+    }
+  }
+
+  // Search some enum types
+  if (name == "Fact")
+  {
+    option_setter(
+        name, _options->Fact,
+        {DOFACT, SamePattern, SamePattern_SameRowPerm, FACTORED},
+        {"DOFACT", "SamePattern", "SamePattern_SameRowPerm", "FACTORED"},
+        value);
+  }
+  else if (name == "Trans")
+  {
+    option_setter(name, _options->Trans, {NOTRANS, TRANS, CONJ},
+                  {"NOTRANS", "TRANS", "CONJ"}, value);
+  }
+  else if (name == "ColPerm")
+  {
+    option_setter(name, _options->ColPerm,
+                  {NATURAL, MMD_ATA, MMD_AT_PLUS_A, COLAMD, METIS_AT_PLUS_A,
+                   PARMETIS, METIS_ATA, ZOLTAN, MY_PERMC},
+                  {"NATURAL", "MMD_ATA", "MMD_AT_PLUS_A", "COLAMD",
+                   "METIS_AT_PLUS_A", "PARMETIS", "METIS_ATA", "ZOLTAN",
+                   "MY_PERMC"},
+                  value);
+  }
+  else if (name == "RowPerm")
+  {
+    option_setter(name, _options->RowPerm,
+                  {NOROWPERM, LargeDiag_MC64, LargeDiag_HWPM, MY_PERMR},
+                  {"NOROWPERM", "LargeDiag_MC64", "LargeDiag_HWPM", "MY_PERMR"},
+                  value);
+  }
+  else
+  {
+    std::runtime_error("Unsupported name");
+  }
 }
 //----------------------------------------------------------------------------
 template <typename T>
 int SuperLUDistSolver<T>::solve(const la::Vector<T>& b, la::Vector<T>& u) const
 {
+  common::Timer tsolve("SuperLU Solve");
   int_t m = _superlu_matA->supermatrix()->nrow;
   int_t m_loc = ((NRformat_loc*)(_superlu_matA->supermatrix()->Store))->m_loc;
 
   // RHS
   int_t ldb = m_loc;
   int_t nrhs = 1;
 
-  superlu_dist_options_t options;
-  set_default_options_dist(&options);
-  options.DiagInv = YES;
-  options.ReplaceTinyPivot = YES;
-  if (!_verbose)
-    options.PrintStat = NO;
-
   int info = 0;
   SuperLUStat_t stat;
   PStatInit(&stat);
@@ -232,12 +342,12 @@ int SuperLUDistSolver<T>::solve(const la::Vector<T>& b, la::Vector<T>& u) const
     dSOLVEstruct_t SOLVEstruct;
 
     spdlog::info("Call SuperLU_DIST pdgssvx()");
-    pdgssvx(&options, _superlu_matA->supermatrix(), &ScalePermstruct,
+    pdgssvx(_options.get(), _superlu_matA->supermatrix(), &ScalePermstruct,
             u.array().data(), ldb, nrhs, _gridinfo.get(), &LUstruct,
             &SOLVEstruct, berr.data(), &stat, &info);
 
     spdlog::info("Finalize solve");
-    dSolveFinalize(&options, &SOLVEstruct);
+    dSolveFinalize(_options.get(), &SOLVEstruct);
     dScalePermstructFree(&ScalePermstruct);
     dLUstructFree(&LUstruct);
   }
@@ -251,12 +361,12 @@ int SuperLUDistSolver<T>::solve(const la::Vector<T>& b, la::Vector<T>& u) const
     sSOLVEstruct_t SOLVEstruct;
 
     spdlog::info("Call SuperLU_DIST psgssvx()");
-    psgssvx(&options, _superlu_matA->supermatrix(), &ScalePermstruct,
+    psgssvx(_options.get(), _superlu_matA->supermatrix(), &ScalePermstruct,
             u.array().data(), ldb, nrhs, _gridinfo.get(), &LUstruct,
             &SOLVEstruct, berr.data(), &stat, &info);
 
     spdlog::info("Finalize solve");
-    sSolveFinalize(&options, &SOLVEstruct);
+    sSolveFinalize(_options.get(), &SOLVEstruct);
     sScalePermstructFree(&ScalePermstruct);
     sLUstructFree(&LUstruct);
   }
@@ -270,13 +380,13 @@ int SuperLUDistSolver<T>::solve(const la::Vector<T>& b, la::Vector<T>& u) const
     zSOLVEstruct_t SOLVEstruct;
 
     spdlog::info("Call SuperLU_DIST pzgssvx()");
-    pzgssvx(&options, _superlu_matA->supermatrix(), &ScalePermstruct,
+    pzgssvx(_options.get(), _superlu_matA->supermatrix(), &ScalePermstruct,
             reinterpret_cast<doublecomplex*>(u.array().data()), ldb, nrhs,
             _gridinfo.get(), &LUstruct, &SOLVEstruct, berr.data(), &stat,
             &info);
 
     spdlog::info("Finalize solve");
-    zSolveFinalize(&options, &SOLVEstruct);
+    zSolveFinalize(_options.get(), &SOLVEstruct);
     zScalePermstructFree(&ScalePermstruct);
     zLUstructFree(&LUstruct);
   }
@@ -287,8 +397,7 @@ 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);
 
-  if (_verbose)
-    PStatPrint(&options, &stat, _gridinfo.get());
+  PStatPrint(_options.get(), &stat, _gridinfo.get());
   PStatFree(&stat);
 
   return info;

cpp/dolfinx/la/superlu_dist.h

@@ -11,6 +11,7 @@
 #include <dolfinx/la/MatrixCSR.h>
 #include <dolfinx/la/Vector.h>
 #include <memory>
+#include <string>
 
 namespace dolfinx::la
 {
@@ -19,8 +20,9 @@ class SuperLUDistStructs
 {
 public:
   struct SuperMatrix;
-  struct gridinfo_t;
   struct vec_int_t;
+  struct gridinfo_t;
+  struct superlu_dist_options_t;
 };
 
 /// Call library cleanup and delete pointer. For use with
@@ -39,21 +41,17 @@ class SuperLUDistMatrix
 public:
   /// @brief Create SuperLU_DIST matrix operator.
   ///
-  /// Handles RAII-type memory management of underlying C objects.
-  ///
   /// @tparam T Scalar type.
   /// @param A Matrix.
-  /// @param verbose Verbose output.
-  SuperLUDistMatrix(std::shared_ptr<const MatrixCSR<T>> A,
-                    bool verbose = false);
+  SuperLUDistMatrix(std::shared_ptr<const MatrixCSR<T>> A);
 
   /// Copy constructor
   SuperLUDistMatrix(const SuperLUDistMatrix&) = delete;
 
   /// Copy assignment
   SuperLUDistMatrix& operator=(const SuperLUDistMatrix&) = delete;
 
-  /// Get non-const pointer to SuperLU_DIST SuperMatrix.
+  /// Get pointer to SuperLU_DIST SuperMatrix (non-const).
   SuperLUDistStructs::SuperMatrix* supermatrix() const;
 
   /// Get MatrixCSR (const).
@@ -70,9 +68,6 @@ class SuperLUDistMatrix
   // Pointer to native SuperMatrix
   std::unique_ptr<SuperLUDistStructs::SuperMatrix, SuperMatrixDeleter>
       _supermatrix;
-
-  // Flag for diagnostic output
-  bool _verbose;
 };
 
 /// Call library cleanup and delete pointer. For use with
@@ -91,38 +86,71 @@ class SuperLUDistSolver
 public:
   /// @brief Create solver for a SuperLU_DIST matrix operator.
   ///
-  /// Solves Au = b using SuperLU_DIST.
+  /// Solves linear system Au = b via LU decomposition.
+  ///
+  /// The SuperLU_DIST solver has options set to upstream defaults,
+  /// except PrintStat (verbose solver output) set to NO.
   ///
   /// @tparam T Scalar type.
-  /// @param A Matrix to solve for.
-  /// @param verbose Verbose output.
-  SuperLUDistSolver(std::shared_ptr<const SuperLUDistMatrix<T>> A,
-                    bool verbose = false);
+  /// @param A Assembled left-hand side matrix.
+  SuperLUDistSolver(std::shared_ptr<const SuperLUDistMatrix<T>> A);
 
   /// Copy constructor
   SuperLUDistSolver(const SuperLUDistSolver&) = delete;
 
   /// Copy assignment
   SuperLUDistSolver& operator=(const SuperLUDistSolver&) = delete;
 
+  /// @brief Set solver option (name, value)
+  ///
+  /// See SuperLU_DIST User's Guide for option names and values.
+  ///
+  /// @param name Option name.
+  /// @param value Option value.
+  void set_option(std::string name, std::string value);
+
+  /// @brief Set all solver options (native struct)
+  ///
+  /// See SuperLU_DIST User's Guide for option names and values.
+  ///
+  /// Callers must complete the forward declared struct, e.g.:
+  ///
+  /// ```cpp
+  /// #include <superlu_defs.h>
+  /// struct dolfinx::la::SuperLUDistStructs::superlu_dist_options_t
+  ///  : public ::superlu_dist_options_t
+  /// {
+  /// };
+  ///
+  /// SuperLUDistStructs::superlu_dist_options_t options;
+  /// set_default_options_dist(&options);
+  /// options.PrintStat = YES;
+  /// // Setup SuperLUDistMatrix and SuperLUDistSolver
+  /// solver.set_options(options);
+  /// ```
+  ///
+  /// @param options SuperLU_DIST option struct.
+  void set_options(SuperLUDistStructs::superlu_dist_options_t options);
+
   /// @brief Solve linear system Au = b.
   ///
   /// @param b Right-hand side vector.
-  /// @param u Solution vector.
-  /// @returns SuperLU_DIST info flag.
-  /// @note Vectors must have size and parallel layout that is
-  /// compatible with `A`.
+  /// @param u Solution vector, overwritten during solve.
+  /// @returns SuperLU_DIST info integer.
+  /// @note The caller must check the return code for success `(== 0)`.
+  /// @note The caller must `u.scatter_forward()` after the solve.
+  /// @note Vectors must have size and parallel layout compatible with `A`.
   int solve(const Vector<T>& b, Vector<T>& u) const;
 
 private:
-  // Wrapped SuperLU SuperMatrix
+  // Assembled left-hand side matrix
   std::shared_ptr<const SuperLUDistMatrix<T>> _superlu_matA;
 
+  // Pointer to struct superlu_dist_options_t
+  std::unique_ptr<SuperLUDistStructs::superlu_dist_options_t> _options;
+
   // Pointer to struct gridinfo_t
   std::unique_ptr<SuperLUDistStructs::gridinfo_t, GridInfoDeleter> _gridinfo;
-
-  // Flag for diagnostic output
-  bool _verbose;
 };
 } // namespace dolfinx::la
 #endif