Skip to main content

a/math/matrix/
lu.rs

1//! The LU decomposition with partial pivoting, and what it gives: determinants, solutions of
2//! linear systems and inverses.
3//!
4//! # The decomposition
5//!
6//! Square matrices up to 16 × 16 are factored by [`tiny`]: a copy on the stack, its order known
7//! so that its loops unroll, the pivots chosen and the rows swapped by selects rather than
8//! branches, which keeps the whole matrix in registers. Up to 24 columns, [`unblocked`]
9//! eliminates a column at a time by rank-one updates of whole columns. Larger matrices go
10//! Toledo's way ([`factor`]): the left half, then the right half updated by it, down to panels
11//! of [`NARROW`] = 16 columns. The row swaps of each half are applied to the other a column at a
12//! time, each column read once, rather than a row at a time across all of them, which took a
13//! third of the time at 512 × 512; and `U12 = L11⁻¹·A12` comes from the triangular solve with a
14//! dense copy of it that the product updating `A22` reads, in working memory taken once.
15//!
16//! Every path chooses the same pivots, the first element of the largest `|re| + |im|` in its
17//! column, as LAPACK does, and divides by a pivot so small that its reciprocal overflows rather
18//! than multiplying by the reciprocal, as LAPACK's `getf2` does.
19//!
20//! # The inverse
21//!
22//! Up to 12 × 12, [`tiny_inverse`] inverts on the stack, in unrolled loops, and allocates only
23//! the inverse. Larger matrices are factored in working memory, then inverted as LAPACK's
24//! `getri` does: `U⁻¹` by [`upper_inverse`], then `X·L = U⁻¹` by one triangular solve on the
25//! right, then the columns swapped back. That is `2n³` operations in all, against `8n³/3` for
26//! the solve against the identity that it replaced, nearly all of them in products of blocks.
27//!
28//! # Speed
29//!
30//! Measured on an Apple M4 Max, one core, `f64`, in µs, the median of 15 to 21 interleaved
31//! batches, against faer 0.24 (`partial_piv_lu`, then its `inverse`, `Par::Seq`) and nalgebra
32//! 0.35 (`lu`, `try_inverse`):
33//!
34//! | n | LU | faer | nalgebra | inverse | faer | nalgebra |
35//! | --- | --- | --- | --- | --- | --- | --- |
36//! | 2 | 0.034 | 0.122 | 0.036 | 0.014 | 0.233 | 0.015 |
37//! | 4 | 0.061 | 0.223 | 0.071 | 0.054 | 0.490 | 0.042 |
38//! | 8 | 0.140 | 0.381 | 0.173 | 0.209 | 0.970 | 0.582 |
39//! | 16 | 0.525 | 0.933 | 0.541 | 1.21 | 2.57 | 2.43 |
40//! | 32 | 2.03 | 3.19 | 2.59 | 5.10 | 9.24 | 13.8 |
41//! | 64 | 9.45 | 14.2 | 14.7 | 24.0 | 38.1 | 70.8 |
42//! | 128 | 52.5 | 71.2 | 99.8 | 146 | 183 | 439 |
43//! | 256 | 302 | 400 | 1 109 | 903 | 1 061 | 4 071 |
44//! | 512 | 2 024 | 2 697 | 8 417 | 5 769 | 7 675 | 31 030 |
45//! | 1024 | 14 424 | 18 348 | 63 443 | 43 363 | 53 809 | 231 502 |
46//!
47//! nalgebra inverts matrices up to 4 × 4 by their cofactors, a formula with half the operations
48//! and none of the pivoting, at 3 × 3 and 4 × 4 a little faster than the decomposition. Apple's
49//! Accelerate, on the matrix unit (SME) rather than NEON, ran 2 to 5 times faster from 128 × 128
50//! up (`dgetrf` 836 µs and `dgetrf` with `dgetri` 1 785 µs at 512).
51
52use super::blas::{Op, Side, Uplo, gemm, trsm, trsm_into, with_scratch};
53use super::{Matrix, Scalar};
54use crate::math::complex::Float;
55use core::hint::select_unpredictable;
56
57/// The LU decomposition `P·A = L·U` of an `m × n` matrix, by Gaussian elimination with partial
58/// pivoting: `P` a permutation, `L` of `m × min(m, n)` lower triangular with ones on its
59/// diagonal, `U` of `min(m, n) × n` upper triangular. From [`Matrix::lu`].
60///
61/// # Examples
62///
63/// ```
64/// use a::math::matrix::Matrix;
65///
66/// let a = Matrix::from_rows(&[[2.0, 1.0], [4.0, 3.0]]);
67/// let lu = a.lu();
68/// // The larger first column's element is the pivot: the rows are swapped.
69/// assert_eq!(lu.permutation(), [1, 0]);
70/// assert_eq!(lu.l(), Matrix::from_rows(&[[1.0, 0.0], [0.5, 1.0]]));
71/// assert_eq!(lu.u(), Matrix::from_rows(&[[4.0, 3.0], [0.0, -0.5]]));
72/// assert_eq!(&lu.p() * &a, &lu.l() * &lu.u());
73/// assert_eq!(lu.det(), 2.0);
74/// let x = lu.solve(&Matrix::from_rows(&[[3.0], [7.0]])).unwrap();
75/// assert_eq!(x, Matrix::from_rows(&[[1.0], [1.0]]));
76/// ```
77#[derive(Clone, Debug)]
78pub struct Lu<T> {
79    /// `L` below the diagonal and `U` on and above it.
80    lu: Matrix<T>,
81    /// Row `i` of `P·A` is row `perm[i]` of `A`.
82    perm: Vec<usize>,
83    /// Whether `P` is an odd permutation.
84    odd: bool,
85}
86
87impl<T: Scalar> Matrix<T> {
88    /// The LU decomposition with partial pivoting, `P·A = L·U`: see [`Lu`]. It takes about
89    /// `2n³/3` operations for an `n × n` matrix, nearly all of them in products of blocks.
90    pub fn lu(&self) -> Lu<T> {
91        let (m, n) = self.shape();
92        let mut lu = self.clone();
93        let mut perm: Vec<usize> = (0..m).collect();
94        let mut odd = false;
95        if m == n && n <= TINY {
96            // Tiny squares unrolled, their loops of known lengths.
97            match n {
98                0 => {}
99                1 => odd = tiny::<T, 1>(&mut lu.data, &mut perm),
100                2 => odd = tiny::<T, 2>(&mut lu.data, &mut perm),
101                3 => odd = tiny::<T, 3>(&mut lu.data, &mut perm),
102                4 => odd = tiny::<T, 4>(&mut lu.data, &mut perm),
103                5 => odd = tiny::<T, 5>(&mut lu.data, &mut perm),
104                6 => odd = tiny::<T, 6>(&mut lu.data, &mut perm),
105                7 => odd = tiny::<T, 7>(&mut lu.data, &mut perm),
106                8 => odd = tiny::<T, 8>(&mut lu.data, &mut perm),
107                9 => odd = tiny::<T, 9>(&mut lu.data, &mut perm),
108                10 => odd = tiny::<T, 10>(&mut lu.data, &mut perm),
109                11 => odd = tiny::<T, 11>(&mut lu.data, &mut perm),
110                12 => odd = tiny::<T, 12>(&mut lu.data, &mut perm),
111                13 => odd = tiny::<T, 13>(&mut lu.data, &mut perm),
112                14 => odd = tiny::<T, 14>(&mut lu.data, &mut perm),
113                15 => odd = tiny::<T, 15>(&mut lu.data, &mut perm),
114                _ => odd = tiny::<T, 16>(&mut lu.data, &mut perm),
115            }
116            return Lu { lu, perm, odd };
117        }
118        with_pivots(m.min(n), |pivots| {
119            factor_all(&mut lu.data, m, n, pivots);
120            for (i, &p) in pivots.iter().enumerate() {
121                if p != i {
122                    perm.swap(i, p);
123                    odd = !odd;
124                }
125            }
126        });
127        Lu { lu, perm, odd }
128    }
129
130    /// The determinant `det(A)`, from the LU decomposition: the product of the pivots, with the
131    /// sign of the permutation. The determinant of the `0 × 0` matrix is 1.
132    ///
133    /// # Panics
134    ///
135    /// Panics if the matrix is not square.
136    #[track_caller]
137    pub fn det(&self) -> T {
138        assert!(self.is_square(), "the matrix is not square");
139        self.lu().det()
140    }
141
142    /// The inverse `A⁻¹`, from the LU decomposition, or `None` when a pivot is zero: the matrix
143    /// is singular. A matrix that is singular only up to rounding gives an inverse with huge
144    /// elements; its [condition number](Matrix::cond) tells how far to trust it.
145    ///
146    /// # Panics
147    ///
148    /// Panics if the matrix is not square.
149    #[track_caller]
150    pub fn inverse(&self) -> Option<Self> {
151        assert!(self.is_square(), "the matrix is not square");
152        match self.rows() {
153            0 => Some(self.clone()),
154            1 => tiny_inverse_inline::<T, 1>(self),
155            2 => tiny_inverse_inline::<T, 2>(self),
156            3 => tiny_inverse_inline::<T, 3>(self),
157            4 => tiny_inverse_inline::<T, 4>(self),
158            5 => tiny_inverse::<T, 5>(self),
159            6 => tiny_inverse::<T, 6>(self),
160            7 => tiny_inverse::<T, 7>(self),
161            8 => tiny_inverse::<T, 8>(self),
162            9 => tiny_inverse::<T, 9>(self),
163            10 => tiny_inverse::<T, 10>(self),
164            11 => tiny_inverse::<T, 11>(self),
165            12 => tiny_inverse::<T, 12>(self),
166            n => {
167                // The factors in working memory, their row swaps kept as they are, so that the
168                // inverse is the only allocation.
169                with_scratch(n * n, |lu| {
170                    lu.copy_from_slice(&self.data);
171                    with_pivots(n, |pivots| {
172                        factor_all(lu, n, n, pivots);
173                        if (0..n).any(|i| lu[i + i * n] == T::zero()) {
174                            return None;
175                        }
176                        let mut x = Matrix::zeros(n, n);
177                        inverse_from(lu, n, &mut x.data);
178                        // Column `j` of `X·P` is column `pivots[j]` of `X` and the other way
179                        // round, from the last swap to the first.
180                        for (j, &p) in pivots.iter().enumerate().rev() {
181                            if p != j {
182                                let (a, b) = x.data.split_at_mut(p * n);
183                                a[j * n..j * n + n].swap_with_slice(&mut b[..n]);
184                            }
185                        }
186                        Some(x)
187                    })
188                })
189            }
190        }
191    }
192
193    /// The solution `X` of `A·X = B`, from the LU decomposition, or `None` when the matrix is
194    /// singular.
195    ///
196    /// # Panics
197    ///
198    /// Panics if the matrix is not square, or `B` has not as many rows.
199    #[track_caller]
200    pub fn solve(&self, b: &Matrix<T>) -> Option<Self> {
201        assert!(self.is_square(), "the matrix is not square");
202        self.lu().solve(b)
203    }
204
205    /// `Aⁿ` for any integer `n`: [`pow`](Matrix::pow) of the inverse for a negative `n`, `None`
206    /// when that doesn't exist.
207    ///
208    /// # Panics
209    ///
210    /// Panics if the matrix is not square.
211    #[track_caller]
212    pub fn powi(&self, n: i32) -> Option<Self> {
213        if n >= 0 {
214            Some(self.pow(n.unsigned_abs()))
215        } else {
216            Some(self.inverse()?.pow(n.unsigned_abs()))
217        }
218    }
219}
220
221impl<T: Scalar> Lu<T> {
222    /// `L`, `m × min(m, n)`, lower triangular with ones on the diagonal.
223    pub fn l(&self) -> Matrix<T> {
224        let (m, n) = self.lu.shape();
225        Matrix::from_fn(m, m.min(n), |i, j| match i.cmp(&j) {
226            core::cmp::Ordering::Greater => self.lu[(i, j)],
227            core::cmp::Ordering::Equal => T::one(),
228            core::cmp::Ordering::Less => T::zero(),
229        })
230    }
231
232    /// `U`, `min(m, n) × n`, upper triangular.
233    pub fn u(&self) -> Matrix<T> {
234        let (m, n) = self.lu.shape();
235        Matrix::from_fn(m.min(n), n, |i, j| {
236            if i <= j { self.lu[(i, j)] } else { T::zero() }
237        })
238    }
239
240    /// `P`, the `m × m` permutation matrix with `P·A = L·U`.
241    pub fn p(&self) -> Matrix<T> {
242        let m = self.perm.len();
243        Matrix::from_fn(m, m, |i, j| {
244            if self.perm[i] == j {
245                T::one()
246            } else {
247                T::zero()
248            }
249        })
250    }
251
252    /// The permutation: row `i` of `P·A` is row `permutation()[i]` of `A`.
253    pub fn permutation(&self) -> &[usize] {
254        &self.perm
255    }
256
257    /// Whether a pivot is zero, which for a square matrix makes it singular.
258    pub fn is_singular(&self) -> bool {
259        let (m, n) = self.lu.shape();
260        (0..m.min(n)).any(|i| self.lu[(i, i)] == T::zero())
261    }
262
263    /// The determinant of the square matrix: the product of the pivots, negated for an odd
264    /// permutation.
265    ///
266    /// # Panics
267    ///
268    /// Panics if the matrix is not square.
269    #[track_caller]
270    pub fn det(&self) -> T {
271        assert!(self.lu.is_square(), "the matrix is not square");
272        let mut d = T::one();
273        for i in 0..self.lu.rows() {
274            d *= self.lu[(i, i)];
275        }
276        if self.odd { -d } else { d }
277    }
278
279    /// The solution `X` of `A·X = B` for the square matrix, or `None` when it is singular.
280    ///
281    /// # Panics
282    ///
283    /// Panics if the matrix is not square, or `B` has not as many rows.
284    #[track_caller]
285    pub fn solve(&self, b: &Matrix<T>) -> Option<Matrix<T>> {
286        assert!(self.lu.is_square(), "the matrix is not square");
287        let n = self.lu.rows();
288        assert_eq!(b.rows(), n, "the right-hand side has not as many rows");
289        if self.is_singular() {
290            return None;
291        }
292        let mut x = Matrix::from_fn(n, b.cols(), |i, j| b[(self.perm[i], j)]);
293        let cols = x.cols();
294        trsm(
295            Side::Left,
296            Uplo::Lower,
297            Op::N,
298            true,
299            n,
300            cols,
301            &self.lu.data,
302            n,
303            &mut x.data,
304            n,
305        );
306        trsm(
307            Side::Left,
308            Uplo::Upper,
309            Op::N,
310            false,
311            n,
312            cols,
313            &self.lu.data,
314            n,
315            &mut x.data,
316            n,
317        );
318        Some(x)
319    }
320
321    /// The inverse of the square matrix, or `None` when it is singular.
322    ///
323    /// # Panics
324    ///
325    /// Panics if the matrix is not square.
326    #[track_caller]
327    pub fn inverse(&self) -> Option<Matrix<T>> {
328        assert!(self.lu.is_square(), "the matrix is not square");
329        if self.is_singular() {
330            return None;
331        }
332        // `A⁻¹ = U⁻¹·L⁻¹·P`, as LAPACK's `getri`: `U⁻¹`, then `X·L = U⁻¹` for `X`, then the
333        // columns of `X` permuted; about `4n³/3` operations, nearly all of them in products of
334        // blocks.
335        let n = self.lu.rows();
336        let mut x = Matrix::zeros(n, n);
337        inverse_from(&self.lu.data, n, &mut x.data);
338        // Column `perm[j]` of `X·P` is column `j` of `X`: the cycles of the permutation, a swap
339        // of whole columns at a time.
340        let mut to = self.perm.clone();
341        for j in 0..n {
342            while to[j] != j {
343                let k = to[j];
344                let (a, b) = x.data.split_at_mut(j.max(k) * n);
345                a[j.min(k) * n..j.min(k) * n + n].swap_with_slice(&mut b[..n]);
346                to.swap(j, k);
347            }
348        }
349        Some(x)
350    }
351}
352
353/// Writes `X = U⁻¹·L⁻¹` for the factors `lu` of order `n` to `x`, which holds zeros: `U⁻¹`,
354/// then `X·L = U⁻¹` for `X`.
355fn inverse_from<T: Scalar>(lu: &[T], n: usize, x: &mut [T]) {
356    upper_inverse(lu, n, x, n, n);
357    trsm(Side::Right, Uplo::Lower, Op::N, true, n, n, lu, n, x, n);
358}
359
360/// Triangles up to this order are inverted by a solve against the identity.
361const INVERSE_LEAF: usize = 16;
362
363/// Writes the inverse of the upper triangle of order `n` at the start of `u` to the upper
364/// triangle of `w`, recursively: `W11 = U11⁻¹`, `W22 = U22⁻¹` and `W12 = -U11⁻¹·U12·U22⁻¹`, the
365/// last by two triangular solves from the triangles of `U` itself, so that no block is read
366/// where another is written. The lower triangle of `w` is left as it is.
367fn upper_inverse<T: Scalar>(u: &[T], ldu: usize, w: &mut [T], ldw: usize, n: usize) {
368    if n <= INVERSE_LEAF {
369        for j in 0..n {
370            for i in 0..=j {
371                w[i + j * ldw] = if i == j { T::one() } else { T::zero() };
372            }
373        }
374        trsm(Side::Left, Uplo::Upper, Op::N, false, n, n, u, ldu, w, ldw);
375        return;
376    }
377    let n1 = split(n);
378    let n2 = n - n1;
379    let (w1, w2) = w.split_at_mut(n1 * ldw);
380    for j in 0..n2 {
381        let (src, dst) = (&u[(n1 + j) * ldu..], &mut w2[j * ldw..]);
382        for (d, s) in dst[..n1].iter_mut().zip(&src[..n1]) {
383            *d = -*s;
384        }
385    }
386    let u22 = &u[n1 + n1 * ldu..];
387    trsm(
388        Side::Left,
389        Uplo::Upper,
390        Op::N,
391        false,
392        n1,
393        n2,
394        u,
395        ldu,
396        w2,
397        ldw,
398    );
399    trsm(
400        Side::Right,
401        Uplo::Upper,
402        Op::N,
403        false,
404        n1,
405        n2,
406        u22,
407        ldu,
408        w2,
409        ldw,
410    );
411    upper_inverse(u, ldu, w1, ldw, n1);
412    upper_inverse(u22, ldu, &mut w2[n1..], ldw, n2);
413}
414
415/// [`Matrix::inverse`] of an `N × N` matrix, `N` known, so that its loops unroll: the LU
416/// decomposition of a copy on the stack, as [`tiny`] finds it, then `U⁻¹`, `X·L = U⁻¹` and the
417/// permutation of the columns, all on the stack, then the only allocation, the inverse.
418///
419/// The decomposition stops at the first zero pivot, and does not skip the updates by zeros: these
420/// change nothing that is finite, and the branches cost more than they save here. The pivots are
421/// chosen and the rows swapped by selects, not branches, which keeps the whole matrix in
422/// registers. Orders from 5 are functions of their own, as one function holding them all compiled
423/// worse; up to 4, the call itself cost a fifth of the time, so they are inlined.
424#[inline(never)]
425fn tiny_inverse<T: Scalar, const N: usize>(a: &Matrix<T>) -> Option<Matrix<T>> {
426    tiny_inverse_inline::<T, N>(a)
427}
428
429/// [`tiny_inverse`], inlined. Its loops index rows and columns as the elimination is written: every
430/// index is known once they unroll.
431#[inline(always)]
432#[allow(clippy::needless_range_loop)]
433fn tiny_inverse_inline<T: Scalar, const N: usize>(a: &Matrix<T>) -> Option<Matrix<T>> {
434    let mut lu = [[T::zero(); N]; N];
435    lu.as_flattened_mut().copy_from_slice(&a.data);
436    let mut perm: [usize; N] = core::array::from_fn(|i| i);
437    let mut r = [T::zero(); N];
438    for j in 0..N {
439        // The first element of the largest `|re| + |im|`, the first itself when it is NaN.
440        let mut p = j;
441        let mut best = lu[j][j].abs1();
442        for i in j + 1..N {
443            let v = lu[j][i].abs1();
444            let greater = v > best;
445            best = select_unpredictable(greater, v, best);
446            p = select_unpredictable(greater, i, p);
447        }
448        // The swap by selects, every index known: the whole matrix stays in registers.
449        for i in j + 1..N {
450            let swap = p == i;
451            for col in &mut lu {
452                let (x, y) = (col[j], col[i]);
453                col[j] = select_unpredictable(swap, y, x);
454                col[i] = select_unpredictable(swap, x, y);
455            }
456            let (x, y) = (perm[j], perm[i]);
457            perm[j] = select_unpredictable(swap, y, x);
458            perm[i] = select_unpredictable(swap, x, y);
459        }
460        let d = lu[j][j];
461        if d == T::zero() {
462            return None;
463        }
464        r[j] = T::one() / d;
465        if r[j].is_finite() {
466            for i in j + 1..N {
467                lu[j][i] *= r[j];
468            }
469        } else {
470            // A reciprocal that overflows, of a pivot near the bottom of the range: divide.
471            for i in j + 1..N {
472                lu[j][i] /= d;
473            }
474        }
475        for c in j + 1..N {
476            let u = lu[c][j];
477            for i in j + 1..N {
478                lu[c][i] = lu[c][i].mul_sub(lu[j][i], u);
479            }
480        }
481    }
482    // `W = U⁻¹` by columns, each by back substitution with the reciprocals of the diagonal.
483    let mut w = [[T::zero(); N]; N];
484    for j in 0..N {
485        w[j][j] = r[j];
486        for i in (0..j).rev() {
487            let mut s = T::zero();
488            for k in i + 1..=j {
489                s = s.mul_sub(lu[k][i], w[j][k]);
490            }
491            w[j][i] = s * r[i];
492        }
493    }
494    // `X·L = W` from the last column: `X[:, j] = W[:, j] - Σ_{k > j} X[:, k]·L[k, j]`.
495    let mut x = [[T::zero(); N]; N];
496    for j in (0..N).rev() {
497        let mut col = w[j];
498        for k in j + 1..N {
499            let l = lu[j][k];
500            for (c, x) in col.iter_mut().zip(&x[k]) {
501                *c = c.mul_sub(*x, l);
502            }
503        }
504        x[j] = col;
505    }
506    // Column `perm[j]` of `X·P` is column `j` of `X`.
507    let mut inverse = vec![T::zero(); N * N];
508    for (j, &p) in perm.iter().enumerate() {
509        inverse[p * N..p * N + N].copy_from_slice(&x[j]);
510    }
511    Some(Matrix::from_column_major(N, N, inverse))
512}
513
514/// Pivots up to this many are kept on the stack.
515const STACK_PIVOTS: usize = 32;
516
517/// Runs `f` on `k` zeros, for the row swaps of a factorization.
518fn with_pivots<X>(k: usize, f: impl FnOnce(&mut [usize]) -> X) -> X {
519    if k <= STACK_PIVOTS {
520        f(&mut [0; STACK_PIVOTS][..k])
521    } else {
522        f(&mut vec![0; k])
523    }
524}
525
526/// Factors the `m × n` matrix `a` in place, `P·A = L·U`: `pivots[i]` is the row swapped with
527/// row `i`, in turn, for `i < min(m, n)`.
528fn factor_all<T: Scalar>(a: &mut [T], m: usize, n: usize, pivots: &mut [usize]) {
529    let k = m.min(n);
530    if k == 0 {
531        return;
532    }
533    if k <= UNBLOCKED {
534        unblocked(a, m, k, m, pivots);
535    } else {
536        with_scratch(work_size(k), |work| factor(a, m, k, m, pivots, work));
537    }
538    if n > k {
539        // The columns right of the square part: the same row swaps, then `L⁻¹`.
540        let (left, right) = a.split_at_mut(k * m);
541        swap_rows(right, m, n - k, pivots, 0);
542        trsm(
543            Side::Left,
544            Uplo::Lower,
545            Op::N,
546            true,
547            k,
548            n - k,
549            left,
550            m,
551            right,
552            m,
553        );
554    }
555}
556
557/// Columns narrower than this are factored without recursion.
558const NARROW: usize = 16;
559
560/// Matrices of up to this many columns are factored without recursion at all.
561const UNBLOCKED: usize = 24;
562
563/// Where the recursion splits `n` columns: the first part a multiple of [`NARROW`], so that the
564/// unblocked panels are full.
565fn split(n: usize) -> usize {
566    (n / 2).next_multiple_of(NARROW).min(n - 1)
567}
568
569/// The working memory of [`factor`] of `n` columns: the dense copy of the largest `U12` of the
570/// recursion, which an uneven split may leave below its top.
571fn work_size(n: usize) -> usize {
572    if n <= NARROW {
573        return 0;
574    }
575    let n1 = split(n);
576    (n1 * (n - n1)).max(work_size(n1)).max(work_size(n - n1))
577}
578
579/// Factors the `m × n` block at the start of `a`, `n ≤ m`, in place, Toledo's way: the left
580/// half, then the right half updated by it, each half the same way down to narrow panels.
581/// `pivots[i]` is the row swapped with row `i`, within the block. `work` holds the dense copy of
582/// `U12` for the product that updates `A22`: at least [`work_size`] elements.
583fn factor<T: Scalar>(
584    a: &mut [T],
585    m: usize,
586    n: usize,
587    ld: usize,
588    pivots: &mut [usize],
589    work: &mut [T],
590) {
591    if n <= NARROW {
592        unblocked(a, m, n, ld, pivots);
593        return;
594    }
595    let n1 = split(n);
596    let n2 = n - n1;
597    factor(a, m, n1, ld, &mut pivots[..n1], work);
598    let (left, right) = a.split_at_mut(n1 * ld);
599    swap_rows(right, ld, n2, &pivots[..n1], 0);
600    // `U12 = L11⁻¹·A12`, then `A22 -= L21·U12`, `U12` read from its copy.
601    let u12 = &mut work[..n1 * n2];
602    trsm_into(
603        Uplo::Lower,
604        Op::N,
605        true,
606        n1,
607        n2,
608        left,
609        ld,
610        right,
611        ld,
612        u12,
613        n1,
614    );
615    gemm(
616        m - n1,
617        n2,
618        n1,
619        -T::one(),
620        &left[n1..],
621        ld,
622        Op::N,
623        u12,
624        n1,
625        Op::N,
626        T::one(),
627        &mut right[n1..],
628        ld,
629    );
630    factor(&mut right[n1..], m - n1, n2, ld, &mut pivots[n1..], work);
631    for p in &mut pivots[n1..] {
632        *p += n1;
633    }
634    swap_rows(left, ld, n1, &pivots[n1..], n1);
635}
636
637/// Factors a narrow block by rank-one updates, choosing as pivot the element of the largest
638/// `|re| + |im|` in each column, as LAPACK does. A zero column is left as it is.
639fn unblocked<T: Scalar>(a: &mut [T], m: usize, n: usize, ld: usize, pivots: &mut [usize]) {
640    for j in 0..n {
641        let p = j + pivot(&a[j * ld + j..j * ld + m]);
642        pivots[j] = p;
643        if p != j {
644            for c in 0..n {
645                a.swap(j + c * ld, p + c * ld);
646            }
647        }
648        let (left, right) = a.split_at_mut(((j + 1) * ld).min(a.len()));
649        let col = &mut left[j * ld + j..j * ld + m];
650        let (d, l) = col.split_first_mut().unwrap();
651        if *d == T::zero() {
652            continue;
653        }
654        let r = T::one() / *d;
655        if r.is_finite() {
656            for x in l.iter_mut() {
657                *x *= r;
658            }
659        } else {
660            // A reciprocal that overflows, of a pivot near the bottom of the range: divide, as
661            // LAPACK's `getf2` does.
662            for x in l.iter_mut() {
663                *x /= *d;
664            }
665        }
666        for c in 0..n - j - 1 {
667            let col = &mut right[c * ld + j..c * ld + m];
668            let (u, x) = col.split_first_mut().unwrap();
669            let u = *u;
670            if u == T::zero() {
671                continue;
672            }
673            for (x, l) in x.iter_mut().zip(l.iter()) {
674                *x = x.mul_sub(*l, u);
675            }
676        }
677    }
678}
679
680/// The order up to which square matrices are factored by [`tiny`]: past it, [`unblocked`] ran
681/// as fast.
682const TINY: usize = 16;
683
684/// [`unblocked`] of an `N × N` matrix, `N` known, so that its loops unroll; the row swaps are
685/// applied to `perm`, and the result is whether they make an odd permutation. The matrix is
686/// copied to the stack, its pivots chosen and its rows swapped by selects, not branches, every
687/// index known, so that it stays in registers: a 5 × 5 decomposition so took 67 ns against 88.
688/// Each order is a function of its own. Its loops index rows and columns as the elimination is
689/// written: every index is known once they unroll.
690#[inline(never)]
691#[allow(clippy::needless_range_loop)]
692fn tiny<T: Scalar, const N: usize>(data: &mut [T], perm: &mut [usize]) -> bool {
693    let mut a = [[T::zero(); N]; N];
694    a.as_flattened_mut().copy_from_slice(&data[..N * N]);
695    let perm: &mut [usize; N] = (&mut perm[..N]).try_into().unwrap();
696    let mut odd = false;
697    for j in 0..N {
698        // The first element of the largest `|re| + |im|`, the first itself when it is NaN.
699        let mut p = j;
700        let mut best = a[j][j].abs1();
701        for i in j + 1..N {
702            let v = a[j][i].abs1();
703            let greater = v > best;
704            best = select_unpredictable(greater, v, best);
705            p = select_unpredictable(greater, i, p);
706        }
707        odd ^= p != j;
708        for i in j + 1..N {
709            let swap = p == i;
710            for col in &mut a {
711                let (x, y) = (col[j], col[i]);
712                col[j] = select_unpredictable(swap, y, x);
713                col[i] = select_unpredictable(swap, x, y);
714            }
715            let (x, y) = (perm[j], perm[i]);
716            perm[j] = select_unpredictable(swap, y, x);
717            perm[i] = select_unpredictable(swap, x, y);
718        }
719        let d = a[j][j];
720        if d == T::zero() {
721            continue;
722        }
723        let r = T::one() / d;
724        if r.is_finite() {
725            for i in j + 1..N {
726                a[j][i] *= r;
727            }
728        } else {
729            // A reciprocal that overflows, of a pivot near the bottom of the range: divide.
730            for i in j + 1..N {
731                a[j][i] /= d;
732            }
733        }
734        for c in j + 1..N {
735            let u = a[c][j];
736            if u == T::zero() {
737                continue;
738            }
739            for i in j + 1..N {
740                a[c][i] = a[c][i].mul_sub(a[j][i], u);
741            }
742        }
743    }
744    data[..N * N].copy_from_slice(a.as_flattened());
745    odd
746}
747
748/// The index of the pivot of a column: the first element of the largest `|re| + |im|`, the
749/// first element itself when it is NaN, as the scan from it with `>` finds.
750fn pivot<T: Scalar>(col: &[T]) -> usize {
751    let first = col[0].abs1();
752    if first.is_nan() {
753        return 0;
754    }
755    // The largest by `max`, which passes over NaNs, then where it first is.
756    let best = col.iter().fold(first, |b, x| b.max(x.abs1()));
757    col.iter().position(|x| x.abs1() == best).unwrap_or(0)
758}
759
760/// Swaps rows `first + i` and `pivots[i]` (an absolute row) of the `n` columns of `a`, in order,
761/// a column at a time.
762fn swap_rows<T>(a: &mut [T], ld: usize, n: usize, pivots: &[usize], first: usize) {
763    for c in 0..n {
764        let col = &mut a[c * ld..];
765        for (i, &p) in pivots.iter().enumerate() {
766            let i = first + i;
767            if p != i {
768                col.swap(i, p);
769            }
770        }
771    }
772}