Remove "atom specific" functions from matrix classes - PR 1

include/utils/matrix.h

@@ -82,12 +82,12 @@ typedef CSR_matrix *(*matrix_to_csr_fn)(matrix *A);
    cache already matches values_version, so it is cheap to call when fresh. */
 typedef void (*matrix_refresh_csc_values_fn)(matrix *A);
 
-/* Allocate C = A[indices, :] */
-typedef matrix *(*matrix_index_alloc_fn)(matrix *A, const int *indices, int n_idxs);
+/* Allocate C = A[map, :]. C stores map internally, so the fill takes none. */
+typedef matrix *(*matrix_row_gather_alloc_fn)(const matrix *A, const int *map,
+                                              int m_out);
 
-/* Fill values of C = A[indices, :] */
-typedef void (*matrix_index_fill_values_fn)(matrix *A, const int *indices,
-                                            int n_idxs, matrix *C);
+/* Fill values of C = A[map, :] */
+typedef void (*matrix_row_gather_fill_values_fn)(const matrix *A, matrix *C);
 
 /* Row-tiling for the promote atom: A must be a 1-row matrix; returns
    a new matrix of shape (size, A->n) where every row is a copy of A's
@@ -156,8 +156,8 @@ struct matrix
     matrix_refresh_csc_values_fn refresh_csc_values;
 
     /* Atom-specific ops */
-    matrix_index_alloc_fn index_alloc;
-    matrix_index_fill_values_fn index_fill_values;
+    matrix_row_gather_alloc_fn row_gather_alloc;
+    matrix_row_gather_fill_values_fn row_gather_fill_values;
     matrix_promote_alloc_fn promote_alloc;
     matrix_promote_fill_values_fn promote_fill_values;
     matrix_broadcast_alloc_fn broadcast_alloc;

include/utils/permuted_dense.h

@@ -66,6 +66,12 @@ typedef struct permuted_dense
     int *kernel_iwork;
     size_t kernel_iwork_size;
 
+    /* Int state bound by the alloc that produced this PD and read by the
+       matching fill: row_gather_alloc stores, per dense row of this PD, the
+       dense row of the source it copies (length m0). NULL otherwise; never
+       touched by any other kernel (unlike kernel_iwork). */
+    int *bound_iwork;
+
     /* Cached transpose of this PD as another permuted_dense, allocated lazily
        on first call to permuted_dense_ensure_transpose_cache. On the cache PD
        itself, transpose_seen records the source's base.values_version whose
@@ -96,12 +102,12 @@ matrix *broadcast_pd_alloc(const permuted_dense *A, broadcast_type type, int d1,
 void broadcast_pd_fill_values(const permuted_dense *A, broadcast_type type, int d1,
                               int d2, permuted_dense *C);
 
-/* Allocate C = A[indices, :], where A and C are permuted dense. */
-matrix *index_pd_alloc(const permuted_dense *A, const int *indices, int n_idxs);
+/* Allocate C = A[map, :], where A and C are permuted dense. C stores map
+   internally, so the fill takes none. */
+matrix *row_gather_pd_alloc(const permuted_dense *A, const int *map, int m_out);
 
-/* Fill values of C = A[indices, :]. */
-void index_pd_fill_values(const permuted_dense *A, const int *indices, int n_idxs,
-                          permuted_dense *C);
+/* Fill values of C = A[map, :], where A and C are permuted dense. */
+void row_gather_pd_fill_values(const permuted_dense *A, permuted_dense *C);
 
 /* Allocate C = promote(A, size), where A and C are permuted dense. */
 matrix *promote_pd_alloc(const permuted_dense *A, int size);

include/utils/sparse_matrix.h

@@ -30,10 +30,12 @@ typedef struct sparse_matrix
     CSC_matrix *csc_cache;
     uint64_t csc_seen; /* base.values_version the csc_cache values reflect */
     int *csc_iwork;
-    int *transpose_iwork; /* sized csr->n; allocated by sparse_transpose_alloc
-                             on the output sm and reused by
-                             sparse_transpose_fill_values. NULL when this
-                             sm wasn't produced by transpose_alloc. */
+    /* Int state bound by the alloc that produced this matrix and read by the
+       matching fill: transpose_alloc stores csr->n scratch for
+       transpose_fill_values; row_gather_alloc stores the base.m-long row map
+       for row_gather_fill_values. NULL otherwise; never touched by any other
+       kernel. */
+    int *bound_iwork;
 } sparse_matrix;
 
 /* Constructor. Takes ownership of A; the caller must not free A separately

src/atoms/affine/index.c

@@ -65,20 +65,19 @@ static void jacobian_init_impl(expr *node)
     jacobian_init(x);
 
     /* allocate sparsity pattern for the matrix consisting of rows
-       'idx->indices' of the child's Jacobian */
+       'idx->indices' of the child's Jacobian; the gather map is bound to
+       node->jacobian, so eval only needs the two matrices */
     node->jacobian =
-        x->jacobian->index_alloc(x->jacobian, idx->indices, idx->n_idxs);
+        x->jacobian->row_gather_alloc(x->jacobian, idx->indices, idx->n_idxs);
 }
 
 static void eval_jacobian_impl(expr *node)
 {
     expr *x = node->left;
-    index_expr *idx = (index_expr *) node;
     eval_jacobian(x);
 
     /* copy values of the selected rows into the preallocated output */
-    x->jacobian->index_fill_values(x->jacobian, idx->indices, idx->n_idxs,
-                                   node->jacobian);
+    x->jacobian->row_gather_fill_values(x->jacobian, node->jacobian);
 }
 
 static void wsum_hess_init_impl(expr *node)

src/atoms/affine/transpose.c

@@ -50,18 +50,18 @@ static void jacobian_init_impl(expr *node)
         indices[r] = (r / d1) + (r % d1) * d2;
     }
 
-    node->jacobian = child->jacobian->index_alloc(child->jacobian, indices, n_out);
+    node->jacobian =
+        child->jacobian->row_gather_alloc(child->jacobian, indices, n_out);
 
-    /* save indices for eval_jacobian */
-    node->work->iwork = indices;
+    /* the gather map is bound to node->jacobian; nothing to keep */
+    sp_free(indices);
 }
 
 static void eval_jacobian_impl(expr *node)
 {
     expr *child = node->left;
     eval_jacobian(child);
-    child->jacobian->index_fill_values(child->jacobian, node->work->iwork,
-                                       node->size, node->jacobian);
+    child->jacobian->row_gather_fill_values(child->jacobian, node->jacobian);
 }
 
 static void wsum_hess_init_impl(expr *node)

src/utils/permuted_dense.c

@@ -23,6 +23,7 @@
 #include "utils/tracked_alloc.h"
 #include "utils/utils.h"
 #include <assert.h>
+#include <limits.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -47,6 +48,7 @@ static void permuted_dense_free(matrix *self)
     }
     sp_free(pd->kernel_dwork);
     sp_free(pd->kernel_iwork);
+    sp_free(pd->bound_iwork);
     free_matrix((matrix *) pd->transpose_cache);
     sp_free(pd);
 }
@@ -108,49 +110,58 @@ static void permuted_dense_vtable_transpose_fill_values(const matrix *self,
     transpose_pd_fill_values((const permuted_dense *) self, (permuted_dense *) out);
 }
 
-matrix *index_pd_alloc(const permuted_dense *A, const int *indices, int n_idxs)
+matrix *row_gather_pd_alloc(const permuted_dense *A, const int *map, int m_out)
 {
-    /* Scan indices: which output positions i hit a row in A->row_perm? */
-    int *new_row_perm = (int *) sp_malloc(n_idxs * sizeof(int));
+    /* Output position i is dense iff map[i] hits a row in A->row_perm. The kept
+       positions form C's row_perm (strictly increasing by construction, repeats
+       in map included); src[k] is the dense row of A that C's row k copies. */
+    int *new_row_perm = (int *) sp_malloc(m_out * sizeof(int));
+    int *src = (int *) sp_malloc(m_out * sizeof(int));
     int new_m0 = 0;
-    for (int i = 0; i < n_idxs; i++)
+    for (int i = 0; i < m_out; i++)
     {
-        if (A->row_inv[indices[i]] >= 0)
+        int ii = A->row_inv[map[i]];
+        if (ii >= 0)
         {
-            new_row_perm[new_m0++] = i;
+            new_row_perm[new_m0] = i;
+            src[new_m0] = ii;
+            new_m0++;
         }
     }
 
-    matrix *out = new_permuted_dense(n_idxs, A->base.n, new_m0, A->n0, new_row_perm,
+    matrix *out = new_permuted_dense(m_out, A->base.n, new_m0, A->n0, new_row_perm,
                                      A->col_perm, NULL);
+    if (new_m0 > 0)
+    {
+        permuted_dense *C = (permuted_dense *) out;
+        C->bound_iwork = (int *) sp_malloc(new_m0 * sizeof(int));
+        memcpy(C->bound_iwork, src, new_m0 * sizeof(int));
+    }
     sp_free(new_row_perm);
+    sp_free(src);
     return out;
 }
 
-void index_pd_fill_values(const permuted_dense *A, const int *indices, int n_idxs,
-                          permuted_dense *C)
+void row_gather_pd_fill_values(const permuted_dense *A, permuted_dense *C)
 {
-    (void) n_idxs;
+    assert(A->base.n == C->base.n && (C->m0 == 0 || C->bound_iwork != NULL));
     int n0 = A->n0;
     for (int k = 0; k < C->m0; k++)
     {
-        int i = C->row_perm[k];
-        int old_ii = A->row_inv[indices[i]];
-        memcpy(C->X + k * n0, A->X + old_ii * n0, n0 * sizeof(double));
+        memcpy(C->X + k * n0, A->X + C->bound_iwork[k] * n0, n0 * sizeof(double));
     }
 }
 
-static matrix *permuted_dense_vtable_index_alloc(matrix *self, const int *indices,
-                                                 int n_idxs)
+static matrix *permuted_dense_vtable_row_gather_alloc(const matrix *self,
+                                                      const int *map, int m_out)
 {
-    return index_pd_alloc((const permuted_dense *) self, indices, n_idxs);
+    return row_gather_pd_alloc((const permuted_dense *) self, map, m_out);
 }
 
-static void permuted_dense_vtable_index_fill_values(matrix *self, const int *indices,
-                                                    int n_idxs, matrix *out)
+static void permuted_dense_vtable_row_gather_fill_values(const matrix *self,
+                                                         matrix *out)
 {
-    index_pd_fill_values((const permuted_dense *) self, indices, n_idxs,
-                         (permuted_dense *) out);
+    row_gather_pd_fill_values((const permuted_dense *) self, (permuted_dense *) out);
 }
 
 matrix *promote_pd_alloc(const permuted_dense *A, int size)
@@ -554,8 +565,8 @@ static void wire_vtable(permuted_dense *pd)
     pd->base.to_csr = permuted_dense_to_csr;
     pd->base.transpose_alloc = permuted_dense_vtable_transpose_alloc;
     pd->base.transpose_fill_values = permuted_dense_vtable_transpose_fill_values;
-    pd->base.index_alloc = permuted_dense_vtable_index_alloc;
-    pd->base.index_fill_values = permuted_dense_vtable_index_fill_values;
+    pd->base.row_gather_alloc = permuted_dense_vtable_row_gather_alloc;
+    pd->base.row_gather_fill_values = permuted_dense_vtable_row_gather_fill_values;
     pd->base.promote_alloc = permuted_dense_vtable_promote_alloc;
     pd->base.promote_fill_values = permuted_dense_vtable_promote_fill_values;
     pd->base.broadcast_alloc = permuted_dense_vtable_broadcast_alloc;
@@ -590,6 +601,12 @@ matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm,
     permuted_dense *pd = (permuted_dense *) sp_calloc(1, sizeof(permuted_dense));
     pd->base.m = m;
     pd->base.n = n;
+    if (m0 > 0 && n0 > INT_MAX / m0)
+    {
+        fprintf(stderr, "Error in new_permuted_dense: dense block m0 * n0 "
+                        "exceeds INT_MAX.\n");
+        exit(1);
+    }
     pd->base.nnz = m0 * n0;
     wire_vtable(pd);
 

src/utils/sparse_matrix.c

@@ -24,6 +24,7 @@
 #include "utils/mini_numpy.h"
 #include "utils/tracked_alloc.h"
 #include "utils/utils.h"
+#include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -55,7 +56,7 @@ static void sparse_free(matrix *self)
     free_CSR_matrix(sm->csr);
     free_CSC_matrix(sm->csc_cache);
     sp_free(sm->csc_iwork);
-    sp_free(sm->transpose_iwork);
+    sp_free(sm->bound_iwork);
     sp_free(sm);
 }
 
@@ -117,60 +118,70 @@ static matrix *sparse_transpose_alloc(const matrix *self)
     int *iwork = (int *) sp_malloc(sm->csr->n * sizeof(int));
     CSR_matrix *AT = AT_alloc(sm->csr, iwork);
     sparse_matrix *out = (sparse_matrix *) new_sparse_matrix(AT);
-    out->transpose_iwork = iwork;
+    out->bound_iwork = iwork;
     return &out->base;
 }
 
 static void sparse_transpose_fill_values(const matrix *self, matrix *out)
 {
     const sparse_matrix *sm_in = (const sparse_matrix *) self;
     sparse_matrix *sm_out = (sparse_matrix *) out;
-    AT_fill_values(sm_in->csr, sm_out->csr, sm_out->transpose_iwork);
+    AT_fill_values(sm_in->csr, sm_out->csr, sm_out->bound_iwork);
 }
 
-static matrix *sparse_index_alloc(matrix *self, const int *indices, int n_idxs)
+static matrix *sparse_row_gather_alloc(const matrix *self, const int *map, int m_out)
 {
-    CSR_matrix *Jx = ((sparse_matrix *) self)->csr;
+    const CSR_matrix *Jx = ((const sparse_matrix *) self)->csr;
 
     /* Exact output nnz: sum the selected rows' nnz. Jx->nnz is NOT an upper
-       bound — duplicated indices select the same source row more than once
-       (cvxpy#3442). Duplicated gathers of dense rows can push the true count
+       bound — repeated map entries select the same source row more than once
+       (cvxpy#3442). Repeated gathers of dense rows can push the true count
        past INT_MAX, which a CSR cannot represent, so fail before wrapping. */
     int nnz = 0;
-    for (int i = 0; i < n_idxs; i++)
+    for (int i = 0; i < m_out; i++)
     {
-        int len = Jx->p[indices[i] + 1] - Jx->p[indices[i]];
+        int row = map[i];
+        int len = Jx->p[row + 1] - Jx->p[row];
         if (len > INT_MAX - nnz)
         {
-            fprintf(stderr, "Error in sparse_index_alloc: gathered nnz "
+            fprintf(stderr, "Error in sparse_row_gather_alloc: gathered nnz "
                             "exceeds INT_MAX.\n");
             exit(1);
         }
         nnz += len;
     }
-    CSR_matrix *J = new_CSR_matrix(n_idxs, self->n, nnz);
+    CSR_matrix *J = new_CSR_matrix(m_out, self->n, nnz);
 
     J->p[0] = 0;
-    for (int i = 0; i < n_idxs; i++)
+    for (int i = 0; i < m_out; i++)
     {
-        int row = indices[i];
+        int row = map[i];
         int len = Jx->p[row + 1] - Jx->p[row];
         memcpy(J->i + J->p[i], Jx->i + Jx->p[row], len * sizeof(int));
         J->p[i + 1] = J->p[i] + len;
     }
-    J->nnz = J->p[n_idxs];
-    return new_sparse_matrix(J);
+    J->nnz = J->p[m_out];
+
+    sparse_matrix *out = (sparse_matrix *) new_sparse_matrix(J);
+    out->bound_iwork = (int *) sp_malloc(m_out * sizeof(int));
+    if (m_out > 0)
+    {
+        memcpy(out->bound_iwork, map, m_out * sizeof(int));
+    }
+    return &out->base;
 }
 
-static void sparse_index_fill_values(matrix *self, const int *indices, int n_idxs,
-                                     matrix *out)
+static void sparse_row_gather_fill_values(const matrix *self, matrix *out)
 {
-    CSR_matrix *Jx = ((sparse_matrix *) self)->csr;
-    CSR_matrix *J = ((sparse_matrix *) out)->csr;
-    for (int i = 0; i < n_idxs; i++)
+    const CSR_matrix *Jx = ((const sparse_matrix *) self)->csr;
+    sparse_matrix *sm_out = (sparse_matrix *) out;
+    CSR_matrix *J = sm_out->csr;
+    const int *map = sm_out->bound_iwork;
+    assert(map != NULL && self->n == out->n);
+    for (int i = 0; i < J->m; i++)
     {
         int len = J->p[i + 1] - J->p[i];
-        memcpy(J->x + J->p[i], Jx->x + Jx->p[indices[i]], len * sizeof(double));
+        memcpy(J->x + J->p[i], Jx->x + Jx->p[map[i]], len * sizeof(double));
     }
 }
 
@@ -389,8 +400,8 @@ static void wire_vtable(sparse_matrix *sm)
     sm->base.to_csr = sparse_to_csr;
     sm->base.transpose_alloc = sparse_transpose_alloc;
     sm->base.transpose_fill_values = sparse_transpose_fill_values;
-    sm->base.index_alloc = sparse_index_alloc;
-    sm->base.index_fill_values = sparse_index_fill_values;
+    sm->base.row_gather_alloc = sparse_row_gather_alloc;
+    sm->base.row_gather_fill_values = sparse_row_gather_fill_values;
     sm->base.promote_alloc = sparse_promote_alloc;
     sm->base.promote_fill_values = sparse_promote_fill_values;
     sm->base.broadcast_alloc = sparse_broadcast_alloc;