Fix HyKKT solver reuse and add export support

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## HyKKT Release changes
 
+- Exported HyKKT libraries and fixed solver reuse by refreshing numerical data, reusing allocations, resizing GPU transpose workspaces, and handling zero residuals.
 - Added classes and tests for permutation, Ruiz scaling, Cholesky factorization, Schur complement conjugate gradient and matrix multiplication and addition.
 
 - Changed random number generation int tests to be C++ style and fixed-seed, to avoid random failures.

resolve/CMakeLists.txt

@@ -157,6 +157,16 @@ endif()
 # Add HyKKT solver
 if(RESOLVE_USE_KLU)
   add_subdirectory(hykkt)
+  list(
+    APPEND
+    ReSolve_Targets_List
+    resolve_hykkt
+    resolve_hykkt_ruiz
+    resolve_hykkt_chol
+    resolve_hykkt_spgemm
+    resolve_hykkt_sccg
+    resolve_hykkt_solver
+  )
 endif()
 
 # Set installable targets

resolve/hykkt/HyKKTSolver.cpp

@@ -7,10 +7,12 @@
 #include "HyKKTSolver.hpp"
 
 #include <resolve/matrix/io.hpp>
+#include <resolve/utilities/logger/Logger.hpp>
 
 namespace ReSolve
 {
   using namespace constants;
+  using out = io::Logger;
 
   /**
    * @brief basic constructor
@@ -62,6 +64,9 @@ namespace ReSolve
    * values. It will only set pointers to user provided data; it is user's
    * responsibility to supply and later delete that memory.
    *
+   * Reusing the solver requires the sparsity patterns of the matrix blocks
+   * to remain unchanged. Changing J_d between empty and nonempty is rejected.
+   *
    * @param[in] H_plus_D_x - Pointer to the Hessian matrix block (n_x x n_x),
    * corresponding to H + D_x in the HyKKT paper.
    * @param[in] D_s - Pointer to the slack variables derivatives matrix block
@@ -70,18 +75,32 @@ namespace ReSolve
    * (m_c x n_x).
    * @param[in] J_d - Pointer to the inequality constraints Jacobian block
    * (m_d x n_x)
+   *
+   * @return 0 if successful, 1 if the supplied blocks are incompatible with
+   *         the allocated solver state. On failure, the previously stored
+   *         blocks are left unchanged.
    */
-  void hykkt::HyKKTSolver::setMatrixBlocks(matrix::Csr* H_plus_D_x, matrix::Csr* D_s, matrix::Csr* J, matrix::Csr* J_d)
+  int hykkt::HyKKTSolver::setMatrixBlocks(matrix::Csr* H_plus_D_x, matrix::Csr* D_s, matrix::Csr* J, matrix::Csr* J_d)
   {
+    const bool J_d_flag = J_d->getNnz() > 0;
+
+    if (allocated_ && J_d_flag_ != J_d_flag)
+    {
+      out::error() << "Changing J_d between empty and nonempty is not "
+                      "supported when reusing HyKKT.\n";
+      return 1;
+    }
+
     H_   = H_plus_D_x;
     D_s_ = D_s;
     J_   = J;
     J_d_ = J_d;
 
-    bool J_d_flag = J_d->getNnz() > 0;
-    // status_ = (J_d_flag_ == J_d_flag);
-    status_   = true; // when using API, we can't check if sparsity pattern changed
-    J_d_flag_ = J_d_flag;
+    if (!allocated_)
+    {
+      J_d_flag_ = J_d_flag;
+    }
+    return 0;
   }
 
   /**
@@ -156,16 +175,6 @@ namespace ReSolve
    */
   real_type hykkt::HyKKTSolver::solve()
   {
-    // TODO: Review sparsity pattern checking in HyKKT
-    if (!status_ && allocated_)
-    {
-      printf("\n\nERROR: USING HYKKT WITH NEW NONZERO STRUCTURE\n\n");
-      std::cout << "status = " << status_
-                << ", allocated = " << allocated_
-                << "\n";
-      return 1;
-    }
-
     setupParameters();
 
     if (!allocated_)
@@ -200,6 +209,16 @@ namespace ReSolve
     }
     computeHgammaFactorization();
 
+    if (!allocated_)
+    {
+      sccg_ = new SchurComplementConjugateGradient(J_->getNumRows(),
+                                                   J_->getNumColumns(),
+                                                   cholesky_,
+                                                   matrixHandler_,
+                                                   vectorHandler_,
+                                                   memspace_);
+      sccg_->setup();
+    }
     setupConjugateGradient();
     computeConjugateGradient();
 
@@ -210,8 +229,8 @@ namespace ReSolve
   /**
    * @brief allocates and initiates variables for KKT system
    *
-   * @pre jd_flag_ determines if variables used for Spgemm H_tilde
-   *      should be initiated
+   * @pre J_d_flag_ determines whether variables used to form H_tilde with
+   *      SpGEMM should be initialized.
    *
    * @post all variables used for hykkt are allocated for; J_d-
    *       related variables are not initiated if J_d nnz == 0
@@ -242,7 +261,6 @@ namespace ReSolve
       J_tr_perm_   = new matrix::Csr(J_tr_->getNumRows(), J_tr_->getNumColumns(), J_tr_->getNnz());
       J_d_scaled_  = new matrix::Csr(J_d_->getNumRows(), J_d_->getNumColumns(), J_d_->getNnz());
 
-      D_s_vals_->setData(D_s_->getValues(memspace_), memspace_);
       r_yd_scaled_->allocate(memspace_);
       r_x_perm_->allocate(memspace_);
       omega_perm_->allocate(memspace_);
@@ -258,14 +276,12 @@ namespace ReSolve
       J_d_scaled_->allocateWithExternalSparsityPattern(J_d_->getRowData(memspace_), J_d_->getColData(memspace_), J_d_->getNnz(), memspace_);
       // H_tilde_ does not need to be allocated because loadResultMatrix() does it later
     }
-    else if (memspace_ == memory::DEVICE)
-    {
-      J_d_->syncData(memory::DEVICE); // check if this is redundant
-    }
 
+    // D_s may be replaced between solves, so refresh the external value pointer.
+    D_s_vals_->setData(D_s_->getValues(memspace_), memspace_);
     r_y_copy_->copyFromExternal(r_y_, memspace_, memspace_);
 
-    // check if this is redundant in later iterations
+    // Matrix values may change between solves, so refresh the transpose.
     matrixHandler_->transpose(J_, J_tr_, memspace_);
     if (J_d_flag_)
     {
@@ -314,7 +330,10 @@ namespace ReSolve
     else
     {
       H_tilde_->setNnz(H_->getNnz());
-      H_tilde_->allocateMatrixData(memspace_);
+      if (!allocated_)
+      {
+        H_tilde_->allocateMatrixData(memspace_);
+      }
       H_tilde_->copyFromExternal(H_->getRowData(memspace_),
                                  H_->getColData(memspace_),
                                  H_->getValues(memspace_),
@@ -388,9 +407,6 @@ namespace ReSolve
   void hykkt::HyKKTSolver::setupSpGEMMHgamma()
   {
     spgemm_hgamma_ = new SpGEMM(memspace_, gamma_, ONE);
-    spgemm_hgamma_->loadProductMatrices(J_tr_, J_);
-    spgemm_hgamma_->loadSumMatrix(H_tilde_);
-    spgemm_hgamma_->loadResultMatrix(&H_gamma_); // H_gamma_ will be created when calling SpGEMM->compute()
   }
 
   /*
@@ -403,6 +419,14 @@ namespace ReSolve
    */
   void hykkt::HyKKTSolver::computeSpGEMMHgamma()
   {
+    // Numerical values can change between solves while the sparsity pattern
+    // remains fixed, so refresh the SpGEMM inputs before recomputing H_gamma.
+    spgemm_hgamma_->setCoefficients(gamma_, ONE);
+    spgemm_hgamma_->loadProductMatrices(J_tr_, J_);
+    spgemm_hgamma_->loadSumMatrix(H_tilde_);
+    // HIP initializes the result descriptor using dimensions established by
+    // the product and sum inputs, so load the result matrix after both inputs.
+    spgemm_hgamma_->loadResultMatrix(&H_gamma_);
     spgemm_hgamma_->compute();
     r_x_hat_->copyFromExternal(r_x_til_, memspace_, memspace_);
     matrixHandler_->matvec(J_tr_, r_y_, r_x_hat_, &gamma_, &ONE, memspace_);
@@ -514,11 +538,9 @@ namespace ReSolve
     schur_->copyFromExternal(r_y_, memspace_, memspace_);
     matrixHandler_->matvec(J_perm_, omega_perm_, schur_, &ONE, &MINUS_ONE, memspace_);
 
-    sccg_ = new SchurComplementConjugateGradient(J_->getNumRows(), J_->getNumColumns(), cholesky_, matrixHandler_, vectorHandler_, memspace_);
     sccg_->addMatrixInfo(J_perm_, J_tr_perm_);
     y_->setToZero(memspace_);
     sccg_->addVectorInfo(y_, schur_);
-    sccg_->setup();
   }
 
   /**
@@ -553,10 +575,6 @@ namespace ReSolve
     // block-recovering the solution to the original system by parts
     // this part is to recover delta_x
     cholesky_->solve(z_, r_x_perm_);
-    if (memspace_ == memory::DEVICE)
-    {
-      x_->syncData(memory::DEVICE);
-    }
     permutation_->mapIndex(REV_PERM_V, z_->getData(memspace_), x_->getData(memspace_));
     x_->setDataUpdated(memspace_);
 
@@ -627,10 +645,17 @@ namespace ReSolve
     matrixHandler_->matvec(J_copy_, x_, r_y_copy_, &MINUS_ONE, &ONE, memspace_);
     norm_resy_sq = vectorHandler_->dot(r_y_copy_, r_y_copy_, memspace_);
 
-    // Calculate final relative norm
     norm_resx_sq += norm_resy_sq;
-    real_type norm_res = sqrt(norm_resx_sq) / sqrt(norm_r_x_sq);
-    printf("||Ax-b||/||b|| = %32.32g\n\n", norm_res);
+    real_type norm_res = sqrt(norm_resx_sq);
+    if (norm_r_x_sq > 0)
+    {
+      norm_res /= sqrt(norm_r_x_sq);
+      printf("||Ax-b||/||b|| = %32.32g\n\n", norm_res);
+    }
+    else
+    {
+      printf("||Ax-b|| = %32.32g\n\n", norm_res);
+    }
 
     allocated_ = true;
 

resolve/hykkt/HyKKTSolver.hpp

@@ -41,7 +41,7 @@ namespace ReSolve
           std::istream& r_y_file,
           std::istream& r_yd_file);
 
-      void setMatrixBlocks(matrix::Csr* H_plus_D_x, matrix::Csr* D_s, matrix::Csr* J, matrix::Csr* J_d);
+      int  setMatrixBlocks(matrix::Csr* H_plus_D_x, matrix::Csr* D_s, matrix::Csr* J, matrix::Csr* J_d);
       void setRHSBlocks(vector::Vector* r_x, vector::Vector* r_s, vector::Vector* r_y, vector::Vector* r_yd);
       void setLHSPointers(vector::Vector* x, vector::Vector* s, vector::Vector* y, vector::Vector* y_d);
 
@@ -78,10 +78,6 @@ namespace ReSolve
       bool allocated_ = false;
       bool J_d_flag_  = false;
 
-      // Whether the solver is correctly used with matrices of
-      // the same nonzero structure
-      bool status_ = true;
-
       RuizScaling*                      ruiz_{nullptr};
       SpGEMM*                           spgemm_htil_{nullptr};
       SpGEMM*                           spgemm_hgamma_{nullptr};

resolve/hykkt/cholesky/CholeskySolverCpu.cpp

@@ -39,11 +39,12 @@ namespace ReSolve
 
     void CholeskySolverCpu::addMatrixInfo(matrix::Csr* A)
     {
+      A_ = A;
       if (A_chol_)
       {
         cholmod_free_sparse(&A_chol_, &Common_);
       }
-      A_chol_ = convertToCholmod(A);
+      A_chol_ = convertToCholmod(A_);
     }
 
     /**
@@ -72,6 +73,12 @@ namespace ReSolve
     {
       (void) tol; // Mark tol as unused
 
+      // CHOLMOD stores its own copy of the matrix values. Refresh that copy
+      // before numerical refactorization when the solver is reused.
+      mem_.copyArrayHostToHost(static_cast<real_type*>(A_chol_->x),
+                               A_->getValues(memory::HOST),
+                               A_->getNnz());
+
       cholmod_factorize(A_chol_, factorization_, &Common_);
       if (Common_.status < 0)
       {

resolve/hykkt/cholesky/CholeskySolverCpu.hpp

@@ -27,6 +27,7 @@ namespace ReSolve
       MemoryHandler mem_;
 
       cholmod_common  Common_;
+      matrix::Csr*    A_ = nullptr;
       cholmod_sparse* A_chol_; // cholmod sparse matrix representation
       cholmod_factor* factorization_;
 

resolve/hykkt/sccg/SchurComplementConjugateGradient.cpp

@@ -99,6 +99,11 @@ namespace ReSolve
       p_->allocate(memspace_);
       s_->allocate(memspace_);
       w_->allocate(memspace_);
+    }
+
+    int SchurComplementConjugateGradient::solve()
+    {
+      using namespace constants;
 
       y_->setToZero(memspace_);
       z_->setToZero(memspace_);
@@ -110,17 +115,18 @@ namespace ReSolve
       x_0_->setToZero(memspace_);
 
       beta_ = 0;
-    }
-
-    int SchurComplementConjugateGradient::solve()
-    {
-      using namespace constants;
 
       matrix_handler_->matvec(J_tr_, x_0_, y_, &ONE, &ZERO, memspace_);
       choleskySolver_->solve(z_, y_);
       matrix_handler_->matvec(J_, z_, r_, &MINUS_ONE, &ONE, memspace_);
       gamma_i_ = vector_handler_->dot(r_, r_, memspace_);
 
+      if (sqrt(gamma_i_) < tol_)
+      {
+        gamma_i1_ = gamma_i_;
+        return 0;
+      }
+
       matrix_handler_->matvec(J_tr_, r_, y_, &ONE, &ZERO, memspace_);
       choleskySolver_->solve(z_, y_);
       matrix_handler_->matvec(J_, z_, w_, &ONE, &ZERO, memspace_);

resolve/hykkt/sccg/SchurComplementConjugateGradient.hpp

@@ -49,6 +49,11 @@ namespace ReSolve
       void setSolverTolerance(double tol);
       void setSolverItmax(int itmax);
 
+      /**
+       * @brief Allocates internal work vectors.
+       *
+       * @pre Must be called exactly once before the first call to solve().
+       */
       void setup();
       int  solve();
 

resolve/hykkt/spgemm/SpGEMM.cpp

@@ -53,6 +53,17 @@ namespace ReSolve
       delete impl_;
     }
 
+    /**
+     * Updates the coefficients for the SpGEMM operation.
+     *
+     * @param[in] alpha - Scalar multiplier for the matrix product.
+     * @param[in] beta  - Scalar multiplier for the sum matrix.
+     */
+    void SpGEMM::setCoefficients(real_type alpha, real_type beta)
+    {
+      impl_->setCoefficients(alpha, beta);
+    }
+
     /**
      * Loads the two matrices for the product
      * @param A[in] - Pointer to CSR matrix

resolve/hykkt/spgemm/SpGEMM.hpp

@@ -22,6 +22,8 @@ namespace ReSolve
       SpGEMM(memory::MemorySpace memspace, real_type alpha, real_type beta);
       ~SpGEMM();
 
+      void setCoefficients(real_type alpha, real_type beta);
+
       void loadProductMatrices(matrix::Csr* A, matrix::Csr* B);
       void loadSumMatrix(matrix::Csr* D);
       void loadResultMatrix(matrix::Csr** E_ptr);

resolve/hykkt/spgemm/SpGEMMCpu.cpp

@@ -47,6 +47,12 @@ namespace ReSolve
       cholmod_finish(&Common_);
     }
 
+    void SpGEMMCpu::setCoefficients(real_type alpha, real_type beta)
+    {
+      alpha_ = alpha;
+      beta_  = beta;
+    }
+
     void SpGEMMCpu::loadProductMatrices(matrix::Csr* A, matrix::Csr* B)
     {
       if (!A_)

resolve/hykkt/spgemm/SpGEMMCpu.hpp

@@ -16,11 +16,13 @@ namespace ReSolve
       SpGEMMCpu(real_type alpha, real_type beta);
       ~SpGEMMCpu();
 
-      void loadProductMatrices(matrix::Csr* A, matrix::Csr* B);
-      void loadSumMatrix(matrix::Csr* D);
-      void loadResultMatrix(matrix::Csr** E_ptr);
+      void setCoefficients(real_type alpha, real_type beta) override;
 
-      void compute();
+      void loadProductMatrices(matrix::Csr* A, matrix::Csr* B) override;
+      void loadSumMatrix(matrix::Csr* D) override;
+      void loadResultMatrix(matrix::Csr** E_ptr) override;
+
+      void compute() override;
 
     private:
       real_type alpha_;