kryst 3.2.1

Krylov subspace and preconditioned iterative solvers for dense and sparse linear systems, with shared and distributed memory parallelism.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

//! Scalar sparse matrix-vector multiplication kernels.

#[allow(unused_imports)]
use crate::algebra::prelude::*;
use crate::matrix::csr::CsrMatrix;

/// Portable CSR SpMV: computes `y = A * x` without conjugation.
///
/// This helper keeps the control flow simple to act as the scalar fallback for
/// both real and complex builds.  The loops intentionally mirror the
/// conventional row-wise CSR traversal so it is easy to audit and compare
/// against the existing `spmv_scaled_csr` entry point which supports
/// alpha/beta scaling factors.
#[inline]
pub fn spmv_csr_scalar<S: KrystScalar>(a: &CsrMatrix<S>, x: &[S], y: &mut [S]) {
    debug_assert_eq!(x.len(), a.ncols);
    debug_assert_eq!(y.len(), a.nrows);

    for i in 0..a.nrows {
        let start = a.rowptr[i];
        let end = a.rowptr[i + 1];
        let mut acc = S::zero();
        for k in start..end {
            let j = a.colind[k];
            acc = acc + a.values[k] * x[j];
        }
        y[i] = acc;
    }
}

/// Computes `y = alpha * A * x + beta * y` for a matrix stored in CSR format.
///
/// This is the baseline SpMV kernel and remains the reference implementation
/// for the higher-level SIMD paths that will be added in later iterations. It
/// keeps the loop order deterministic and performs a small unrolled
/// accumulation to improve ILP on modern CPUs without depending on explicit
/// SIMD instructions.
pub fn spmv_scaled_csr<S: KrystScalar>(
    m: usize,
    row_ptr: &[usize],
    col_idx: &[usize],
    vals: &[S],
    alpha: S,
    x: &[S],
    beta: S,
    y: &mut [S],
) {
    assert_eq!(row_ptr.len(), m + 1);
    assert_eq!(col_idx.len(), vals.len());
    if let Some(max_col) = col_idx.iter().copied().max() {
        assert!(x.len() > max_col);
    }
    assert!(y.len() >= m);

    if beta == S::zero() {
        y[..m].fill(S::zero());
    } else if beta != S::one() {
        for yi in &mut y[..m] {
            *yi = (*yi) * beta;
        }
    }

    for i in 0..m {
        let mut acc = S::zero();
        let mut p = row_ptr[i];
        let end = row_ptr[i + 1];

        while p + 3 < end {
            acc = vals[p].mul_add(x[col_idx[p]], acc);
            acc = vals[p + 1].mul_add(x[col_idx[p + 1]], acc);
            acc = vals[p + 2].mul_add(x[col_idx[p + 2]], acc);
            acc = vals[p + 3].mul_add(x[col_idx[p + 3]], acc);
            p += 4;
        }

        while p < end {
            acc = vals[p].mul_add(x[col_idx[p]], acc);
            p += 1;
        }

        y[i] = y[i] + alpha * acc;
    }
}

/// Computes `y = alpha * A^T * x + beta * y` for a CSR matrix.
///
/// In complex builds, this applies the conjugate-transpose (`A^H`) by
/// conjugating the matrix entries.
pub fn spmv_t_scaled_csr<S: KrystScalar>(
    m: usize,
    row_ptr: &[usize],
    col_idx: &[usize],
    vals: &[S],
    alpha: S,
    x: &[S],
    beta: S,
    y: &mut [S],
) {
    assert_eq!(row_ptr.len(), m + 1);
    assert!(col_idx.len() == vals.len());

    if beta == S::zero() {
        y.fill(S::zero());
    } else if beta != S::one() {
        for yi in y.iter_mut() {
            *yi = (*yi) * beta;
        }
    }

    for i in 0..m {
        let xi = x[i];
        if xi == S::zero() {
            continue;
        }
        let rs = row_ptr[i];
        let re = row_ptr[i + 1];
        for p in rs..re {
            let idx = col_idx[p];
            y[idx] = y[idx] + alpha * vals[p].conj() * xi;
        }
    }
}