bola 0.1.2

A lightweight native rust linear algebra library.
Documentation 
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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

use rayon::prelude::*;
use std::{ptr, slice};

#[cfg(feature = "profiling")]
use crate::profiling;

const PARALLEL_THRESHOLD: usize = 256;

/// Performs a series of row interchanges on a matrix.
///
/// # Arguments
/// * `n`    - The number of columns of the matrix `A` to which the interchanges will be applied.
/// * `a`    - A mutable slice representing the matrix `A` in column-major order.
/// * `lda`  - The leading dimension of `A`, i.e., the stride between consecutive columns.
/// * `k1`   - The first element of `ipiv` to be used for interchanges (0-based).
/// * `k2`   - The last element of `ipiv` to be used for interchanges (0-based).
/// * `ipiv` - The slice containing the pivot indices. For `i` in `k1..=k2`, row `i` is
///   swapped with row `ipiv[i]`.
/// * `incx` - The increment between successive values of `ipiv`. If `incx` is positive,
///   the loop is forward; if negative, it's backward. If zero, no swaps are performed.
#[allow(unsafe_op_in_unsafe_fn, clippy::missing_safety_doc)]
pub unsafe fn laswp(n: usize, a: &mut [f64], lda: usize, k1: i32, k2: i32, ipiv: &[i32], incx: i32) {
    #[cfg(feature = "profiling")]
    let _timer = profiling::ScopedTimer::new("LASWP");

    if n == 0 || incx == 0 || k1 > k2 {
        return;
    }

    let (ix0, i1, i2, i_inc) = if incx > 0 {
        ((k1 - 1) as usize, k1, k2, 1)
    } else {
        let ix_start_0based = (k1 - 1) as usize;
        let num_pivots = (k2 - k1) as usize;
        let ix0_final = ix_start_0based + num_pivots * incx.unsigned_abs() as usize;
        (ix0_final, k2, k1, -1)
    };

    if n < PARALLEL_THRESHOLD {
        let a_ptr = a.as_mut_ptr();
        let n_unrolled = n - (n % 4);

        for k_base in (0..n_unrolled).step_by(4) {
            let col0_ptr = a_ptr.add(k_base * lda);
            let col1_ptr = a_ptr.add((k_base + 1) * lda);
            let col2_ptr = a_ptr.add((k_base + 2) * lda);
            let col3_ptr = a_ptr.add((k_base + 3) * lda);

            let mut ix = ix0;
            let mut i = i1;
            while if i_inc > 0 { i <= i2 } else { i >= i2 } {
                let ip = ipiv[ix] as usize;
                if ip != i as usize {
                    let row_i_m1 = i as usize - 1;
                    let row_ip_m1 = ip - 1;
                    ptr::swap(col0_ptr.add(row_i_m1), col0_ptr.add(row_ip_m1));
                    ptr::swap(col1_ptr.add(row_i_m1), col1_ptr.add(row_ip_m1));
                    ptr::swap(col2_ptr.add(row_i_m1), col2_ptr.add(row_ip_m1));
                    ptr::swap(col3_ptr.add(row_i_m1), col3_ptr.add(row_ip_m1));
                }
                ix = (ix as i32 + incx) as usize;
                i += i_inc;
            }
        }

        for k in n_unrolled..n {
            let col_ptr = a_ptr.add(k * lda);
            let mut ix = ix0;
            let mut i = i1;
            while if i_inc > 0 { i <= i2 } else { i >= i2 } {
                let ip = ipiv[ix] as usize;
                if ip != i as usize {
                    ptr::swap(col_ptr.add(i as usize - 1), col_ptr.add(ip - 1));
                }
                ix = (ix as i32 + incx) as usize;
                i += i_inc;
            }
        }
        return;
    }

    let a_addr = a.as_mut_ptr() as usize;
    let ipiv_addr = ipiv.as_ptr() as usize;
    let n32 = (n / 32) * 32;

    if n32 != 0 {
        (0..n32).into_par_iter().step_by(32).for_each(|j_chunk_start| {
            let a_ptr = a_addr as *mut f64;
            let ipiv = slice::from_raw_parts(ipiv_addr as *const i32, ipiv.len());
            for k_base in (j_chunk_start..j_chunk_start + 32).step_by(4) {
                let col0_ptr = a_ptr.add(k_base * lda);
                let col1_ptr = a_ptr.add((k_base + 1) * lda);
                let col2_ptr = a_ptr.add((k_base + 2) * lda);
                let col3_ptr = a_ptr.add((k_base + 3) * lda);
                let mut ix = ix0;
                let mut i = i1;
                while if i_inc > 0 { i <= i2 } else { i >= i2 } {
                    let ip = ipiv[ix] as usize;
                    if ip != i as usize {
                        let row_i_m1 = i as usize - 1;
                        let row_ip_m1 = ip - 1;
                        ptr::swap(col0_ptr.add(row_i_m1), col0_ptr.add(row_ip_m1));
                        ptr::swap(col1_ptr.add(row_i_m1), col1_ptr.add(row_ip_m1));
                        ptr::swap(col2_ptr.add(row_i_m1), col2_ptr.add(row_ip_m1));
                        ptr::swap(col3_ptr.add(row_i_m1), col3_ptr.add(row_ip_m1));
                    }
                    ix = (ix as i32 + incx) as usize;
                    i += i_inc;
                }
            }
        });
    }

    if n32 != n {
        let a_ptr = a.as_mut_ptr();
        let mut ix = ix0;
        let mut i = i1;
        while if i_inc > 0 { i <= i2 } else { i >= i2 } {
            let ip = ipiv[ix] as usize;
            if ip != i as usize {
                let row1_base_ptr = a_ptr.add(i as usize - 1);
                let row2_base_ptr = a_ptr.add(ip - 1);
                for k in n32..n {
                    let offset = k * lda;
                    ptr::swap(row1_base_ptr.add(offset), row2_base_ptr.add(offset));
                }
            }
            ix = (ix as i32 + incx) as usize;
            i += i_inc;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_laswp_row_swaps() {
        let n_cols_to_swap = 3;
        let lda = 3;
        let mut a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];

        let ipiv = vec![3, 2];

        let a_expected = vec![3.0, 2.0, 1.0, 6.0, 5.0, 4.0, 9.0, 8.0, 7.0];

        unsafe {
            laswp(n_cols_to_swap, &mut a, lda, 1, 2, &ipiv, 1);
        }

        assert_eq!(a.len(), a_expected.len(), "Slices have different lengths");
        for (i, (val_a, val_b)) in a.iter().zip(a_expected.iter()).enumerate() {
            assert!(
                (val_a - val_b).abs() < 1e-8,
                "Mismatch at index {}: evaluated[{}] = {}, expected[{}] = {}",
                i,
                i,
                val_a,
                i,
                val_b
            );
        }
    }
}