first draft of transpose

Author: dance858

include/utils/stacked_pd.h

@@ -39,6 +39,11 @@ typedef struct stacked_pd
     permuted_dense **blocks; /* owned; length n_blocks                      */
     int *src_block_idx_p;    /* owned; length n_blocks + 1                  */
     int *src_block_idx;      /* owned; length src_block_idx_p[n_blocks]     */
+    struct stacked_pd *work; /* owned; NULL except when this spd was        */
+                             /* produced by an op that needs a pre-coalesce */
+                             /* intermediate (currently transpose). When    */
+                             /* set, holds the raw spd that                 */
+                             /* coalesce_spd_fill_values reads from.        */
 } stacked_pd;
 
 /* Constructor.
@@ -100,4 +105,16 @@ matrix *coalesce_spd_alloc(const stacked_pd *src);
    (block count, row_perms, col_perms) as the one passed now. */
 void coalesce_spd_fill_values(const stacked_pd *src, stacked_pd *out);
 
+/* Allocate out = transpose(src). Implementation: transpose each source
+   PD block individually (yielding a raw spd that may have overlapping
+   row perms but pairwise-disjoint cells, by the spd invariant), then
+   coalesce. The raw spd is stored on `out->work` for reuse by
+   `transpose_spd_fill_values`. */
+matrix *transpose_spd_alloc(const stacked_pd *src);
+
+/* Fill values of out = transpose(src). The caller must have produced
+   `out` via transpose_spd_alloc on a `src` with the same structure
+   (block count, row_perms, col_perms) as the one passed now. */
+void transpose_spd_fill_values(const stacked_pd *src, stacked_pd *out);
+
 #endif /* STACKED_PD_H */

src/utils/stacked_pd.c

@@ -36,6 +36,10 @@ static void stacked_pd_free(matrix *self)
     free(spd->blocks);
     free(spd->src_block_idx_p);
     free(spd->src_block_idx);
+    if (spd->work != NULL)
+    {
+        free_matrix((matrix *) spd->work);
+    }
     free(spd);
 }
 
@@ -488,3 +492,41 @@ void coalesce_spd_fill_values(const stacked_pd *src, stacked_pd *out)
         }
     }
 }
+
+/* ------------------------------------------------------------------ */
+/* Transpose                                                           */
+/* ------------------------------------------------------------------ */
+
+matrix *transpose_spd_alloc(const stacked_pd *src)
+{
+    int n_blocks = src->n_blocks;
+    permuted_dense **raw_blocks = (permuted_dense **) SP_MALLOC(
+        (n_blocks > 0 ? n_blocks : 1) * sizeof(permuted_dense *));
+    for (int k = 0; k < n_blocks; k++)
+    {
+        matrix *blk = (matrix *) src->blocks[k];
+        raw_blocks[k] = (permuted_dense *) blk->transpose_alloc(blk);
+    }
+
+    /* Raw spd: dimensions swapped, identity src_block_idx_*; rows may
+       overlap so use the unchecked constructor. */
+    matrix *raw = new_stacked_pd_unchecked(src->base.n, src->base.m, n_blocks,
+                                           raw_blocks, NULL, NULL);
+    free(raw_blocks);
+
+    matrix *out = coalesce_spd_alloc((stacked_pd *) raw);
+    ((stacked_pd *) out)->work = (stacked_pd *) raw;
+    return out;
+}
+
+void transpose_spd_fill_values(const stacked_pd *src, stacked_pd *out)
+{
+    stacked_pd *raw = out->work;
+    for (int k = 0; k < src->n_blocks; k++)
+    {
+        matrix *src_blk = (matrix *) src->blocks[k];
+        matrix *raw_blk = (matrix *) raw->blocks[k];
+        src_blk->transpose_fill_values(src_blk, raw_blk);
+    }
+    coalesce_spd_fill_values(raw, out);
+}

tests/all_tests.c

@@ -413,6 +413,11 @@ int main(void)
     mu_run_test(test_coalesce_shared_signature_merges_rows, tests_run);
     mu_run_test(test_coalesce_alloc_then_fill_values, tests_run);
     mu_run_test(test_coalesce_empty_input, tests_run);
+    mu_run_test(test_transpose_spd_no_overlap, tests_run);
+    mu_run_test(test_transpose_spd_overlap_coalesces, tests_run);
+    mu_run_test(test_transpose_spd_alloc_then_fill_values, tests_run);
+    mu_run_test(test_transpose_spd_empty, tests_run);
+    mu_run_test(test_transpose_spd_single_block_full, tests_run);
     mu_run_test(test_YT_kron_I, tests_run);
     mu_run_test(test_YT_kron_I_larger, tests_run);
     mu_run_test(test_I_kron_X, tests_run);

tests/utils/test_stacked_pd.h

@@ -462,4 +462,232 @@ const char *test_coalesce_empty_input(void)
     return 0;
 }
 
+/* Transpose where source blocks have disjoint rows AND disjoint cols.
+   Transposed blocks also have disjoint rows; coalesce leaves them
+   unchanged. */
+const char *test_transpose_spd_no_overlap(void)
+{
+    /* block 0: rows {0, 1}, cols {0, 1}, X = [[1, 2], [3, 4]]
+       block 1: rows {2, 3}, cols {2, 3}, X = [[5, 6], [7, 8]]          */
+    int row_perm_0[2] = {0, 1};
+    int col_perm_0[2] = {0, 1};
+    double X0[4] = {1.0, 2.0, 3.0, 4.0};
+    matrix *blk0 = new_permuted_dense(4, 4, 2, 2, row_perm_0, col_perm_0, X0);
+
+    int row_perm_1[2] = {2, 3};
+    int col_perm_1[2] = {2, 3};
+    double X1[4] = {5.0, 6.0, 7.0, 8.0};
+    matrix *blk1 = new_permuted_dense(4, 4, 2, 2, row_perm_1, col_perm_1, X1);
+
+    permuted_dense *blocks[2] = {(permuted_dense *) blk0, (permuted_dense *) blk1};
+    matrix *src_m = new_stacked_pd(4, 4, 2, blocks, NULL, NULL);
+
+    matrix *out_m = transpose_spd_alloc((stacked_pd *) src_m);
+    transpose_spd_fill_values((stacked_pd *) src_m, (stacked_pd *) out_m);
+    stacked_pd *out = (stacked_pd *) out_m;
+
+    mu_assert("m", out_m->m == 4);
+    mu_assert("n", out_m->n == 4);
+    mu_assert("n_blocks", out->n_blocks == 2);
+    int expected_src_p[3] = {0, 1, 2};
+    int expected_src[2] = {0, 1};
+    mu_assert("src_block_idx_p",
+              cmp_int_array(out->src_block_idx_p, expected_src_p, 3));
+    mu_assert("src_block_idx", cmp_int_array(out->src_block_idx, expected_src, 2));
+
+    /* Transpose of [[1,2],[3,4]] is [[1,3],[2,4]] (row-major flatten:
+       [1, 3, 2, 4]); similarly for the second block. */
+    permuted_dense *O0 = out->blocks[0];
+    double O0_X_expected[4] = {1.0, 3.0, 2.0, 4.0};
+    mu_assert("O0 row_perm", cmp_int_array(O0->row_perm, row_perm_0, 2));
+    mu_assert("O0 col_perm", cmp_int_array(O0->col_perm, col_perm_0, 2));
+    mu_assert("O0 X", cmp_double_array(O0->X, O0_X_expected, 4));
+
+    permuted_dense *O1 = out->blocks[1];
+    double O1_X_expected[4] = {5.0, 7.0, 6.0, 8.0};
+    mu_assert("O1 row_perm", cmp_int_array(O1->row_perm, row_perm_1, 2));
+    mu_assert("O1 col_perm", cmp_int_array(O1->col_perm, col_perm_1, 2));
+    mu_assert("O1 X", cmp_double_array(O1->X, O1_X_expected, 4));
+
+    free_matrix(out_m);
+    free_matrix(src_m);
+    return 0;
+}
+
+/* Transpose where source blocks share a column. After per-block
+   transpose two raw blocks share a row, so coalesce merges and produces
+   three output PDs (the worked example from the design conversation). */
+const char *test_transpose_spd_overlap_coalesces(void)
+{
+    /* Source (4x5):
+         block 0: rows {0, 1}, cols {2, 4}, X = [[1, 2], [3, 4]]
+         block 1: rows {3},    cols {0, 2}, X = [10, 11]
+       (Disjoint rows; cells {(0,2),(0,4),(1,2),(1,4)} and {(3,0),(3,2)}
+       are disjoint.)                                                    */
+    int row_perm_0[2] = {0, 1};
+    int col_perm_0[2] = {2, 4};
+    double X0[4] = {1.0, 2.0, 3.0, 4.0};
+    matrix *blk0 = new_permuted_dense(4, 5, 2, 2, row_perm_0, col_perm_0, X0);
+
+    int row_perm_1[1] = {3};
+    int col_perm_1[2] = {0, 2};
+    double X1[2] = {10.0, 11.0};
+    matrix *blk1 = new_permuted_dense(4, 5, 1, 2, row_perm_1, col_perm_1, X1);
+
+    permuted_dense *blocks[2] = {(permuted_dense *) blk0, (permuted_dense *) blk1};
+    matrix *src_m = new_stacked_pd(4, 5, 2, blocks, NULL, NULL);
+
+    matrix *out_m = transpose_spd_alloc((stacked_pd *) src_m);
+    transpose_spd_fill_values((stacked_pd *) src_m, (stacked_pd *) out_m);
+    stacked_pd *out = (stacked_pd *) out_m;
+
+    /* Transposed (5x4):
+         raw block 0^T: rows {2, 4}, cols {0, 1}, X = [[1, 3], [2, 4]]
+         raw block 1^T: rows {0, 2}, cols {3},    X = [10, 11]
+       Row 2 shared. Coalesce produces (ordered by min row):
+         sig {1}:    row {0},  col {3},        X = [10]
+         sig {0, 1}: row {2},  cols {0, 1, 3}, X = [1, 3, 11]
+         sig {0}:    row {4},  cols {0, 1},    X = [2, 4]               */
+    mu_assert("m", out_m->m == 5);
+    mu_assert("n", out_m->n == 4);
+    mu_assert("n_blocks", out->n_blocks == 3);
+
+    int expected_src_p[4] = {0, 1, 3, 4};
+    int expected_src[4] = {1, 0, 1, 0};
+    mu_assert("src_block_idx_p",
+              cmp_int_array(out->src_block_idx_p, expected_src_p, 4));
+    mu_assert("src_block_idx", cmp_int_array(out->src_block_idx, expected_src, 4));
+
+    permuted_dense *O0 = out->blocks[0];
+    int O0_row[1] = {0};
+    int O0_col[1] = {3};
+    double O0_X[1] = {10.0};
+    mu_assert("O0 row", cmp_int_array(O0->row_perm, O0_row, 1));
+    mu_assert("O0 col", cmp_int_array(O0->col_perm, O0_col, 1));
+    mu_assert("O0 X", cmp_double_array(O0->X, O0_X, 1));
+
+    permuted_dense *O1 = out->blocks[1];
+    int O1_row[1] = {2};
+    int O1_col[3] = {0, 1, 3};
+    double O1_X[3] = {1.0, 3.0, 11.0};
+    mu_assert("O1 row", cmp_int_array(O1->row_perm, O1_row, 1));
+    mu_assert("O1 col", cmp_int_array(O1->col_perm, O1_col, 3));
+    mu_assert("O1 X", cmp_double_array(O1->X, O1_X, 3));
+
+    permuted_dense *O2 = out->blocks[2];
+    int O2_row[1] = {4};
+    int O2_col[2] = {0, 1};
+    double O2_X[2] = {2.0, 4.0};
+    mu_assert("O2 row", cmp_int_array(O2->row_perm, O2_row, 1));
+    mu_assert("O2 col", cmp_int_array(O2->col_perm, O2_col, 2));
+    mu_assert("O2 X", cmp_double_array(O2->X, O2_X, 2));
+
+    free_matrix(out_m);
+    free_matrix(src_m);
+    return 0;
+}
+
+/* Two-phase: alloc once, mutate src values, refill, re-check. */
+const char *test_transpose_spd_alloc_then_fill_values(void)
+{
+    int row_perm_0[2] = {0, 1};
+    int col_perm_0[2] = {2, 4};
+    double X0[4] = {1.0, 2.0, 3.0, 4.0};
+    matrix *blk0 = new_permuted_dense(4, 5, 2, 2, row_perm_0, col_perm_0, X0);
+
+    int row_perm_1[1] = {3};
+    int col_perm_1[2] = {0, 2};
+    double X1[2] = {10.0, 11.0};
+    matrix *blk1 = new_permuted_dense(4, 5, 1, 2, row_perm_1, col_perm_1, X1);
+
+    permuted_dense *blocks[2] = {(permuted_dense *) blk0, (permuted_dense *) blk1};
+    matrix *src_m = new_stacked_pd(4, 5, 2, blocks, NULL, NULL);
+    stacked_pd *src = (stacked_pd *) src_m;
+
+    matrix *out_m = transpose_spd_alloc(src);
+    transpose_spd_fill_values(src, (stacked_pd *) out_m);
+
+    /* Mutate src values then refill:
+         blk0 X -> [100, 200, 300, 400]
+         blk1 X -> [50, 60]                                             */
+    src->blocks[0]->X[0] = 100.0;
+    src->blocks[0]->X[1] = 200.0;
+    src->blocks[0]->X[2] = 300.0;
+    src->blocks[0]->X[3] = 400.0;
+    src->blocks[1]->X[0] = 50.0;
+    src->blocks[1]->X[1] = 60.0;
+
+    transpose_spd_fill_values(src, (stacked_pd *) out_m);
+    stacked_pd *out = (stacked_pd *) out_m;
+
+    /* Expected after refill:
+         O0: X = [50]
+         O1: X = [100, 300, 60]
+         O2: X = [200, 400]                                             */
+    double O0_X[1] = {50.0};
+    mu_assert("O0 X refilled", cmp_double_array(out->blocks[0]->X, O0_X, 1));
+    double O1_X[3] = {100.0, 300.0, 60.0};
+    mu_assert("O1 X refilled", cmp_double_array(out->blocks[1]->X, O1_X, 3));
+    double O2_X[2] = {200.0, 400.0};
+    mu_assert("O2 X refilled", cmp_double_array(out->blocks[2]->X, O2_X, 2));
+
+    free_matrix(out_m);
+    free_matrix(src_m);
+    return 0;
+}
+
+/* Empty source: output also empty; work is an empty spd; free safe. */
+const char *test_transpose_spd_empty(void)
+{
+    matrix *src_m = new_stacked_pd(4, 5, 0, NULL, NULL, NULL);
+    matrix *out_m = transpose_spd_alloc((stacked_pd *) src_m);
+    transpose_spd_fill_values((stacked_pd *) src_m, (stacked_pd *) out_m);
+    stacked_pd *out = (stacked_pd *) out_m;
+
+    mu_assert("m", out_m->m == 5);
+    mu_assert("n", out_m->n == 4);
+    mu_assert("n_blocks", out->n_blocks == 0);
+    mu_assert("nnz", out_m->nnz == 0);
+    mu_assert("work not NULL", out->work != NULL);
+    mu_assert("work n_blocks", out->work->n_blocks == 0);
+
+    free_matrix(out_m);
+    free_matrix(src_m);
+    return 0;
+}
+
+/* Single full-dense block: trivial case, just a dense transpose. */
+const char *test_transpose_spd_single_block_full(void)
+{
+    /* 2x3 dense block, X = [[1, 2, 3], [4, 5, 6]] */
+    double data[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
+    matrix *blk = new_permuted_dense_full(2, 3, data);
+
+    permuted_dense *blocks[1] = {(permuted_dense *) blk};
+    matrix *src_m = new_stacked_pd(2, 3, 1, blocks, NULL, NULL);
+
+    matrix *out_m = transpose_spd_alloc((stacked_pd *) src_m);
+    transpose_spd_fill_values((stacked_pd *) src_m, (stacked_pd *) out_m);
+    stacked_pd *out = (stacked_pd *) out_m;
+
+    mu_assert("m", out_m->m == 3);
+    mu_assert("n", out_m->n == 2);
+    mu_assert("n_blocks", out->n_blocks == 1);
+
+    /* Transposed 3x2 dense: [[1, 4], [2, 5], [3, 6]] -> [1, 4, 2, 5, 3, 6]. */
+    permuted_dense *O = out->blocks[0];
+    int O_row[3] = {0, 1, 2};
+    int O_col[2] = {0, 1};
+    double O_X[6] = {1.0, 4.0, 2.0, 5.0, 3.0, 6.0};
+    mu_assert("O m0", O->m0 == 3);
+    mu_assert("O n0", O->n0 == 2);
+    mu_assert("O row", cmp_int_array(O->row_perm, O_row, 3));
+    mu_assert("O col", cmp_int_array(O->col_perm, O_col, 2));
+    mu_assert("O X", cmp_double_array(O->X, O_X, 6));
+
+    free_matrix(out_m);
+    free_matrix(src_m);
+    return 0;
+}
+
 #endif /* TEST_STACKED_PD_H */