1use super::*;
2use crate::{Operation, symbolic::{symbolic_product, symbolic_sum}};
3
4pub fn csr_sum_pattern(a: CsrMatrix<'_>, b: CsrMatrix<'_>) -> Result<(CsrMatrixOwned, Vec<u32>, Vec<u32>), SparseError> {
5 if a.rows() != b.rows() || a.columns() != b.columns() {
6 return Err(SparseError::DimensionMismatch("CSR sum operands must have the same shape"));
7 }
8 let symbolic = symbolic_sum(a.into(), b.into())?;
9 let left = crate::symbolic::sum_entry_mapping(a.into(), &symbolic, a.index_base())?;
10 let right = crate::symbolic::sum_entry_mapping(b.into(), &symbolic, a.index_base())?;
11 let values = vec![0.0f32; symbolic.column_indices.len()];
12 let matrix = CsrMatrixOwned::new(a.rows(), a.columns(), symbolic.row_offsets, symbolic.column_indices, values, a.index_base())?;
13 Ok((matrix, left, right))
14}
15
16pub fn csrgeam(
17 operation_a: Operation, operation_b: Operation,
18 alpha: f32, a: CsrMatrix<'_>, beta: f32, b: CsrMatrix<'_>,
19) -> Result<CsrMatrixOwned, SparseError> {
20 let transposed_a = transpose(a, operation_a)?;
21 let transposed_b = transpose(b, operation_b)?;
22 let a = transposed_a.as_ref().map_or(a, CsrMatrixOwned::as_ref);
23 let b = transposed_b.as_ref().map_or(b, CsrMatrixOwned::as_ref);
24 if a.rows() != b.rows() || a.columns() != b.columns() {
25 return Err(SparseError::DimensionMismatch("op(A) and op(B) must have the same shape"));
26 }
27 let symbolic = symbolic_sum(a.into(), b.into())?;
28 let mut values = vec![0.0_f32; symbolic.column_indices.len()];
29 let output_base = a.index_base().value();
30 for row in 0..a.rows() {
31 let start = (symbolic.row_offsets[row] - output_base) as usize;
32 let end = (symbolic.row_offsets[row + 1] - output_base) as usize;
33 let columns = &symbolic.column_indices[start..end];
34 for (matrix, scale) in [(a, alpha), (b, beta)] {
35 let base = matrix.index_base().value();
36 for entry in (matrix.row_offsets()[row] - base) as usize..(matrix.row_offsets()[row + 1] - base) as usize {
37 let column = matrix.column_indices()[entry] - base + output_base;
38 let destination = start + columns.binary_search(&column)
39 .map_err(|_| SparseError::InvalidColumnIndex)?;
40 values[destination] = scale.mul_add(matrix.values()[entry], values[destination]);
41 }
42 }
43 }
44 CsrMatrixOwned::new(a.rows(), a.columns(), symbolic.row_offsets, symbolic.column_indices, values, a.index_base())
45}
46
47pub fn csrgemm(
48 operation_a: Operation, operation_b: Operation,
49 alpha: f32, a: CsrMatrix<'_>, b: CsrMatrix<'_>,
50) -> Result<CsrMatrixOwned, SparseError> {
51 let transposed_a = transpose(a, operation_a)?;
52 let transposed_b = transpose(b, operation_b)?;
53 let a = transposed_a.as_ref().map_or(a, CsrMatrixOwned::as_ref);
54 let b = transposed_b.as_ref().map_or(b, CsrMatrixOwned::as_ref);
55 if a.columns() != b.rows() {
56 return Err(SparseError::DimensionMismatch("op(A) columns must equal op(B) rows"));
57 }
58 let symbolic = symbolic_product(a.into(), b.into())?;
59 let mut values = vec![0.0_f32; symbolic.column_indices.len()];
60 let a_base = a.index_base().value();
61 let b_base = b.index_base().value();
62 for row in 0..a.rows() {
63 let start = (symbolic.row_offsets[row] - a_base) as usize;
64 let end = (symbolic.row_offsets[row + 1] - a_base) as usize;
65 let columns = &symbolic.column_indices[start..end];
66 for entry in (a.row_offsets()[row] - a_base) as usize..(a.row_offsets()[row + 1] - a_base) as usize {
67 let inner = (a.column_indices()[entry] - a_base) as usize;
68 for other in (b.row_offsets()[inner] - b_base) as usize..(b.row_offsets()[inner + 1] - b_base) as usize {
69 let column = b.column_indices()[other] - b_base + a_base;
70 let destination = start + columns.binary_search(&column)
71 .map_err(|_| SparseError::InvalidColumnIndex)?;
72 values[destination] = a.values()[entry].mul_add(b.values()[other], values[destination]);
73 }
74 }
75 }
76 for value in &mut values { *value *= alpha; }
77 CsrMatrixOwned::new(a.rows(), b.columns(), symbolic.row_offsets, symbolic.column_indices, values, a.index_base())
78}
79
80pub(super) fn transpose(matrix: CsrMatrix<'_>, operation: Operation) -> Result<Option<CsrMatrixOwned>, SparseError> {
81 match operation {
82 Operation::None => Ok(None),
83 Operation::Transpose | Operation::ConjugateTranspose => matrix.transpose().map(Some),
84 }
85}
86
87pub fn csr_product_pattern(a: CsrMatrix<'_>, b: CsrMatrix<'_>) -> Result<CsrMatrixOwned, SparseError> {
88 if a.columns() != b.rows() {
89 return Err(SparseError::DimensionMismatch("CSR product inner dimensions differ"));
90 }
91 let symbolic = symbolic_product(a.into(), b.into())?;
92 let values = vec![0.0f32; symbolic.column_indices.len()];
93 CsrMatrixOwned::new(a.rows(), b.columns(), symbolic.row_offsets, symbolic.column_indices, values, a.index_base())
94}