scs 0.3.1 → 0.3.2

data/vendor/scs/linsys/cpu/indirect/private.c

@@ -17,7 +17,8 @@ char *SCS(get_lin_sys_summary)(ScsLinSysWork *p, const ScsInfo *info) {
 }
 */
 
-/* set M = inv ( diag ( rho_x * I + P + A' R_y^{-1} A ) ) */
+/* Not possible to do this on the fly due to M_ii += a_i' (R_y)^-1 a_i */
+/* set M = inv ( diag ( R_x + P + A' R_y^{-1} A ) ) */
 static void set_preconditioner(ScsLinSysWork *p) {
   scs_int i, k;
   scs_float *M = p->M;
@@ -28,22 +29,26 @@ static void set_preconditioner(ScsLinSysWork *p) {
   scs_printf("getting pre-conditioner\n");
 #endif
 
+  /* M_ii = (R_x)_i + P_ii + a_i' (R_y)^-1 a_i */
   for (i = 0; i < A->n; ++i) { /* cols */
-    M[i] = p->rho_x;
-    /* diag(A' R_y^{-1} A) */
+    /* M_ii = (R_x)_i */
+    M[i] = p->diag_r[i];
+    /* M_ii += a_i' (R_y)^-1 a_i */
     for (k = A->p[i]; k < A->p[i + 1]; ++k) {
       /* A->i[k] is row of entry k with value A->x[k] */
-      M[i] += A->x[k] * A->x[k] / p->rho_y_vec[A->i[k]];
+      M[i] += A->x[k] * A->x[k] / p->diag_r[A->n + A->i[k]];
     }
     if (P) {
       for (k = P->p[i]; k < P->p[i + 1]; k++) {
         /* diagonal element only */
         if (P->i[k] == i) { /* row == col */
+          /* M_ii += P_ii */
           M[i] += P->x[k];
           break;
         }
       }
     }
+    /* finally invert for pre-conditioner */
     M[i] = 1. / M[i];
   }
 #if VERBOSITY > 0
@@ -111,10 +116,18 @@ void SCS(free_lin_sys_work)(ScsLinSysWork *p) {
 }
 
 /* vec -> R_y^{-1} vec */
-static void scale_by_diag_r(scs_float *vec, ScsLinSysWork *p) {
+static void scale_by_r_y_inv(scs_float *vec, ScsLinSysWork *p) {
   scs_int i;
   for (i = 0; i < p->m; ++i) {
-    vec[i] /= p->rho_y_vec[i];
+    vec[i] /= p->diag_r[p->n + i];
+  }
+}
+
+/* y += R_x * x  */
+static void accum_by_r_x(scs_float *y, const scs_float *x, ScsLinSysWork *p) {
+  scs_int i;
+  for (i = 0; i < p->n; ++i) {
+    y[i] += p->diag_r[i] * x[i];
   }
 }
 
@@ -123,7 +136,7 @@ static void accum_by_a(ScsLinSysWork *p, const scs_float *x, scs_float *y) {
   SCS(accum_by_atrans)(p->At, x, y);
 }
 
-/* y = (rho_x * I + P + A' R_y^{-1} A) x */
+/* y = (R_x + P + A' R_y^{-1} A) x */
 static void mat_vec(const ScsMatrix *A, const ScsMatrix *P, ScsLinSysWork *p,
                     const scs_float *x, scs_float *y) {
   scs_float *z = p->tmp;
@@ -133,10 +146,10 @@ static void mat_vec(const ScsMatrix *A, const ScsMatrix *P, ScsLinSysWork *p,
     SCS(accum_by_p)(P, x, y); /* y = Px */
   }
   accum_by_a(p, x, z);           /* z = Ax */
-  scale_by_diag_r(z, p);         /* z = R_y^{-1} A x */
+  scale_by_r_y_inv(z, p);        /* z = R_y^{-1} A x */
   SCS(accum_by_atrans)(A, z, y); /* y += A'z, y = Px + A' R_y^{-1} Ax */
-  /* y = rho_x * x + Px + A' R_y^{-1} A x */
-  SCS(add_scaled_array)(y, x, A->n, p->rho_x);
+  /* y = R_x * x + Px + A' R_y^{-1} A * x */
+  accum_by_r_x(y, x, p);
 }
 
 static void apply_pre_conditioner(scs_float *z, scs_float *r, scs_int n,
@@ -149,36 +162,35 @@ static void apply_pre_conditioner(scs_float *z, scs_float *r, scs_int n,
 }
 
 /* no need to update anything in this case */
-void SCS(update_lin_sys_rho_y_vec)(ScsLinSysWork *p, scs_float *rho_y_vec) {
-  p->rho_y_vec = rho_y_vec; /* this isn't needed but do it to be safe */
+void SCS(update_lin_sys_diag_r)(ScsLinSysWork *p, const scs_float *diag_r) {
+  p->diag_r = diag_r; /* this isn't needed but do it to be safe */
   set_preconditioner(p);
 }
 
 ScsLinSysWork *SCS(init_lin_sys_work)(const ScsMatrix *A, const ScsMatrix *P,
-                                      scs_float *rho_y_vec, scs_float rho_x) {
+                                      const scs_float *diag_r) {
   ScsLinSysWork *p = (ScsLinSysWork *)scs_calloc(1, sizeof(ScsLinSysWork));
   p->A = A;
   p->P = P;
   p->m = A->m;
   p->n = A->n;
-  p->rho_x = rho_x;
 
-  p->p = (scs_float *)scs_malloc((A->n) * sizeof(scs_float));
-  p->r = (scs_float *)scs_malloc((A->n) * sizeof(scs_float));
-  p->Gp = (scs_float *)scs_malloc((A->n) * sizeof(scs_float));
-  p->tmp = (scs_float *)scs_malloc((A->m) * sizeof(scs_float));
+  p->p = (scs_float *)scs_calloc((A->n), sizeof(scs_float));
+  p->r = (scs_float *)scs_calloc((A->n), sizeof(scs_float));
+  p->Gp = (scs_float *)scs_calloc((A->n), sizeof(scs_float));
+  p->tmp = (scs_float *)scs_calloc((A->m), sizeof(scs_float));
 
   /* memory for A transpose */
-  p->At = (ScsMatrix *)scs_malloc(sizeof(ScsMatrix));
+  p->At = (ScsMatrix *)scs_calloc(1, sizeof(ScsMatrix));
   p->At->m = A->n;
   p->At->n = A->m;
-  p->At->i = (scs_int *)scs_malloc((A->p[A->n]) * sizeof(scs_int));
-  p->At->p = (scs_int *)scs_malloc((A->m + 1) * sizeof(scs_int));
-  p->At->x = (scs_float *)scs_malloc((A->p[A->n]) * sizeof(scs_float));
+  p->At->i = (scs_int *)scs_calloc((A->p[A->n]), sizeof(scs_int));
+  p->At->p = (scs_int *)scs_calloc((A->m + 1), sizeof(scs_int));
+  p->At->x = (scs_float *)scs_calloc((A->p[A->n]), sizeof(scs_float));
   transpose(A, p);
 
   /* preconditioner memory */
-  p->rho_y_vec = rho_y_vec;
+  p->diag_r = diag_r;
   p->z = (scs_float *)scs_calloc(A->n, sizeof(scs_float));
   p->M = (scs_float *)scs_calloc(A->n, sizeof(scs_float));
   set_preconditioner(p);
@@ -192,8 +204,7 @@ ScsLinSysWork *SCS(init_lin_sys_work)(const ScsMatrix *A, const ScsMatrix *P,
   return p;
 }
 
-/* solves (rho_x * I + P + A' R_y^{-1} A)x = b, s warm start, solution stored in
- * b */
+/* solves (R_x * I + P + A' R_y^{-1} A)x = b, s warm start, solution in b */
 static scs_int pcg(ScsLinSysWork *pr, const scs_float *s, scs_float *b,
                    scs_int max_its, scs_float tol) {
   scs_int i, n = pr->n;
@@ -268,14 +279,12 @@ static scs_int pcg(ScsLinSysWork *pr, const scs_float *s, scs_float *b,
 /* solves Mx = b, for x but stores result in b */
 /* s contains warm-start (if available) */
 /*
- * [x] = [rho_x I + P     A' ]^{-1} [rx]
- * [y]   [     A        -R_y ]      [ry]
- *
- * R_y = diag(rho_y_vec)
+ * [x] = [R_x + P     A' ]^{-1} [rx]
+ * [y]   [   A      -R_y ]      [ry]
  *
  * becomes:
  *
- * x = (rho_x I + P + A' R_y^{-1} A)^{-1} (rx + A' R_y^{-1} ry)
+ * x = (R_x + P + A' R_y^{-1} A)^{-1} (rx + A' R_y^{-1} ry)
  * y = R_y^{-1} (Ax - ry)
  *
  */
@@ -299,12 +308,12 @@ scs_int SCS(solve_lin_sys)(ScsLinSysWork *p, scs_float *b, const scs_float *s,
   /* tmp = ry */
   memcpy(p->tmp, &(b[p->n]), p->m * sizeof(scs_float));
   /* tmp = R_y^{-1} * ry */
-  scale_by_diag_r(p->tmp, p);
+  scale_by_r_y_inv(p->tmp, p);
   /* b[:n] = rx + A' R_y^{-1} ry */
   SCS(accum_by_atrans)(p->A, p->tmp, b);
   /* set max_iters to 10 * n (though in theory n is enough for any tol) */
   max_iters = 10 * p->n;
-  /* solves (rho_x I + P + A' R_y^{-1} A)x = b, s warm start, solution stored in
+  /* solves (R_x + P + A' R_y^{-1} A)x = b, s warm start, solution stored in
    * b */
   cg_its = pcg(p, s, b, max_iters, tol); /* b[:n] = x */
 
@@ -313,7 +322,7 @@ scs_int SCS(solve_lin_sys)(ScsLinSysWork *p, scs_float *b, const scs_float *s,
   /* b[n:] = Ax - ry */
   accum_by_a(p, b, &(b[p->n]));
   /* b[n:] = R_y^{-1} (Ax - ry) = y */
-  scale_by_diag_r(&(b[p->n]), p);
+  scale_by_r_y_inv(&(b[p->n]), p);
   p->tot_cg_its += cg_its;
 #if VERBOSITY > 1
   scs_printf("tol %.3e\n", tol);

data/vendor/scs/linsys/cpu/indirect/private.h

@@ -25,8 +25,7 @@ struct SCS_LIN_SYS_WORK {
   scs_float *M;
   /* reporting */
   scs_int tot_cg_its;
-  scs_float *rho_y_vec;
-  scs_float rho_x;
+  const scs_float *diag_r;
 };
 
 #ifdef __cplusplus

data/vendor/scs/linsys/csparse.c

@@ -12,9 +12,9 @@ csc *SCS(cs_spalloc)(scs_int m, scs_int n, scs_int nzmax, scs_int values,
   A->n = n;
   A->nzmax = nzmax = MAX(nzmax, 1);
   A->nz = triplet ? 0 : -1; /* allocate triplet or comp.col */
-  A->p = (scs_int *)scs_malloc((triplet ? nzmax : n + 1) * sizeof(scs_int));
-  A->i = (scs_int *)scs_malloc(nzmax * sizeof(scs_int));
-  A->x = values ? (scs_float *)scs_malloc(nzmax * sizeof(scs_float)) : SCS_NULL;
+  A->p = (scs_int *)scs_calloc((triplet ? nzmax : n + 1), sizeof(scs_int));
+  A->i = (scs_int *)scs_calloc(nzmax, sizeof(scs_int));
+  A->x = values ? (scs_float *)scs_calloc(nzmax, sizeof(scs_float)) : SCS_NULL;
   return (!A->p || !A->i || (values && !A->x)) ? SCS(cs_spfree)(A) : A;
 }
 

data/vendor/scs/linsys/external/amd/SuiteSparse_config.c

@@ -18,7 +18,6 @@
 
 #ifdef MATLAB_MEX_FILE
 #include "mex.h"
-#include "scs_matrix.h"
 #endif
 
 #ifndef NULL
@@ -51,7 +50,7 @@
 
 struct SuiteSparse_config_struct SuiteSparse_config =
 {
-    _scs_malloc, _scs_calloc, _scs_realloc, _scs_free, _scs_printf,
+    scs_malloc, scs_calloc, scs_realloc, scs_free, _scs_printf,
     SuiteSparse_hypot,
     SuiteSparse_divcomplex
 
@@ -73,12 +72,13 @@ struct SuiteSparse_config_struct SuiteSparse_config =
    SuiteSparse_start be called prior to calling any SuiteSparse function.
  */
 
+
 void SuiteSparse_start ( void )
 {
-    SuiteSparse_config.malloc_func  = _scs_malloc ;
-    SuiteSparse_config.calloc_func  = _scs_calloc ;
-    SuiteSparse_config.realloc_func = _scs_realloc ;
-    SuiteSparse_config.free_func    = _scs_free ;
+    SuiteSparse_config.malloc_func  = scs_malloc ;
+    SuiteSparse_config.calloc_func  = scs_calloc ;
+    SuiteSparse_config.realloc_func = scs_realloc ;
+    SuiteSparse_config.free_func    = scs_free ;
     SuiteSparse_config.printf_func  = _scs_printf ;
     /* math functions */
     SuiteSparse_config.hypot_func = SuiteSparse_hypot ;

data/vendor/scs/linsys/external/amd/SuiteSparse_config.h

@@ -44,7 +44,7 @@ extern "C" {
 
 #include <limits.h>
 #include <stdlib.h>
-#include "scs.h"
+#include "glbopts.h"
 #include "ctrlc.h"
 
 /* ========================================================================== */
@@ -71,6 +71,11 @@ extern "C" {
 #define SuiteSparse_long_id "%" SuiteSparse_long_idd
 #endif
 
+#ifndef _scs_printf
+#define _scs_printf scs_printf
+#endif
+
+
 /* ========================================================================== */
 /* === SuiteSparse_config parameters and functions ========================== */
 /* ========================================================================== */

data/vendor/scs/linsys/external/amd/amd_order.c

@@ -89,8 +89,8 @@ GLOBAL Int AMD_order
     }
 
     /* allocate two size-n integer workspaces */
-    Len  = SuiteSparse_malloc (n, sizeof (Int)) ;
-    Pinv = SuiteSparse_malloc (n, sizeof (Int)) ;
+    Len  = (Int *)SuiteSparse_malloc (n, sizeof (Int)) ;
+    Pinv = (Int *)SuiteSparse_malloc (n, sizeof (Int)) ;
     mem += n ;
     mem += n ;
     if (!Len || !Pinv)
@@ -106,8 +106,8 @@ GLOBAL Int AMD_order
     {
 	/* sort the input matrix and remove duplicate entries */
 	AMD_DEBUG1 (("Matrix is jumbled\n")) ;
-	Rp = SuiteSparse_malloc (n+1, sizeof (Int)) ;
-	Ri = SuiteSparse_malloc (nz,  sizeof (Int)) ;
+	Rp = (Int *)SuiteSparse_malloc (n+1, sizeof (Int)) ;
+	Ri = (Int *)SuiteSparse_malloc (nz,  sizeof (Int)) ;
 	mem += (n+1) ;
 	mem += MAX (nz,1) ;
 	if (!Rp || !Ri)
@@ -160,7 +160,7 @@ GLOBAL Int AMD_order
     ok = ok && (slen < Int_MAX) ;	/* S[i] for Int i must be OK */
     if (ok)
     {
-	S = SuiteSparse_malloc (slen, sizeof (Int)) ;
+	S = (Int *)SuiteSparse_malloc (slen, sizeof (Int)) ;
     }
     AMD_DEBUG1 (("slen %g\n", (scs_float) slen)) ;
     if (!S)

data/vendor/scs/linsys/gpu/gpu.h

@@ -1,12 +1,17 @@
-#ifndef SCSGPU_H_GUARD
-#define SCSGPU_H_GUARD
+#ifndef SCS_GPU_H_GUARD
+#define SCS_GPU_H_GUARD
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-#include <cublas_v2.h>
+/* TODO: Do we need this?
+
 #include <cuda.h>
+
+*/
+
+#include <cublas_v2.h>
 #include <cuda_runtime_api.h>
 #include <cusparse.h>
 
@@ -31,11 +36,9 @@ extern "C" {
 #ifndef SFLOAT
 #define CUBLAS(x) cublasD##x
 #define CUBLASI(x) cublasId##x
-#define CUSPARSE(x) cusparseD##x
 #else
 #define CUBLAS(x) cublasS##x
 #define CUBLASI(x) cublasIs##x
-#define CUSPARSE(x) cusparseS##x
 #endif
 #define CUSPARSE_GEN(x) cusparse##x
 #else
@@ -46,9 +49,6 @@ extern "C" {
 #define CUBLASI(x)                                                             \
   CUDA_CHECK_ERR;                                                              \
   cublasId##x
-#define CUSPARSE(x)                                                            \
-  CUDA_CHECK_ERR;                                                              \
-  cusparseD##x
 #else
 #define CUBLAS(x)                                                              \
   CUDA_CHECK_ERR;                                                              \
@@ -56,9 +56,6 @@ extern "C" {
 #define CUBLASI(x)                                                             \
   CUDA_CHECK_ERR;                                                              \
   cublasIs##x
-#define CUSPARSE(x)                                                            \
-  CUDA_CHECK_ERR;                                                              \
-  cusparseS##x
 #endif
 #define CUSPARSE_GEN(x)                                                        \
   CUDA_CHECK_ERR;                                                              \

data/vendor/scs/linsys/gpu/indirect/private.c

@@ -35,63 +35,77 @@ char *SCS(get_lin_sys_summary)(ScsLinSysWork *p, const ScsInfo *info) {
 }
 */
 
-/* set M = inv ( diag ( rho_x * I + P + A' R_y^{-1} A ) ) */
-static void set_preconditioner(ScsLinSysWork *p, scs_float *rho_y_vec) {
+/* Not possible to do this on the fly due to M_ii += a_i' (R_y)^-1 a_i */
+/* set M = inv ( diag ( R_x + P + A' R_y^{-1} A ) ) */
+static void set_preconditioner(ScsLinSysWork *p, const scs_float *diag_r) {
   scs_int i, k;
   const ScsMatrix *A = p->A;
   const ScsMatrix *P = p->P;
-  scs_float *M = (scs_float *)scs_calloc(A->n, sizeof(scs_float));
+  scs_float *M = p->M;
 
 #if VERBOSITY > 0
   scs_printf("getting pre-conditioner\n");
 #endif
 
+  /* M_ii = (R_x)_i + P_ii + a_i' (R_y)^-1 a_i */
   for (i = 0; i < A->n; ++i) { /* cols */
-    M[i] = p->rho_x;
-    /* diag(A' R_y^{-1} A) */
+    /* M_ii = (R_x)_i */
+    M[i] = diag_r[i];
+    /* M_ii += a_i' (R_y)^-1 a_i */
     for (k = A->p[i]; k < A->p[i + 1]; ++k) {
       /* A->i[k] is row of entry k with value A->x[k] */
-      M[i] += A->x[k] * A->x[k] / rho_y_vec[A->i[k]];
+      M[i] += A->x[k] * A->x[k] / diag_r[A->n + A->i[k]];
     }
     if (P) {
       for (k = P->p[i]; k < P->p[i + 1]; k++) {
         /* diagonal element only */
         if (P->i[k] == i) { /* row == col */
+          /* M_ii += P_ii */
           M[i] += P->x[k];
           break;
         }
       }
     }
+    /* finally invert for pre-conditioner */
     M[i] = 1. / M[i];
   }
-  cudaMemcpy(p->M, M, A->n * sizeof(scs_float), cudaMemcpyHostToDevice);
-  scs_free(M);
+  cudaMemcpy(p->M_gpu, M, A->n * sizeof(scs_float), cudaMemcpyHostToDevice);
 #if VERBOSITY > 0
   scs_printf("finished getting pre-conditioner\n");
 #endif
 }
 
 /* no need to update anything in this case */
-void SCS(update_lin_sys_rho_y_vec)(ScsLinSysWork *p, scs_float *rho_y_vec) {
+void SCS(update_lin_sys_diag_r)(ScsLinSysWork *p, const scs_float *diag_r) {
   scs_int i;
+
+  /* R_x to gpu */
+  cudaMemcpy(p->r_x_gpu, diag_r, p->n * sizeof(scs_float),
+             cudaMemcpyHostToDevice);
+
+  /* 1/R_y to gpu */
   for (i = 0; i < p->m; ++i)
-    p->inv_rho_y_vec[i] = 1. / rho_y_vec[i];
-  cudaMemcpy(p->inv_rho_y_vec_gpu, p->inv_rho_y_vec, p->m * sizeof(scs_float),
+    p->inv_r_y[i] = 1. / diag_r[p->n + i];
+  cudaMemcpy(p->inv_r_y_gpu, p->inv_r_y, p->m * sizeof(scs_float),
              cudaMemcpyHostToDevice);
-  set_preconditioner(p, rho_y_vec);
+
+  /* set preconditioner M on gpu */
+  set_preconditioner(p, diag_r);
 }
 
 void SCS(free_lin_sys_work)(ScsLinSysWork *p) {
   if (p) {
-    scs_free(p->inv_rho_y_vec);
+    scs_free(p->M);
+    scs_free(p->inv_r_y);
     cudaFree(p->p);
     cudaFree(p->r);
     cudaFree(p->Gp);
     cudaFree(p->bg);
     cudaFree(p->tmp_m);
     cudaFree(p->z);
-    cudaFree(p->M);
-    cudaFree(p->inv_rho_y_vec_gpu);
+    cudaFree(p->M_gpu);
+    cudaFree(p->r_x_gpu);
+    cudaFree(p->inv_r_y_gpu);
     if (p->Pg) {
       SCS(free_gpu_matrix)(p->Pg);
       scs_free(p->Pg);
@@ -126,22 +140,23 @@ static void scale_by_diag(cublasHandle_t cublas_handle, scs_float *M,
    0, M, 1, z, 1);
 }
 
-/* y = (rho_x * I + P + A' R_y^{-1} A) x */
+/* y = (R_x + P + A' R_y^{-1} A) x */
 static void mat_vec(ScsLinSysWork *p, const scs_float *x, scs_float *y) {
   /* x and y MUST already be loaded to GPU */
   scs_float *z = p->tmp_m; /* temp memory */
-  cudaMemset(y, 0, p->n * sizeof(scs_float));
   cudaMemset(z, 0, p->m * sizeof(scs_float));
 
   cusparseDnVecSetValues(p->dn_vec_m, (void *)z);
   cusparseDnVecSetValues(p->dn_vec_n, (void *)x);
   cusparseDnVecSetValues(p->dn_vec_n_p, (void *)y);
 
-  /* y = rho_x * x */
-  CUBLAS(axpy)(p->cublas_handle, p->n, &(p->rho_x), x, 1, y, 1);
+  /* y = x */
+  cudaMemcpy(y, x, p->n * sizeof(scs_float), cudaMemcpyHostToDevice);
+  /* y = R_x * x */
+  scale_by_diag(p->cublas_handle, p->r_x_gpu, y, p->n);
 
   if (p->Pg) {
-    /* y = rho_x * x + Px */
+    /* y = R_x * x + P x */
     SCS(accum_by_p_gpu)
     (p->Pg, p->dn_vec_n, p->dn_vec_n_p, p->cusparse_handle, &p->buffer_size,
      &p->buffer);
@@ -158,9 +173,9 @@ static void mat_vec(ScsLinSysWork *p, const scs_float *x, scs_float *y) {
    &p->buffer);
 #endif
   /* z = R_y^{-1} A x */
-  scale_by_diag(p->cublas_handle, p->inv_rho_y_vec_gpu, z, p->m);
+  scale_by_diag(p->cublas_handle, p->inv_r_y_gpu, z, p->m);
 
-  /* y += A'z => y = rho_x * x + Px + A' R_y^{-1} Ax */
+  /* y += A'z => y = R_x * x + P x + A' R_y^{-1} Ax */
   SCS(accum_by_atrans_gpu)
   (p->Ag, p->dn_vec_m, p->dn_vec_n_p, p->cusparse_handle, &p->buffer_size,
    &p->buffer);
@@ -201,19 +216,35 @@ static csc *fill_p_matrix(const ScsMatrix *P) {
 }
 
 ScsLinSysWork *SCS(init_lin_sys_work)(const ScsMatrix *A, const ScsMatrix *P,
-                                      scs_float *rho_y_vec, scs_float rho_x) {
+                                      const scs_float *diag_r) {
   cudaError_t err;
-  scs_int i;
   csc *P_full;
-  ScsLinSysWork *p = (ScsLinSysWork *)scs_calloc(1, sizeof(ScsLinSysWork));
-  ScsGpuMatrix *Ag = (ScsGpuMatrix *)scs_calloc(1, sizeof(ScsGpuMatrix));
+  ScsLinSysWork *p = SCS_NULL;
+  ScsGpuMatrix *Ag = SCS_NULL;
   ScsGpuMatrix *Pg = SCS_NULL;
+  int device_count;
+
+  err = cudaGetDeviceCount(&device_count);
+  if (err > 0) {
+    scs_printf("cudaError: %i (100 indicates no device)\n", (int)err);
+    return SCS_NULL;
+  }
+
+  p = (ScsLinSysWork *)scs_calloc(1, sizeof(ScsLinSysWork));
+  Ag = (ScsGpuMatrix *)scs_calloc(1, sizeof(ScsGpuMatrix));
+
+  p->inv_r_y = (scs_float *)scs_calloc(A->m, sizeof(scs_float));
+  p->M = (scs_float *)scs_calloc(A->n, sizeof(scs_float));
+
+  p->A = A;
+  p->P = P;
+  p->m = A->m;
+  p->n = A->n;
 
 #if GPU_TRANSPOSE_MAT > 0
   size_t new_buffer_size = 0;
 #endif
 
-  p->rho_x = rho_x;
   p->cublas_handle = 0;
   p->cusparse_handle = 0;
 
@@ -242,8 +273,9 @@ ScsLinSysWork *SCS(init_lin_sys_work)(const ScsMatrix *A, const ScsMatrix *P,
   cudaMalloc((void **)&p->bg, (A->n + A->m) * sizeof(scs_float));
   cudaMalloc((void **)&p->tmp_m, A->m * sizeof(scs_float));
   cudaMalloc((void **)&p->z, A->n * sizeof(scs_float));
-  cudaMalloc((void **)&p->M, A->n * sizeof(scs_float));
-  cudaMalloc((void **)&p->inv_rho_y_vec_gpu, A->m * sizeof(scs_float));
+  cudaMalloc((void **)&p->M_gpu, A->n * sizeof(scs_float));
+  cudaMalloc((void **)&p->r_x_gpu, A->n * sizeof(scs_float));
+  cudaMalloc((void **)&p->inv_r_y_gpu, A->m * sizeof(scs_float));
 
   cudaMemcpy(Ag->i, A->i, (A->p[A->n]) * sizeof(scs_int),
              cudaMemcpyHostToDevice);
@@ -251,12 +283,6 @@ ScsLinSysWork *SCS(init_lin_sys_work)(const ScsMatrix *A, const ScsMatrix *P,
   cudaMemcpy(Ag->x, A->x, (A->p[A->n]) * sizeof(scs_float),
              cudaMemcpyHostToDevice);
 
-  p->inv_rho_y_vec = (scs_float *)scs_malloc(A->m * sizeof(scs_float));
-  for (i = 0; i < A->m; ++i)
-    p->inv_rho_y_vec[i] = 1. / rho_y_vec[i];
-  cudaMemcpy(p->inv_rho_y_vec_gpu, p->inv_rho_y_vec, A->m * sizeof(scs_float),
-             cudaMemcpyHostToDevice);
-
   cusparseCreateCsr(&Ag->descr, Ag->n, Ag->m, Ag->nnz, Ag->p, Ag->i, Ag->x,
                     SCS_CUSPARSE_INDEX, SCS_CUSPARSE_INDEX,
                     CUSPARSE_INDEX_BASE_ZERO, SCS_CUDA_FLOAT);
@@ -297,7 +323,8 @@ ScsLinSysWork *SCS(init_lin_sys_work)(const ScsMatrix *A, const ScsMatrix *P,
   cusparseCreateDnVec(&p->dn_vec_n_p, Ag->n, p->tmp_m, SCS_CUDA_FLOAT);
   cusparseCreateDnVec(&p->dn_vec_m, Ag->m, p->tmp_m, SCS_CUDA_FLOAT);
 
-  set_preconditioner(p, rho_y_vec);
+  /* Form preconditioner and copy R_x, 1/R_y to gpu */
+  SCS(update_lin_sys_diag_r)(p, diag_r);
 
 #if GPU_TRANSPOSE_MAT > 0
   p->Agt = (ScsGpuMatrix *)scs_malloc(sizeof(ScsGpuMatrix));
@@ -346,9 +373,8 @@ ScsLinSysWork *SCS(init_lin_sys_work)(const ScsMatrix *A, const ScsMatrix *P,
   return p;
 }
 
-/* solves (rho_x * I + P + A' R_y^{-1} A)x = b, s warm start, solution stored in
- * b */
-/* on GPU */
+/* solves (R_x + P + A' R_y^{-1} A)x = b, s warm start, solution stored in
+ * b, on GPU */
 static scs_int pcg(ScsLinSysWork *pr, const scs_float *s, scs_float *bg,
                    scs_int max_its, scs_float tol) {
   scs_int i, n = pr->n;
@@ -386,7 +412,7 @@ static scs_int pcg(ScsLinSysWork *pr, const scs_float *s, scs_float *bg,
 
   /* z = M r */
   cudaMemcpy(z, r, n * sizeof(scs_float), cudaMemcpyDeviceToDevice);
-  scale_by_diag(cublas_handle, pr->M, z, n);
+  scale_by_diag(cublas_handle, pr->M_gpu, z, n);
   /* ztr = z'r */
   CUBLAS(dot)(cublas_handle, n, r, 1, z, 1, &ztr);
   /* p = z */
@@ -415,7 +441,7 @@ static scs_int pcg(ScsLinSysWork *pr, const scs_float *s, scs_float *bg,
     }
     /* z = M r */
     cudaMemcpy(z, r, n * sizeof(scs_float), cudaMemcpyDeviceToDevice);
-    scale_by_diag(cublas_handle, pr->M, z, n);
+    scale_by_diag(cublas_handle, pr->M_gpu, z, n);
     ztr_prev = ztr;
     /* ztr = z'r */
     CUBLAS(dot)(cublas_handle, n, r, 1, z, 1, &ztr);
@@ -431,14 +457,12 @@ static scs_int pcg(ScsLinSysWork *pr, const scs_float *s, scs_float *bg,
 /* solves Mx = b, for x but stores result in b */
 /* s contains warm-start (if available) */
 /*
- * [x] = [rho_x I + P     A' ]^{-1} [rx]
+ * [x] = [R_x + P        A' ]^{-1} [rx]
  * [y]   [     A        -R_y ]      [ry]
  *
- * R_y = diag(rho_y_vec)
- *
  * becomes:
  *
- * x = (rho_x I + P + A' R_y^{-1} A)^{-1} (rx + A' R_y^{-1} ry)
+ * x = (R_x + P + A' R_y^{-1} A)^{-1} (rx + A' R_y^{-1} ry)
  * y = R_y^{-1} (Ax - ry)
  *
  */
@@ -451,7 +475,6 @@ scs_int SCS(solve_lin_sys)(ScsLinSysWork *p, scs_float *b, const scs_float *s,
   scs_float *bg = p->bg;
   scs_float *tmp_m = p->tmp_m;
   ScsGpuMatrix *Ag = p->Ag;
-  ScsGpuMatrix *Pg = p->Pg;
 
   if (CG_NORM(b, p->n + p->m) <= 1e-12) {
     memset(b, 0, (p->n + p->m) * sizeof(scs_float));
@@ -471,7 +494,7 @@ scs_int SCS(solve_lin_sys)(ScsLinSysWork *p, scs_float *b, const scs_float *s,
   cudaMemcpy(tmp_m, &(bg[Ag->n]), Ag->m * sizeof(scs_float),
              cudaMemcpyDeviceToDevice);
   /* tmp = R_y^{-1} * tmp = R_y^{-1} * ry */
-  scale_by_diag(p->cublas_handle, p->inv_rho_y_vec_gpu, tmp_m, p->Ag->m);
+  scale_by_diag(p->cublas_handle, p->inv_r_y_gpu, tmp_m, p->Ag->m);
 
   cusparseDnVecSetValues(p->dn_vec_m, (void *)tmp_m); /* R * ry */
   cusparseDnVecSetValues(p->dn_vec_n, (void *)bg);    /* rx */
@@ -483,7 +506,7 @@ scs_int SCS(solve_lin_sys)(ScsLinSysWork *p, scs_float *b, const scs_float *s,
   /* set max_iters to 10 * n (though in theory n is enough for any tol) */
   max_iters = 10 * Ag->n;
 
-  /* solves (rho_x I + P + A' R_y^{-1} A)x = bg, s warm start, solution stored
+  /* solves (R_x + P + A' R_y^{-1} A)x = bg, s warm start, solution stored
    * in bg */
   cg_its = pcg(p, s, bg, max_iters, tol); /* bg[:n] = x */
 
@@ -504,7 +527,7 @@ scs_int SCS(solve_lin_sys)(ScsLinSysWork *p, scs_float *b, const scs_float *s,
 #endif
 
   /* bg[n:] = R_y^{-1} bg[n:] = R_y^{-1} (Ax - ry) = y */
-  scale_by_diag(p->cublas_handle, p->inv_rho_y_vec_gpu, &(bg[p->n]), p->Ag->m);
+  scale_by_diag(p->cublas_handle, p->inv_r_y_gpu, &(bg[p->n]), p->Ag->m);
 
   /* copy bg = [x; y] back to b */
   cudaMemcpy(b, bg, (Ag->n + Ag->m) * sizeof(scs_float),