Sparse NT Scaling

algebra/cuda/cudss_backend.cu

@@ -390,8 +390,15 @@ static void csc_to_csr_device(const QOCOCscMatrix* csc, QOCOInt** csr_row_ptr,
 }
 
 static LinSysData* cudss_setup(QOCOProblemData* data, QOCOSettings* settings,
-                               QOCOInt Wnnz, QOCOFloat* analysis_time_sec)
+                               QOCOInt Wnnz, QOCOInt nsoc_sparse,
+                               QOCOInt* soc_is_sparse, QOCOInt nt_sparse_nnz,
+                               QOCOInt* sparse_soc_nt_idx,
+                               QOCOFloat* analysis_time_sec)
 {
+  (void)nsoc_sparse;
+  (void)soc_is_sparse;
+  (void)nt_sparse_nnz;
+  (void)sparse_soc_nt_idx;
   // Load CUDA libraries dynamically
   if (!load_cuda_libraries()) {
     fprintf(stderr, "Failed to load CUDA libraries\n");
@@ -447,9 +454,10 @@ static LinSysData* cudss_setup(QOCOProblemData* data, QOCOSettings* settings,
   QOCOCscMatrix* Kcsc = construct_kkt(
       get_csc_matrix(data->P), get_csc_matrix(data->A), get_csc_matrix(data->G),
       get_csc_matrix(data->At), get_csc_matrix(data->Gt),
-      settings->kkt_static_reg_A, data->n, data->m, data->p, data->l, data->nsoc,
-      get_data_vectori(data->q), linsys_data->PregtoKKT, linsys_data->AttoKKT,
-      linsys_data->GttoKKT, linsys_data->nt2kkt, linsys_data->ntdiag2kkt, Wnnz);
+      settings->kkt_static_reg_A, data->n, data->m, data->p, data->l,
+      data->nsoc, get_data_vectori(data->q), linsys_data->PregtoKKT,
+      linsys_data->AttoKKT, linsys_data->GttoKKT, linsys_data->nt2kkt,
+      linsys_data->ntdiag2kkt, Wnnz, NULL, 0, 0, NULL, NULL, NULL, NULL);
   set_cpu_mode(0);
 
   // Convert KKT matrix from CSC (CPU) to CSR (GPU) for cuDSS
@@ -460,8 +468,8 @@ static LinSysData* cudss_setup(QOCOProblemData* data, QOCOSettings* settings,
   QOCOInt* h_csr_col_ind;
   QOCOInt* csc2csr;
 
-  csc_to_csr_device(Kcsc, &csr_row_ptr, &csr_col_ind, &csr_val,
-                    &h_csr_row_ptr, &h_csr_col_ind, &csc2csr);
+  csc_to_csr_device(Kcsc, &csr_row_ptr, &csr_col_ind, &csr_val, &h_csr_row_ptr,
+                    &h_csr_col_ind, &csc2csr);
 
   // Build nt2kktcsr and ntdiag2kktcsr mappings (CSR indices instead of CSC)
   QOCOInt* nt2kktcsr = NULL;
@@ -610,8 +618,7 @@ static LinSysData* cudss_setup(QOCOProblemData* data, QOCOSettings* settings,
 __global__ void
 update_csr_nt_blocks_kernel(const QOCOFloat* WtW, // NT block values (on GPU)
                             QOCOFloat* csr_val, // CSR values to update (on GPU)
-                            const QOCOInt* nt2kktcsr,
-                            QOCOInt Wnnz)
+                            const QOCOInt* nt2kktcsr, QOCOInt Wnnz)
 {
   QOCOInt idx = blockIdx.x * blockDim.x + threadIdx.x;
 
@@ -711,8 +718,7 @@ static void cudss_solve_system(LinSysData* linsys_data, const QOCOFloat* rhs,
  */
 static QOCOFloat compute_linsys_residual(LinSysData* linsys_data,
                                          QOCOWorkspace* work,
-                                         const QOCOFloat* b,
-                                         const QOCOFloat* x,
+                                         const QOCOFloat* b, const QOCOFloat* x,
                                          QOCOFloat* residual_scratch)
 {
   QOCOFloat* nt_scaling = get_data_vectorf(work->nt_scaling);
@@ -767,14 +773,13 @@ static void cudss_solve(LinSysData* linsys_data, QOCOWorkspace* work,
 #ifdef QOCO_LOGGING
   FILE* log_f = fopen("qoco_log.txt", "a");
   if (log_f) {
-    log_linsys_error(linsys_data, work, b, x, residual, "initial solve",
-                     log_f);
+    log_linsys_error(linsys_data, work, b, x, residual, "initial solve", log_f);
   }
 #endif
 
   QOCOFloat* best_sol = get_data_vectorf(work->xyzbuff1);
-  QOCOFloat best_res = compute_linsys_residual(linsys_data, work, b, x,
-                                               residual);
+  QOCOFloat best_res =
+      compute_linsys_residual(linsys_data, work, b, x, residual);
   copy_arrayf(x, best_sol, linsys_data->Kn);
 
   QOCOInt ir_count = 0;
@@ -791,8 +796,8 @@ static void cudss_solve(LinSysData* linsys_data, QOCOWorkspace* work,
     // x_new = x_old + dx.
     qoco_axpy(linsys_data->d_xyz_matrix_data, x, x, 1.0, linsys_data->Kn);
 
-    QOCOFloat new_res = compute_linsys_residual(linsys_data, work, b, x,
-                                                residual);
+    QOCOFloat new_res =
+        compute_linsys_residual(linsys_data, work, b, x, residual);
 
 #ifdef QOCO_LOGGING
     if (log_f) {
@@ -819,8 +824,9 @@ static void cudss_solve(LinSysData* linsys_data, QOCOWorkspace* work,
 #endif
 }
 
-void cudss_set_nt_identity(LinSysData* linsys_data, QOCOInt m)
+void cudss_set_nt_identity(LinSysData* linsys_data, QOCOWorkspace* work)
 {
+  QOCOInt m = work->data->m;
   int Wnnz = linsys_data->Wnnz;
 
   int N = max(Wnnz, m);
@@ -848,9 +854,11 @@ void cudss_set_nt_identity(LinSysData* linsys_data, QOCOInt m)
   }
 }
 
-static void cudss_update_nt(LinSysData* linsys_data, QOCOVectorf* WtW_vec,
-                            QOCOFloat kkt_static_reg_G, QOCOInt m)
+static void cudss_update_nt(LinSysData* linsys_data, QOCOWorkspace* work,
+                            QOCOFloat kkt_static_reg_G)
 {
+  QOCOVectorf* WtW_vec = work->WtW;
+  QOCOInt m = work->data->m;
   QOCOFloat* WtW = get_data_vectorf(WtW_vec);
   // Update CSR matrix values on GPU directly for NT blocks
   if (linsys_data->Wnnz > 0 && linsys_data->d_nt2kktcsr) {

include/cone.h

@@ -18,6 +18,9 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
+
+/** SOC cones with dim > this threshold use the sparse NT expansion. */
+#define SOC_SPARSE_THRESHOLD 4
 #include "qoco_linalg.h"
 #include "qoco_utils.h"
 #include "structs.h"
@@ -151,4 +154,4 @@ void add_e(QOCOFloat* x, QOCOFloat a, QOCOInt l, QOCOInt nsoc, QOCOVectori* q);
 }
 #endif
 
-#endif
+#endif

include/kkt.h

@@ -49,14 +49,15 @@
  * @param Wnnz Number of nonzeros in upper triangular portion of NT scaling.
  * @return QOCOCscMatrix*
  */
-QOCOCscMatrix* construct_kkt(QOCOCscMatrix* P, QOCOCscMatrix* A,
-                             QOCOCscMatrix* G, QOCOCscMatrix* At,
-                             QOCOCscMatrix* Gt, QOCOFloat kkt_static_reg_A,
-                             QOCOInt n, QOCOInt m, QOCOInt p, QOCOInt l,
-                             QOCOInt nsoc, QOCOInt* q, QOCOInt* PregtoKKT,
-                             QOCOInt* AttoKKT, QOCOInt* GttoKKT,
-                             QOCOInt* nt2kkt, QOCOInt* ntdiag2kkt,
-                             QOCOInt Wnnz);
+QOCOCscMatrix*
+construct_kkt(QOCOCscMatrix* P, QOCOCscMatrix* A, QOCOCscMatrix* G,
+              QOCOCscMatrix* At, QOCOCscMatrix* Gt, QOCOFloat kkt_static_reg_A,
+              QOCOInt n, QOCOInt m, QOCOInt p, QOCOInt l, QOCOInt nsoc,
+              QOCOInt* q, QOCOInt* PregtoKKT, QOCOInt* AttoKKT,
+              QOCOInt* GttoKKT, QOCOInt* nt2kkt, QOCOInt* ntdiag2kkt,
+              QOCOInt Wnnz, QOCOInt* soc_is_sparse, QOCOInt nsoc_sparse,
+              QOCOInt nt_sparse_nnz, QOCOInt* sparse_soc_nt_idx,
+              QOCOInt* nt_u2kkt, QOCOInt* nt_v2kkt, QOCOInt* nt_uvdiag2kkt);
 
 /**
  * @brief Gets initial values for primal and dual variables such that (s,z) \in
@@ -164,4 +165,4 @@ void kkt_multiply(QOCOFloat* x, QOCOFloat* y, QOCOProblemData* data,
                   QOCOFloat* nt_scaling, QOCOInt* nt_scaling_soc_idx,
                   QOCOInt* soc_idx, QOCOFloat* nbuff, QOCOFloat* mbuff1,
                   QOCOFloat* mbuff2);
-#endif /* #ifndef QOCO_KKT_H */
+#endif /* #ifndef QOCO_KKT_H */

include/structs.h

@@ -230,9 +230,13 @@ typedef struct {
   /** Upper triangular part of Nesterov-Todd Scaling */
   QOCOVectorf* W;
 
-  /** NT scaling data used by nt_multiply(). Shared layout across backends:
+  /** NT scaling data used by nt_multiply().
+   *
+   * Builtin backend layout:
    *   LP entries: scalar scales sqrt(s_i / z_i), length l.
    *   SOC i block: [eta, w0, w1[0], ..., w1[q_i - 2]], length q_i + 1.
+   *
+   * CUDA backend currently keeps the historical dense q_i-by-q_i SOC blocks.
    */
   QOCOVectorf* nt_scaling;
 
@@ -249,6 +253,31 @@ typedef struct {
    * block */
   QOCOVectori* soc_idx;
 
+  /** Number of SOCs using sparse NT expansion (dim > SOC_SPARSE_THRESHOLD). */
+  QOCOInt nsoc_sparse;
+
+  /** Per-SOC flag: 1 if cone i uses sparse NT expansion, 0 otherwise. */
+  QOCOInt* soc_is_sparse;
+
+  /** Packed u vectors for sparse NT expansion (length nt_sparse_nnz). */
+  QOCOVectorf* nt_u_sparse;
+
+  /** Packed v vectors for sparse NT expansion (length nt_sparse_nnz). */
+  QOCOVectorf* nt_v_sparse;
+
+  /** eta^2 values for sparse SOCs (length nsoc_sparse). */
+  QOCOVectorf* nt_eta2_sparse;
+
+  /** d scalars for sparse SOCs (length nsoc_sparse). */
+  QOCOVectorf* nt_d_sparse;
+
+  /** Index into nt_u/v_sparse for the k-th sparse SOC (length nsoc_sparse). */
+  QOCOVectori* sparse_soc_nt_idx;
+
+  /** Total elements in nt_u_sparse / nt_v_sparse: sum of q[i] for sparse SOCs.
+   */
+  QOCOInt nt_sparse_nnz;
+
   /** Scaled variables. */
   QOCOVectorf* lambda;
 
@@ -364,10 +393,13 @@ typedef struct LinSysData LinSysData;
 typedef struct {
   const char* (*linsys_name)();
   LinSysData* (*linsys_setup)(QOCOProblemData* data, QOCOSettings* settings,
-                              QOCOInt Wnnz, QOCOFloat* analysis_time_sec);
-  void (*linsys_set_nt_identity)(LinSysData* linsys_data, QOCOInt m);
-  void (*linsys_update_nt)(LinSysData* linsys_data, QOCOVectorf* WtW_vec,
-                           QOCOFloat kkt_static_reg_G, QOCOInt m);
+                              QOCOInt Wnnz, QOCOInt nsoc_sparse,
+                              QOCOInt* soc_is_sparse, QOCOInt nt_sparse_nnz,
+                              QOCOInt* sparse_soc_nt_idx,
+                              QOCOFloat* analysis_time_sec);
+  void (*linsys_set_nt_identity)(LinSysData* linsys_data, QOCOWorkspace* work);
+  void (*linsys_update_nt)(LinSysData* linsys_data, QOCOWorkspace* work,
+                           QOCOFloat kkt_static_reg_G);
   void (*linsys_update_data)(LinSysData* linsys_data, QOCOProblemData* data);
   void (*linsys_factor)(LinSysData* linsys_data, QOCOInt n,
                         QOCOFloat kkt_dynamic_reg);
@@ -400,4 +432,4 @@ typedef struct {
 
 } QOCOSolver;
 
-#endif /* #ifndef QOCO_STRUCTS_H */
+#endif /* #ifndef QOCO_STRUCTS_H */

lib/qdldl/include/qdldl.h

@@ -140,7 +140,8 @@ QDLDL_API QDLDL_int QDLDL_factor(const QDLDL_int n, const QDLDL_int* Ap,
                                  const QDLDL_int* Lnz, const QDLDL_int* etree,
                                  QDLDL_bool* bwork, QDLDL_int* iwork,
                                  QDLDL_float* fwork, QDLDL_int* perm,
-                                 QDLDL_int pos_diags, QDLDL_float dyn_reg);
+                                 QDLDL_int pos_diags, QDLDL_float dyn_reg,
+                                 const QDLDL_bool* positive_diag);
 
 /**
  * Solves LDL'x = b

lib/qdldl/src/qdldl.c

@@ -98,7 +98,8 @@ QDLDL_int QDLDL_factor(const QDLDL_int n, const QDLDL_int* Ap,
                        QDLDL_float* D, QDLDL_float* Dinv, const QDLDL_int* Lnz,
                        const QDLDL_int* etree, QDLDL_bool* bwork,
                        QDLDL_int* iwork, QDLDL_float* fwork, QDLDL_int* perm,
-                       QDLDL_int pos_diags, QDLDL_float dyn_reg)
+                       QDLDL_int pos_diags, QDLDL_float dyn_reg,
+                       const QDLDL_bool* positive_diag)
 {
 
   QDLDL_int i, j, k, nnzY, bidx, cidx, nextIdx, nnzE, tmpIdx;
@@ -138,7 +139,7 @@ QDLDL_int QDLDL_factor(const QDLDL_int n, const QDLDL_int* Ap,
   }
 
   // Dynamic regularization
-  if (perm[0] < pos_diags) {
+  if (positive_diag ? positive_diag[perm[0]] : perm[0] < pos_diags) {
     D[0] = D[0] < 1e-11 ? dyn_reg : D[0];
   }
   else {
@@ -248,7 +249,7 @@ QDLDL_int QDLDL_factor(const QDLDL_int n, const QDLDL_int* Ap,
     } // end for i
 
     // Dynamic regularization
-    if (perm[k] < pos_diags) {
+    if (positive_diag ? positive_diag[perm[k]] : perm[k] < pos_diags) {
       D[k] = D[k] < 1e-11 ? dyn_reg : D[k];
     }
     else {