Add fast NT scaling computations for W*z and W^{-1}*z

.github/workflows/unit_tests.yml

@@ -50,8 +50,8 @@ jobs:
             c_compiler: clang
             cpp_compiler: clang++
           - os: windows-latest
-            c_compiler: MSVC
-            cpp_compiler: MSVC
+            c_compiler: cl
+            cpp_compiler: cl
           - precision: long_double
             long_double_flag: -DQOCO_LONG_DOUBLE_PRECISION:BOOL=True
           - precision: double

algebra/builtin/qdldl_backend.c

@@ -224,7 +224,9 @@ static QOCOFloat compute_linsys_residual(LinSysData* linsys_data,
                                          QOCOWorkspace* work, QOCOFloat* b,
                                          QOCOFloat* x_scratch)
 {
-  QOCOFloat* Wfull = get_data_vectorf(work->Wfull);
+  QOCOFloat* nt_scaling = get_data_vectorf(work->nt_scaling);
+  QOCOInt* nt_scaling_soc_idx = get_data_vectori(work->nt_scaling_soc_idx);
+  QOCOInt* soc_idx = get_data_vectori(work->soc_idx);
   QOCOFloat* xbuff = get_data_vectorf(work->xbuff);
   QOCOFloat* ubuff1 = get_data_vectorf(work->ubuff1);
   QOCOFloat* ubuff2 = get_data_vectorf(work->ubuff2);
@@ -239,8 +241,8 @@ static QOCOFloat compute_linsys_residual(LinSysData* linsys_data,
   // Compute K_true * x_scratch -> xyzbuff2 against the unregularized matrix.
   // data->P stores the regularized P (P + eps_P * I), so subtract the P
   // regularization contribution from the x block to recover the true product.
-  kkt_multiply(x_scratch, linsys_data->xyzbuff2, work->data, Wfull, NULL, NULL,
-               xbuff, ubuff1, ubuff2);
+  kkt_multiply(x_scratch, linsys_data->xyzbuff2, work->data, nt_scaling,
+               nt_scaling_soc_idx, soc_idx, xbuff, ubuff1, ubuff2);
   for (QOCOInt k = 0; k < n; ++k) {
     linsys_data->xyzbuff2[k] -= linsys_data->kkt_static_reg_P * x_scratch[k];
   }

algebra/cuda/cudss_backend.cu

@@ -715,8 +715,8 @@ static QOCOFloat compute_linsys_residual(LinSysData* linsys_data,
                                          const QOCOFloat* x,
                                          QOCOFloat* residual_scratch)
 {
-  QOCOFloat* Wfull = get_data_vectorf(work->Wfull);
-  QOCOInt* Wsoc_idx = get_data_vectori(work->Wsoc_idx);
+  QOCOFloat* nt_scaling = get_data_vectorf(work->nt_scaling);
+  QOCOInt* nt_scaling_soc_idx = get_data_vectori(work->nt_scaling_soc_idx);
   QOCOInt* soc_idx = get_data_vectori(work->soc_idx);
   QOCOFloat* xbuff = get_data_vectorf(work->xbuff);
   QOCOFloat* ubuff1 = get_data_vectorf(work->ubuff1);
@@ -726,8 +726,8 @@ static QOCOFloat compute_linsys_residual(LinSysData* linsys_data,
 
   // d_rhs_matrix_data is scratch here; cudss_solve_system overwrites it before
   // every cuDSS solve.
-  kkt_multiply((QOCOFloat*)x, linsys_data->d_rhs_matrix_data, work->data, Wfull,
-               Wsoc_idx, soc_idx, xbuff, ubuff1, ubuff2);
+  kkt_multiply((QOCOFloat*)x, linsys_data->d_rhs_matrix_data, work->data,
+               nt_scaling, nt_scaling_soc_idx, soc_idx, xbuff, ubuff1, ubuff2);
 
   // data->P stores P + eps_P * I, so remove the P regularization from the
   // product before measuring the true KKT residual.

include/cone.h

@@ -23,16 +23,17 @@ extern "C" {
 #include "structs.h"
 
 /**
- * @brief Sets Wfull to I.
+ * @brief Sets NT scaling data to identity.
  *
- * @param Wfull Full NT scaling matrix.
- * @param Wnnzfull Number of elements in Wfull.
- * @param Wsoc_idx Vector pointing to the start of each SOC block in Wfull.
+ * @param nt_scaling NT scaling data.
+ * @param nt_scaling_nnz Number of elements in nt_scaling.
+ * @param nt_scaling_soc_idx Vector pointing to the start of each SOC block.
  *                 Only used in the GPU (cone.cu) implementation.
  * @param data Pointer to problem data.
  */
-void set_Wfull_identity(QOCOVectorf* Wfull, QOCOInt Wnnzfull,
-                        QOCOVectori* Wsoc_idx, QOCOProblemData* data);
+void set_nt_scaling_identity(QOCOVectorf* nt_scaling, QOCOInt nt_scaling_nnz,
+                             QOCOVectori* nt_scaling_soc_idx,
+                             QOCOProblemData* data);
 
 /**
  * @brief Computes cone product u * v = p with respect to C.
@@ -79,13 +80,13 @@ void cone_division(const QOCOFloat* lambda, const QOCOFloat* v, QOCOFloat* d,
 void bring2cone(QOCOFloat* u, QOCOInt* soc_idx, QOCOProblemData* data);
 
 /**
- * @brief Computes z = W * x where W is a full Nesterov-Todd scaling matrix.
- * The NT scaling array for the LP cones are stored first, then the NT
- * scalings for the second-order cones are stored in column major order.
+ * @brief Computes z = W * x using Nesterov-Todd scaling data.
+ * The LP cone scalings are stored first. Each SOC block stores the fast
+ * scaling parameters [eta, w0, w1...] in the builtin backend.
  *
- * @param W Nesterov Todd scaling matrix.
- * @param Wsoc_idx Vector pointing to the start of each SOC block in W.
- *                 Only used in the GPU (cone.cu) implementation.
+ * @param W Nesterov Todd scaling data.
+ * @param nt_scaling_soc_idx Vector pointing to the start of each SOC block in
+ * W. Only used in the GPU (cone.cu) implementation.
  * @param soc_idx Array pointing to the start of each SOC block in x and z.
  *                Only used in the GPU (cone.cu) implementation.
  * @param x Input vector.
@@ -95,10 +96,18 @@ void bring2cone(QOCOFloat* u, QOCOInt* soc_idx, QOCOProblemData* data);
  * @param nsoc Number of second-order cones in C.
  * @param q Array of second-order cone dimensions.
  */
-void nt_multiply(QOCOFloat* W, QOCOInt* Wsoc_idx, QOCOInt* soc_idx,
+void nt_multiply(QOCOFloat* W, QOCOInt* nt_scaling_soc_idx, QOCOInt* soc_idx,
                  QOCOFloat* x, QOCOFloat* z, QOCOInt l, QOCOInt m, QOCOInt nsoc,
                  QOCOInt* q);
 
+/**
+ * @brief Computes z = W^{-1} * x using the same NT scaling data as
+ * nt_multiply().
+ */
+void nt_multiply_inv(QOCOFloat* W, QOCOInt* nt_scaling_soc_idx,
+                     QOCOInt* soc_idx, QOCOFloat* x, QOCOFloat* z, QOCOInt l,
+                     QOCOInt m, QOCOInt nsoc, QOCOInt* q);
+
 /**
  * @brief Compute Nesterov-Todd scalings and scaled variables.
  *

include/kkt.h

@@ -152,14 +152,16 @@ void predictor_corrector(QOCOSolver* solver);
  * @param x Pointer to input vector.
  * @param y Pointer to output vector.
  * @param data Pointer to problem data.
- * @param Wfull Pointer to full NT scaling matrix W.
- * @param Wsoc_idx Vector pointing to the start of each SOC block in Wfull.
+ * @param nt_scaling Pointer to NT scaling data.
+ * @param nt_scaling_soc_idx Vector pointing to the start of each SOC block in
+ * nt_scaling.
  * @param soc_idx Array pointing to the start of each SOC block in x and y.
  * @param nbuff Temporary buffer of length n.
  * @param mbuff1 Temporary buffer of length m.
  * @param mbuff2 Temporary buffer of length m.
  */
 void kkt_multiply(QOCOFloat* x, QOCOFloat* y, QOCOProblemData* data,
-                  QOCOFloat* Wfull, QOCOInt* Wsoc_idx, QOCOInt* soc_idx,
-                  QOCOFloat* nbuff, QOCOFloat* mbuff1, QOCOFloat* mbuff2);
+                  QOCOFloat* nt_scaling, QOCOInt* nt_scaling_soc_idx,
+                  QOCOInt* soc_idx, QOCOFloat* nbuff, QOCOFloat* mbuff1,
+                  QOCOFloat* mbuff2);
 #endif /* #ifndef QOCO_KKT_H */

include/structs.h

@@ -224,26 +224,26 @@ typedef struct {
   /** Number of nonzeros in upper triangular part of Nesterov-Todd Scaling. */
   QOCOInt Wnnz;
 
-  /** Number of nonzeros in full Nesterov-Todd Scaling. */
-  QOCOInt Wnnzfull;
+  /** Number of entries in nt_scaling. */
+  QOCOInt nt_scaling_nnz;
 
   /** Upper triangular part of Nesterov-Todd Scaling */
   QOCOVectorf* W;
 
-  /** Full Nesterov-Todd Scaling */
-  QOCOVectorf* Wfull;
+  /** NT scaling data used by nt_multiply(). Shared layout across backends:
+   *   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.
+   */
+  QOCOVectorf* nt_scaling;
 
   /** Upper triangular part of inverse of Nesterov-Todd Scaling */
   QOCOVectorf* Winv;
 
-  /** Full inverse of Nesterov-Todd Scaling */
-  QOCOVectorf* Winvfull;
-
   /** Nesterov-Todd Scaling squared */
   QOCOVectorf* WtW;
 
-  /** Vector which points to the start of the ith soc block in Wfull */
-  QOCOVectori* Wsoc_idx;
+  /** Vector which points to the start of the ith SOC block in nt_scaling. */
+  QOCOVectori* nt_scaling_soc_idx;
 
   /** Vector which points to the start of the start of the ith soc variable
    * block */