Skip to content

Compact permuted_dense (no inverse-permutation arrays) for kron-path blocks #108

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Transurgeon wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Compact permuted_dense (no inverse-permutation arrays) for kron-path blocks #108

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 22 additions & 2 deletions include/utils/permuted_dense.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,14 @@ typedef struct permuted_dense
int n0;
int *row_perm;
int *col_perm;
int *col_inv; /* col_inv[col_perm[jj]] = jj, otherwise -1 */
int *row_inv; /* row_inv[row_perm[ii]] = ii, otherwise -1 */

/* Dense inverse permutations, sized by the GLOBAL dims (n and m):
col_inv[col_perm[jj]] = jj and row_inv[row_perm[ii]] = ii, -1
elsewhere. PDs built by new_permuted_dense_compact leave both NULL;
consumers either fall back to sorted scans of the perms or call
permuted_dense_ensure_col_inv / _row_inv to materialize on demand. */
int *col_inv;
int *row_inv;

/* Row-major block of size m0 x n0. Owned by this PD when owns_X == true,
and otherwise X is a view into a buffer (eg., stacked_pd's shared values
Expand Down Expand Up @@ -81,6 +87,20 @@ matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm,
col_perm = [0..n-1], dense block fills the full (m, n) shape. */
matrix *new_permuted_dense_full(int m, int n, const double *data);

/* Compact constructor: like new_permuted_dense, but leaves col_inv and
row_inv NULL instead of allocating the two global-dimension inverse
arrays — storage stays proportional to the block, not the global shape.
Used where those arrays dominate memory (the kron-path blocks, whose
global dims are the full variable space); consumers materialize them on
demand via the ensure helpers below, or scan the sorted perms instead. */
matrix *new_permuted_dense_compact(int m, int n, int m0, int n0, const int *row_perm,
const int *col_perm, const double *X_data);

/* Materialize A->col_inv / A->row_inv if NULL. Same const-cast on-demand
slot convention as permuted_dense_ensure_kernel_dwork. */
void permuted_dense_ensure_col_inv(const permuted_dense *A);
void permuted_dense_ensure_row_inv(const permuted_dense *A);

/* Ensure A->kernel_dwork is sized at least 'size' doubles. Grows in
place; contents are NOT preserved. */
void permuted_dense_ensure_kernel_dwork(const permuted_dense *A, size_t size);
Expand Down
10 changes: 10 additions & 0 deletions include/utils/utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,16 @@ int sorted_intersect_indices(const int *a, int a_len, const int *b, int b_len,
void sorted_union_int_arrays(const int *const *arrs, const int *lens, int n_arrs,
iVec *out);

/* Position of value g in the sorted, strictly-increasing array 'perm' of
length n0 (binary search), or -1 if absent. The scan replacement for a
dense inverse-permutation lookup inv[g]. */
int sorted_pos(const int *perm, int n0, int g);

/* Return true if any of the 'len' integers in 'idxs' (arbitrary order) is a
member of the sorted, strictly-increasing array 'perm' of length n0. The
scan replacement for idxs_hits_set when no dense inverse array exists. */
bool sorted_hits(const int *idxs, int len, const int *perm, int n0);

/* in-place cumulative sum */
void cumsum(int *p, int n);

Expand Down
114 changes: 93 additions & 21 deletions src/utils/permuted_dense.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,16 +113,44 @@ matrix *index_pd_alloc(const permuted_dense *A, const int *indices, int n_idxs)
/* Scan indices: which output positions i hit a row in A->row_perm? */
int *new_row_perm = (int *) sp_malloc(n_idxs * sizeof(int));
int new_m0 = 0;
if (A->row_inv != NULL)
{
for (int i = 0; i < n_idxs; i++)
{
if (A->row_inv[indices[i]] >= 0)
{
new_row_perm[new_m0++] = i;
}
}

matrix *out = new_permuted_dense(n_idxs, A->base.n, new_m0, A->n0,
new_row_perm, A->col_perm, NULL);
sp_free(new_row_perm);
return out;
}

/* Compact source (row_inv == NULL): membership via binary search in the sorted
row_perm. Record each hit's source dense row so index_pd_fill_values needs no
inverse array either — indices are owned by the atom and fixed
between alloc and fill, so the mapping stays valid. */
int *src_rows = (int *) sp_malloc(n_idxs * sizeof(int));
for (int i = 0; i < n_idxs; i++)
{
if (A->row_inv[indices[i]] >= 0)
int old_ii = sorted_pos(A->row_perm, A->m0, indices[i]);
if (old_ii >= 0)
{
src_rows[new_m0] = old_ii;
new_row_perm[new_m0++] = i;
}
}

matrix *out = new_permuted_dense(n_idxs, A->base.n, new_m0, A->n0, new_row_perm,
A->col_perm, NULL);
matrix *out = new_permuted_dense_compact(n_idxs, A->base.n, new_m0, A->n0,
new_row_perm, A->col_perm, NULL);
permuted_dense *C = (permuted_dense *) out;
C->kernel_iwork_size = (size_t) new_m0;
C->kernel_iwork = (int *) sp_malloc(new_m0 * sizeof(int));
memcpy(C->kernel_iwork, src_rows, new_m0 * sizeof(int));
sp_free(src_rows);
sp_free(new_row_perm);
return out;
}
Expand All @@ -132,11 +160,22 @@ void index_pd_fill_values(const permuted_dense *A, const int *indices, int n_idx
{
(void) n_idxs;
int n0 = A->n0;
if (A->row_inv != NULL)
{
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));
}
return;
}

/* Compact source (row_inv == NULL): C->kernel_iwork[k] holds the source dense
row for output dense row k, precomputed by index_pd_alloc. */
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->kernel_iwork[k] * n0, n0 * sizeof(double));
}
}

Expand Down Expand Up @@ -566,8 +605,9 @@ static void wire_vtable(permuted_dense *pd)
pd->base.refresh_csc_values = permuted_dense_refresh_csc_values;
}

matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm,
const int *col_perm, const double *X_data)
static matrix *new_permuted_dense_impl(int m, int n, int m0, int n0,
const int *row_perm, const int *col_perm,
const double *X_data, bool with_inv)
{
/* Validate sorted invariants. */
for (int ii = 1; ii < m0; ii++)
Expand Down Expand Up @@ -602,8 +642,6 @@ matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm,
pd->X = (double *) sp_malloc(sz * sizeof(double));
pd->base.x = pd->X;
pd->owns_X = true;
pd->col_inv = (int *) sp_malloc(n * sizeof(int));
pd->row_inv = (int *) sp_malloc(m * sizeof(int));

if (m0 > 0)
{
Expand All @@ -614,30 +652,64 @@ matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm,
memcpy(pd->col_perm, col_perm, n0 * sizeof(int));
}

if (with_inv)
{
permuted_dense_ensure_col_inv(pd);
permuted_dense_ensure_row_inv(pd);
}

if (X_data != NULL && sz > 0)
{
memcpy(pd->X, X_data, sz * sizeof(double));
}

return &pd->base;
}

matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm,
const int *col_perm, const double *X_data)
{
return new_permuted_dense_impl(m, n, m0, n0, row_perm, col_perm, X_data, true);
}

matrix *new_permuted_dense_compact(int m, int n, int m0, int n0, const int *row_perm,
const int *col_perm, const double *X_data)
{
return new_permuted_dense_impl(m, n, m0, n0, row_perm, col_perm, X_data, false);
}

void permuted_dense_ensure_col_inv(const permuted_dense *A)
{
/* const-cast: lazily populating a cache slot, same convention as
permuted_dense_ensure_kernel_dwork. */
permuted_dense *pd = (permuted_dense *) A;
if (pd->col_inv != NULL) return;
int n = pd->base.n;
pd->col_inv = (int *) sp_malloc(n * sizeof(int));
for (int j = 0; j < n; j++)
{
pd->col_inv[j] = -1;
}
for (int jj = 0; jj < n0; jj++)
for (int jj = 0; jj < pd->n0; jj++)
{
pd->col_inv[col_perm[jj]] = jj;
pd->col_inv[pd->col_perm[jj]] = jj;
}
}

void permuted_dense_ensure_row_inv(const permuted_dense *A)
{
permuted_dense *pd = (permuted_dense *) A;
if (pd->row_inv != NULL) return;
int m = pd->base.m;
pd->row_inv = (int *) sp_malloc(m * sizeof(int));
for (int i = 0; i < m; i++)
{
pd->row_inv[i] = -1;
}
for (int ii = 0; ii < m0; ii++)
for (int ii = 0; ii < pd->m0; ii++)
{
pd->row_inv[row_perm[ii]] = ii;
pd->row_inv[pd->row_perm[ii]] = ii;
}

if (X_data != NULL && sz > 0)
{
memcpy(pd->X, X_data, sz * sizeof(double));
}

return &pd->base;
}

matrix *new_permuted_dense_full(int m, int n, const double *data)
Expand Down
69 changes: 63 additions & 6 deletions src/utils/permuted_dense_linalg.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,15 +27,35 @@
#include <stdlib.h>
#include <string.h>

/* Compactness propagates through structural copies: a compact PD (no
inverse arrays — the kron-path blocks) produces compact copies and
transposes, so the global-dimension inv arrays never materialize anywhere
along its transpose/reshape/copy_sparsity chain. Sources with prebuilt
inv arrays produce full copies. */
static bool pd_is_compact(const permuted_dense *A)
{
return A->col_inv == NULL || A->row_inv == NULL;
}

matrix *copy_sparsity_pd_alloc(const permuted_dense *A)
{
if (pd_is_compact(A))
{
return new_permuted_dense_compact(A->base.m, A->base.n, A->m0, A->n0,
A->row_perm, A->col_perm, NULL);
}
return new_permuted_dense(A->base.m, A->base.n, A->m0, A->n0, A->row_perm,
A->col_perm, NULL);
}

matrix *transpose_pd_alloc(const permuted_dense *A)
{
/* swap (m, n), (m0, n0), and (row_perm, col_perm) */
if (pd_is_compact(A))
{
return new_permuted_dense_compact(A->base.n, A->base.m, A->n0, A->m0,
A->col_perm, A->row_perm, NULL);
}
return new_permuted_dense(A->base.n, A->base.m, A->n0, A->m0, A->col_perm,
A->row_perm, NULL);
}
Expand Down Expand Up @@ -265,14 +285,29 @@ matrix *BA_pd_csc_alloc(const permuted_dense *B, const CSC_matrix *A)
{
int start = A->p[j];
int len = A->p[j + 1] - start;
if (idxs_hits_set(A->i + start, len, B->col_inv))
bool hit = B->col_inv != NULL
? idxs_hits_set(A->i + start, len, B->col_inv)
: sorted_hits(A->i + start, len, B->col_perm, B->n0);
if (hit)
{
iVec_append(col_perm_C, j);
}
}

matrix *C = new_permuted_dense(B->base.m, A->n, B->m0, col_perm_C->len,
B->row_perm, col_perm_C->data, NULL);
/* The output is built compact: its inverse arrays would be sized by the
GLOBAL dims (m = B's global rows, n = A's full column count) and are
the dominant memory cost on the kron path, where this allocator runs
once per kron block. Every consumer of C either scans its sorted perms
or ensures the arrays on demand. */
matrix *C = new_permuted_dense_compact(B->base.m, A->n, B->m0, col_perm_C->len,
B->row_perm, col_perm_C->data, NULL);

/* BA_pd_csc_fill_values reads B->col_inv whenever C is non-empty (a
no-op when B already carries it, e.g. the kron scratch). */
if (col_perm_C->len > 0)
{
permuted_dense_ensure_col_inv(B);
}
iVec_free(col_perm_C);
return C;
}
Expand All @@ -283,7 +318,9 @@ void BA_pd_csc_fill_values(const double *B, int n0_B, const int *inv,
/* C[i, j] = bi^T @ ajj, where bi is the ith row of B_X (length n0_B,
row stride n0_B) and ajj is the jjth column of A's sparse block
(column jj = C->col_perm[j]). inv maps A's row indices to positions
in B_X (entries with inv[r] == -1 are skipped). */
in B_X (entries with inv[r] == -1 are skipped). inv may be NULL only
for an empty C (the alloc twin ensures it whenever C is non-empty). */
assert(inv != NULL || C->n0 == 0);

/* row i of C */
for (int i = 0; i < C->m0; i++)
Expand Down Expand Up @@ -399,14 +436,24 @@ matrix *BTA_pd_csc_alloc(const permuted_dense *B, const CSC_matrix *A)
{
int start = A->p[j];
int len = A->p[j + 1] - start;
if (idxs_hits_set(A->i + start, len, B->row_inv))
bool hit = B->row_inv != NULL
? idxs_hits_set(A->i + start, len, B->row_inv)
: sorted_hits(A->i + start, len, B->row_perm, B->m0);
if (hit)
{
iVec_append(col_active, j);
}
}

matrix *C = new_permuted_dense(B->base.n, A->n, B->n0, col_active->len,
B->col_perm, col_active->data, NULL);

/* BTDA_pd_csc_fill_values reads B->row_inv whenever C is non-empty; a
disjoint (empty) product never materializes it. */
if (col_active->len > 0)
{
permuted_dense_ensure_row_inv(B);
}
iVec_free(col_active);

/* Pre-size B's dwork for the BTDA fill (holds (diag(d) B)^T). */
Expand Down Expand Up @@ -453,14 +500,24 @@ matrix *BTA_csc_pd_alloc(const CSC_matrix *B, const permuted_dense *A)
{
int start = B->p[i];
int len = B->p[i + 1] - start;
if (idxs_hits_set(B->i + start, len, A->row_inv))
bool hit = A->row_inv != NULL
? idxs_hits_set(B->i + start, len, A->row_inv)
: sorted_hits(B->i + start, len, A->row_perm, A->m0);
if (hit)
{
iVec_append(row_active, i);
}
}

matrix *C = new_permuted_dense(B->n, A->base.n, row_active->len, A->n0,
row_active->data, A->col_perm, NULL);

/* BTDA_csc_pd_fill_values reads A->row_inv whenever C is non-empty; a
disjoint (empty) product never materializes it. */
if (row_active->len > 0)
{
permuted_dense_ensure_row_inv(A);
}
iVec_free(row_active);

/* Pre-size A's dwork for the BTDA fill (holds (diag(d_perm) X_A)^T). */
Expand Down
10 changes: 8 additions & 2 deletions src/utils/stacked_pd_coalesce.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -342,7 +342,10 @@ static inline void coalesce_spd_scatter(const stacked_pd *src, stacked_pd *out,
/* for each row in src_k */
for (int i = 0; i < src_k->m0; i++)
{
int out_i = out_k->row_inv[src_k->row_perm[i]];
int out_i =
out_k->row_inv != NULL
? out_k->row_inv[src_k->row_perm[i]]
: sorted_pos(out_k->row_perm, out_k->m0, src_k->row_perm[i]);
if (out_i < 0)
{
/* this row in src_k does not contribute to out_k */
Expand All @@ -355,7 +358,10 @@ static inline void coalesce_spd_scatter(const stacked_pd *src, stacked_pd *out,
/* place each col in src_row at its correct position in out_row */
for (int j = 0; j < src_k->n0; j++)
{
int out_j = out_k->col_inv[src_k->col_perm[j]];
int out_j = out_k->col_inv != NULL
? out_k->col_inv[src_k->col_perm[j]]
: sorted_pos(out_k->col_perm, out_k->n0,
src_k->col_perm[j]);
assert(out_j >= 0);
if (accumulate)
{
Expand Down
Loading
Loading