|
18 | 18 | #include "old-code/old_CSR_sum.h" |
19 | 19 | #include "utils/CSR_matrix.h" |
20 | 20 | #include "utils/int_double_pair.h" |
| 21 | +#include "utils/utils.h" |
21 | 22 | #include <assert.h> |
22 | 23 | #include <stdlib.h> |
23 | 24 | #include <string.h> |
@@ -330,3 +331,198 @@ void sum_spaced_rows_into_row_csr(const CSR_matrix *A, CSR_matrix *C, |
330 | 331 |
|
331 | 332 | C->p[1] = C->nnz; |
332 | 333 | } |
| 334 | + |
| 335 | +// ------------------------------------------------------------------------------------ |
| 336 | +// Row-sum kernels with an idx_map (input nnz -> output position), retired from the |
| 337 | +// engine when the sum atom moved onto the generic row_reduce primitive. Kept as |
| 338 | +// standalone CSR operations. Fill values with accumulator() from utils/CSR_sum.h |
| 339 | +// after zeroing C->x. |
| 340 | +// ------------------------------------------------------------------------------------ |
| 341 | + |
| 342 | +/* iwork must have size max(A->n, A->nnz), and idx_map must have size A->nnz */ |
| 343 | +void sum_all_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C, int *iwork, |
| 344 | + int *idx_map) |
| 345 | +{ |
| 346 | + // ------------------------------------------------------------------- |
| 347 | + // Build sparsity pattern of the summed row |
| 348 | + // ------------------------------------------------------------------- |
| 349 | + int *cols = iwork; |
| 350 | + memcpy(cols, A->i, A->nnz * sizeof(int)); |
| 351 | + sort_int_array(cols, A->nnz); |
| 352 | + |
| 353 | + int unique_nnz = 0; |
| 354 | + int prev_col = -1; |
| 355 | + for (int j = 0; j < A->nnz; j++) |
| 356 | + { |
| 357 | + if (cols[j] != prev_col) |
| 358 | + { |
| 359 | + C->i[unique_nnz] = cols[j]; |
| 360 | + prev_col = cols[j]; |
| 361 | + unique_nnz++; |
| 362 | + } |
| 363 | + } |
| 364 | + |
| 365 | + C->p[0] = 0; |
| 366 | + C->p[1] = unique_nnz; |
| 367 | + C->nnz = unique_nnz; |
| 368 | + |
| 369 | + // ------------------------------------------------------------------- |
| 370 | + // Map child values to summed-row positions. col_to_pos maps |
| 371 | + // column indices to positions in C's row. |
| 372 | + // ------------------------------------------------------------------- |
| 373 | + int *col_to_pos = iwork; |
| 374 | + for (int idx = 0; idx < unique_nnz; idx++) |
| 375 | + { |
| 376 | + col_to_pos[C->i[idx]] = idx; |
| 377 | + } |
| 378 | + |
| 379 | + for (int i = 0; i < A->m; i++) |
| 380 | + { |
| 381 | + for (int j = A->p[i]; j < A->p[i + 1]; j++) |
| 382 | + { |
| 383 | + idx_map[j] = col_to_pos[A->i[j]]; |
| 384 | + } |
| 385 | + } |
| 386 | +} |
| 387 | + |
| 388 | +/* iwork must have size max(A->n, A->nnz), and idx_map must have size A->nnz */ |
| 389 | +void sum_block_of_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C, |
| 390 | + int row_block_size, int *iwork, int *idx_map) |
| 391 | +{ |
| 392 | + assert(A->m % row_block_size == 0); |
| 393 | + int n_blocks = A->m / row_block_size; |
| 394 | + assert(C->m == n_blocks); |
| 395 | + |
| 396 | + C->n = A->n; |
| 397 | + C->p[0] = 0; |
| 398 | + int cursor = 0; |
| 399 | + |
| 400 | + int *cols = iwork; |
| 401 | + int *col_to_pos = iwork; |
| 402 | + |
| 403 | + for (int block = 0; block < n_blocks; block++) |
| 404 | + { |
| 405 | + int start_row = block * row_block_size; |
| 406 | + int end_row = start_row + row_block_size; |
| 407 | + |
| 408 | + // ----------------------------------------------------------------- |
| 409 | + // Build sparsity pattern of the row resulting from summing |
| 410 | + // the block of rows from A |
| 411 | + // ----------------------------------------------------------------- |
| 412 | + C->p[block] = cursor; |
| 413 | + int count = 0; |
| 414 | + for (int row = start_row; row < end_row; row++) |
| 415 | + { |
| 416 | + for (int j = A->p[row]; j < A->p[row + 1]; j++) |
| 417 | + { |
| 418 | + cols[count++] = A->i[j]; |
| 419 | + } |
| 420 | + } |
| 421 | + |
| 422 | + /* Sort columns and write unique pattern into C->i */ |
| 423 | + sort_int_array(cols, count); |
| 424 | + |
| 425 | + int unique_nnz = 0; |
| 426 | + int prev_col = -1; |
| 427 | + for (int t = 0; t < count; t++) |
| 428 | + { |
| 429 | + int col = cols[t]; |
| 430 | + if (t == 0 || col != prev_col) |
| 431 | + { |
| 432 | + C->i[cursor + unique_nnz] = col; |
| 433 | + prev_col = col; |
| 434 | + unique_nnz++; |
| 435 | + } |
| 436 | + } |
| 437 | + |
| 438 | + cursor += unique_nnz; |
| 439 | + C->p[block + 1] = cursor; |
| 440 | + |
| 441 | + // ----------------------------------------------------------------- |
| 442 | + // Build idx_map for all entries in this block |
| 443 | + // ----------------------------------------------------------------- |
| 444 | + int row_start = C->p[block]; |
| 445 | + for (int idx = 0; idx < unique_nnz; idx++) |
| 446 | + { |
| 447 | + col_to_pos[C->i[row_start + idx]] = row_start + idx; |
| 448 | + } |
| 449 | + |
| 450 | + for (int row = start_row; row < end_row; row++) |
| 451 | + { |
| 452 | + for (int j = A->p[row]; j < A->p[row + 1]; j++) |
| 453 | + { |
| 454 | + idx_map[j] = col_to_pos[A->i[j]]; |
| 455 | + } |
| 456 | + } |
| 457 | + } |
| 458 | + |
| 459 | + C->nnz = cursor; |
| 460 | +} |
| 461 | + |
| 462 | +/* iwork must have size max(A->n, A->nnz), and idx_map must have size A->nnz */ |
| 463 | +void sum_evenly_spaced_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C, |
| 464 | + int row_spacing, int *iwork, int *idx_map) |
| 465 | +{ |
| 466 | + assert(C->m == row_spacing); |
| 467 | + C->n = A->n; |
| 468 | + C->p[0] = 0; |
| 469 | + int cursor = 0; |
| 470 | + |
| 471 | + int *cols = iwork; |
| 472 | + int *col_to_pos = iwork; |
| 473 | + |
| 474 | + for (int C_row = 0; C_row < C->m; C_row++) |
| 475 | + { |
| 476 | + // ----------------------------------------------------------------- |
| 477 | + // Build sparsity pattern of the row resulting from summing |
| 478 | + // evenly spaced rows from A |
| 479 | + // ----------------------------------------------------------------- |
| 480 | + C->p[C_row] = cursor; |
| 481 | + int count = 0; |
| 482 | + for (int row = C_row; row < A->m; row += row_spacing) |
| 483 | + { |
| 484 | + for (int j = A->p[row]; j < A->p[row + 1]; j++) |
| 485 | + { |
| 486 | + cols[count++] = A->i[j]; |
| 487 | + } |
| 488 | + } |
| 489 | + |
| 490 | + /* Sort columns and write unique pattern into C->i */ |
| 491 | + sort_int_array(cols, count); |
| 492 | + |
| 493 | + int unique_nnz = 0; |
| 494 | + int prev_col = -1; |
| 495 | + for (int t = 0; t < count; t++) |
| 496 | + { |
| 497 | + int col = cols[t]; |
| 498 | + if (t == 0 || col != prev_col) |
| 499 | + { |
| 500 | + C->i[cursor + unique_nnz] = col; |
| 501 | + prev_col = col; |
| 502 | + unique_nnz++; |
| 503 | + } |
| 504 | + } |
| 505 | + |
| 506 | + cursor += unique_nnz; |
| 507 | + C->p[C_row + 1] = cursor; |
| 508 | + |
| 509 | + // ----------------------------------------------------------------- |
| 510 | + // Build idx_map for all entries in evenly spaced rows |
| 511 | + // ----------------------------------------------------------------- |
| 512 | + int row_start = C->p[C_row]; |
| 513 | + for (int idx = 0; idx < unique_nnz; idx++) |
| 514 | + { |
| 515 | + col_to_pos[C->i[row_start + idx]] = row_start + idx; |
| 516 | + } |
| 517 | + |
| 518 | + for (int row = C_row; row < A->m; row += row_spacing) |
| 519 | + { |
| 520 | + for (int j = A->p[row]; j < A->p[row + 1]; j++) |
| 521 | + { |
| 522 | + idx_map[j] = col_to_pos[A->i[j]]; |
| 523 | + } |
| 524 | + } |
| 525 | + } |
| 526 | + |
| 527 | + C->nnz = cursor; |
| 528 | +} |
0 commit comments