Skip to content

Commit e4f4100

Browse files
dance858dance858
committed
BTA_spd_matrix dispatcher
1 parent db319a7 commit e4f4100

4 files changed

Lines changed: 221 additions & 0 deletions

File tree

‎include/utils/matrix_BTA.h‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,17 @@ matrix *BA_pd_matrices_alloc(const permuted_dense *B, matrix *A);
4747
void BA_pd_matrices_fill_values(const permuted_dense *B, const matrix *A,
4848
permuted_dense *C);
4949

50+
/* Polymorphic dispatcher for C = B^T @ (diag(d) @) A where B is stacked_pd
51+
and A is any matrix type (PD, stacked_pd, or sparse_matrix). C is always
52+
stacked_pd (one block per signature group of overlapping c_k's of B).
53+
For the sparse-A branch the dispatcher ensures sm_A->csc_cache structure
54+
exists at alloc time; the caller must refresh sm_A->csc_cache values via
55+
sm_A->refresh_csc_values before calling _fill_values (same fill-side
56+
contract as BA_pd_matrices_fill_values). */
57+
matrix *BTA_spd_matrices_alloc(const stacked_pd *B, matrix *A);
58+
void BTDA_spd_matrices_fill_values(const stacked_pd *B, const double *d,
59+
const matrix *A, stacked_pd *C);
60+
5061
/* Polymorphic dispatcher: C = kron(I_p, A) @ J as a stacked_pd, where A
5162
is a permuted_dense and J is any matrix type (permuted_dense,
5263
stacked_pd, or sparse_matrix). Output shape: (A->m * p) x J->n. For

‎src/utils/matrix_BTA.c‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,42 @@ void BA_pd_matrices_fill_values(const permuted_dense *B, const matrix *A,
166166
BA_pd_csc_fill_values(B->X, B->n0, B->col_inv, sm_A->csc_cache, C);
167167
}
168168

169+
matrix *BTA_spd_matrices_alloc(const stacked_pd *B, matrix *A)
170+
{
171+
if (A->is_permuted_dense)
172+
{
173+
return BTA_spd_pd_alloc(B, (const permuted_dense *) A);
174+
}
175+
if (A->is_stacked_pd)
176+
{
177+
return BTA_spd_spd_alloc(B, (const stacked_pd *) A);
178+
}
179+
180+
/* A is sparse */
181+
sparse_matrix *sm_A = (sparse_matrix *) A;
182+
sparse_matrix_ensure_csc_cache(sm_A);
183+
return BTA_spd_csc_alloc(B, sm_A->csc_cache);
184+
}
185+
186+
void BTDA_spd_matrices_fill_values(const stacked_pd *B, const double *d,
187+
const matrix *A, stacked_pd *C)
188+
{
189+
if (A->is_permuted_dense)
190+
{
191+
BTDA_spd_pd_fill_values(B, d, (const permuted_dense *) A, C);
192+
return;
193+
}
194+
if (A->is_stacked_pd)
195+
{
196+
BTDA_spd_spd_fill_values(B, d, (const stacked_pd *) A, C);
197+
return;
198+
}
199+
200+
/* A is sparse */
201+
const sparse_matrix *sm_A = (const sparse_matrix *) A;
202+
BTDA_spd_csc_fill_values(B, d, sm_A->csc_cache, C);
203+
}
204+
169205
matrix *BA_pd_kron_matrices_alloc(const permuted_dense *A, int p, matrix *J)
170206
{
171207
if (J->is_permuted_dense)

‎tests/all_tests.c‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -441,6 +441,9 @@ int main(void)
441441
mu_run_test(test_BTDA_spd_pd_overlapping_cp, tests_run);
442442
mu_run_test(test_BTDA_spd_csc_overlapping_cp, tests_run);
443443
mu_run_test(test_BTDA_spd_spd_overlapping, tests_run);
444+
mu_run_test(test_BTA_spd_matrices_pd_A, tests_run);
445+
mu_run_test(test_BTA_spd_matrices_csc_A, tests_run);
446+
mu_run_test(test_BTA_spd_matrices_spd_A, tests_run);
444447
mu_run_test(test_stacked_pd_construct_and_free, tests_run);
445448
mu_run_test(test_coalesce_no_overlap, tests_run);
446449
mu_run_test(test_coalesce_three_signatures, tests_run);

‎tests/utils/test_matrix_BTA.h‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -705,4 +705,175 @@ const char *test_BTDA_spd_spd_overlapping(void)
705705
return 0;
706706
}
707707

708+
/* BTA_spd_matrices dispatcher: A is permuted_dense. Verifies the
709+
PD-branch routes to BTA_spd_pd / BTDA_spd_pd. */
710+
const char *test_BTA_spd_matrices_pd_A(void)
711+
{
712+
/* Same inputs as test_BTDA_spd_pd_overlapping_cp. */
713+
int B0_rp[2] = {0, 1};
714+
int B0_cp[2] = {0, 2};
715+
double B0X[4] = {1, 2, 3, 4};
716+
matrix *B_blk0 = new_permuted_dense(4, 3, 2, 2, B0_rp, B0_cp, B0X);
717+
int B1_rp[2] = {2, 3};
718+
int B1_cp[2] = {1, 2};
719+
double B1X[4] = {5, 6, 7, 8};
720+
matrix *B_blk1 = new_permuted_dense(4, 3, 2, 2, B1_rp, B1_cp, B1X);
721+
permuted_dense *B_blocks[2] = {(permuted_dense *) B_blk0,
722+
(permuted_dense *) B_blk1};
723+
matrix *B_spd = new_stacked_pd(4, 3, 2, B_blocks, NULL, NULL);
724+
725+
int A_rp[4] = {0, 1, 2, 3};
726+
int A_cp[3] = {0, 1, 2};
727+
double AX[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
728+
matrix *A = new_permuted_dense(4, 3, 4, 3, A_rp, A_cp, AX);
729+
730+
double d[4] = {2.0, -1.5, 0.5, 1.25};
731+
732+
/* Route 1: new dispatcher. */
733+
matrix *C_ours = BTA_spd_matrices_alloc((stacked_pd *) B_spd, A);
734+
BTDA_spd_matrices_fill_values((stacked_pd *) B_spd, d, A, (stacked_pd *) C_ours);
735+
736+
/* Route 2: production dispatcher with B flattened to sparse. */
737+
matrix *B_sparse = spd_to_sparse_matrix_copy(B_spd);
738+
matrix *C_ref = BTA_matrices_alloc(A, B_sparse);
739+
B_sparse->refresh_csc_values(B_sparse);
740+
BTDA_matrices_fill_values(A, d, B_sparse, C_ref);
741+
742+
CSR_matrix *csr_ours = C_ours->to_csr(C_ours);
743+
CSR_matrix *csr_ref = C_ref->to_csr(C_ref);
744+
mu_assert("m", csr_ours->m == csr_ref->m);
745+
mu_assert("n", csr_ours->n == csr_ref->n);
746+
mu_assert("nnz", csr_ours->nnz == csr_ref->nnz);
747+
mu_assert("p", cmp_int_array(csr_ours->p, csr_ref->p, csr_ours->m + 1));
748+
mu_assert("i", cmp_int_array(csr_ours->i, csr_ref->i, csr_ours->nnz));
749+
mu_assert("x", cmp_double_array(csr_ours->x, csr_ref->x, csr_ours->nnz));
750+
751+
free_matrix(C_ref);
752+
free_matrix(B_sparse);
753+
free_matrix(C_ours);
754+
free_matrix(A);
755+
free_matrix(B_spd);
756+
return 0;
757+
}
758+
759+
/* BTA_spd_matrices dispatcher: A is sparse_matrix. Verifies the
760+
sparse-branch ensures csc_cache and routes to BTA_spd_csc / BTDA_spd_csc. */
761+
const char *test_BTA_spd_matrices_csc_A(void)
762+
{
763+
int B0_rp[2] = {0, 1};
764+
int B0_cp[2] = {0, 2};
765+
double B0X[4] = {1, 2, 3, 4};
766+
matrix *B_blk0 = new_permuted_dense(4, 3, 2, 2, B0_rp, B0_cp, B0X);
767+
int B1_rp[2] = {2, 3};
768+
int B1_cp[2] = {1, 2};
769+
double B1X[4] = {5, 6, 7, 8};
770+
matrix *B_blk1 = new_permuted_dense(4, 3, 2, 2, B1_rp, B1_cp, B1X);
771+
permuted_dense *B_blocks[2] = {(permuted_dense *) B_blk0,
772+
(permuted_dense *) B_blk1};
773+
matrix *B_spd = new_stacked_pd(4, 3, 2, B_blocks, NULL, NULL);
774+
775+
/* A as sparse_matrix wrapping a CSR (same shape/values as
776+
test_BTDA_spd_csc_overlapping_cp). */
777+
CSR_matrix *A_csr = new_CSR_matrix(4, 3, 7);
778+
int Ap[5] = {0, 2, 3, 6, 7};
779+
int Ai[7] = {0, 2, 1, 0, 1, 2, 2};
780+
double Ax[7] = {1, 2, 3, 4, 5, 6, 7};
781+
memcpy(A_csr->p, Ap, sizeof(Ap));
782+
memcpy(A_csr->i, Ai, sizeof(Ai));
783+
memcpy(A_csr->x, Ax, sizeof(Ax));
784+
matrix *A_sm = new_sparse_matrix(A_csr);
785+
786+
double d[4] = {2.0, -1.5, 0.5, 1.25};
787+
788+
/* Route 1: new dispatcher. Alloc ensures csc_cache structure; per
789+
contract, caller refreshes values before fill. */
790+
matrix *C_ours = BTA_spd_matrices_alloc((stacked_pd *) B_spd, A_sm);
791+
A_sm->refresh_csc_values(A_sm);
792+
BTDA_spd_matrices_fill_values((stacked_pd *) B_spd, d, A_sm,
793+
(stacked_pd *) C_ours);
794+
795+
/* Route 2: production dispatcher with B flattened to sparse. */
796+
matrix *B_sparse = spd_to_sparse_matrix_copy(B_spd);
797+
matrix *C_ref = BTA_matrices_alloc(A_sm, B_sparse);
798+
B_sparse->refresh_csc_values(B_sparse);
799+
BTDA_matrices_fill_values(A_sm, d, B_sparse, C_ref);
800+
801+
CSR_matrix *csr_ours = C_ours->to_csr(C_ours);
802+
CSR_matrix *csr_ref = C_ref->to_csr(C_ref);
803+
mu_assert("m", csr_ours->m == csr_ref->m);
804+
mu_assert("n", csr_ours->n == csr_ref->n);
805+
mu_assert("nnz", csr_ours->nnz == csr_ref->nnz);
806+
mu_assert("p", cmp_int_array(csr_ours->p, csr_ref->p, csr_ours->m + 1));
807+
mu_assert("i", cmp_int_array(csr_ours->i, csr_ref->i, csr_ours->nnz));
808+
mu_assert("x", cmp_double_array(csr_ours->x, csr_ref->x, csr_ours->nnz));
809+
810+
free_matrix(C_ref);
811+
free_matrix(B_sparse);
812+
free_matrix(C_ours);
813+
free_matrix(A_sm);
814+
free_matrix(B_spd);
815+
return 0;
816+
}
817+
818+
/* BTA_spd_matrices dispatcher: A is stacked_pd. Verifies the
819+
spd-branch routes to BTA_spd_spd / BTDA_spd_spd. */
820+
const char *test_BTA_spd_matrices_spd_A(void)
821+
{
822+
int B0_rp[2] = {0, 1};
823+
int B0_cp[2] = {0, 2};
824+
double B0X[4] = {1, 2, 3, 4};
825+
matrix *B_blk0 = new_permuted_dense(4, 3, 2, 2, B0_rp, B0_cp, B0X);
826+
int B1_rp[2] = {2, 3};
827+
int B1_cp[2] = {1, 2};
828+
double B1X[4] = {5, 6, 7, 8};
829+
matrix *B_blk1 = new_permuted_dense(4, 3, 2, 2, B1_rp, B1_cp, B1X);
830+
permuted_dense *B_blocks[2] = {(permuted_dense *) B_blk0,
831+
(permuted_dense *) B_blk1};
832+
matrix *B_spd = new_stacked_pd(4, 3, 2, B_blocks, NULL, NULL);
833+
834+
int A0_rp[2] = {0, 1};
835+
int A0_cp[2] = {0, 1};
836+
double A0X[4] = {10, 11, 12, 13};
837+
matrix *A_blk0 = new_permuted_dense(4, 3, 2, 2, A0_rp, A0_cp, A0X);
838+
int A1_rp[2] = {2, 3};
839+
int A1_cp[2] = {1, 2};
840+
double A1X[4] = {20, 21, 22, 23};
841+
matrix *A_blk1 = new_permuted_dense(4, 3, 2, 2, A1_rp, A1_cp, A1X);
842+
permuted_dense *A_blocks[2] = {(permuted_dense *) A_blk0,
843+
(permuted_dense *) A_blk1};
844+
matrix *A_spd = new_stacked_pd(4, 3, 2, A_blocks, NULL, NULL);
845+
846+
double d[4] = {2.0, -1.5, 0.5, 1.25};
847+
848+
/* Route 1: new dispatcher. */
849+
matrix *C_ours = BTA_spd_matrices_alloc((stacked_pd *) B_spd, A_spd);
850+
BTDA_spd_matrices_fill_values((stacked_pd *) B_spd, d, A_spd,
851+
(stacked_pd *) C_ours);
852+
853+
/* Route 2: production dispatcher with both flattened to sparse. */
854+
matrix *A_sparse = spd_to_sparse_matrix_copy(A_spd);
855+
matrix *B_sparse = spd_to_sparse_matrix_copy(B_spd);
856+
matrix *C_ref = BTA_matrices_alloc(A_sparse, B_sparse);
857+
A_sparse->refresh_csc_values(A_sparse);
858+
B_sparse->refresh_csc_values(B_sparse);
859+
BTDA_matrices_fill_values(A_sparse, d, B_sparse, C_ref);
860+
861+
CSR_matrix *csr_ours = C_ours->to_csr(C_ours);
862+
CSR_matrix *csr_ref = C_ref->to_csr(C_ref);
863+
mu_assert("m", csr_ours->m == csr_ref->m);
864+
mu_assert("n", csr_ours->n == csr_ref->n);
865+
mu_assert("nnz", csr_ours->nnz == csr_ref->nnz);
866+
mu_assert("p", cmp_int_array(csr_ours->p, csr_ref->p, csr_ours->m + 1));
867+
mu_assert("i", cmp_int_array(csr_ours->i, csr_ref->i, csr_ours->nnz));
868+
mu_assert("x", cmp_double_array(csr_ours->x, csr_ref->x, csr_ours->nnz));
869+
870+
free_matrix(C_ref);
871+
free_matrix(B_sparse);
872+
free_matrix(A_sparse);
873+
free_matrix(C_ours);
874+
free_matrix(A_spd);
875+
free_matrix(B_spd);
876+
return 0;
877+
}
878+
708879
#endif /* TEST_MATRIX_BTA_H */

0 commit comments

Comments
 (0)