1use crate::assert;
2use crate::internal_prelude::*;
3use linalg::householder;
4use linalg::matmul::triangular::BlockStructure;
5use linalg::matmul::{self, dot};
6#[derive(Copy, Clone, Debug)]
8pub struct TridiagParams {
9 pub par_threshold: usize,
11 #[doc(hidden)]
12 pub non_exhaustive: NonExhaustive,
13}
14impl<T: ComplexField> Auto<T> for TridiagParams {
15 fn auto() -> Self {
16 Self {
17 par_threshold: 192 * 256,
18 non_exhaustive: NonExhaustive(()),
19 }
20 }
21}
22pub fn tridiag_in_place_scratch<T: ComplexField>(
25 dim: usize,
26 par: Par,
27 params: Spec<TridiagParams, T>,
28) -> StackReq {
29 _ = par;
30 _ = params;
31 StackReq::all_of(&[
32 temp_mat_scratch::<T>(dim, 1).array(2),
33 temp_mat_scratch::<T>(dim, par.degree()),
34 ])
35}
36fn tridiag_fused_op_simd<T: ComplexField>(
37 A: MatMut<'_, T, usize, usize, ContiguousFwd>,
38 y2: ColMut<'_, T, usize>,
39 z2: ColMut<'_, T, usize, ContiguousFwd>,
40 ry2: ColRef<'_, T, usize>,
41 rz2: ColRef<'_, T, usize, ContiguousFwd>,
42 u0: ColRef<'_, T, usize, ContiguousFwd>,
43 u1: ColRef<'_, T, usize>,
44 u2: ColRef<'_, T, usize>,
45 v2: ColRef<'_, T, usize, ContiguousFwd>,
46 f: T,
47 align: usize,
48) {
49 struct Impl<'a, 'M, 'N, T: ComplexField> {
50 A: MatMut<'a, T, Dim<'M>, Dim<'N>, ContiguousFwd>,
51 y2: ColMut<'a, T, Dim<'N>>,
52 z2: ColMut<'a, T, Dim<'M>, ContiguousFwd>,
53 ry2: ColRef<'a, T, Dim<'N>>,
54 rz2: ColRef<'a, T, Dim<'M>, ContiguousFwd>,
55 u0: ColRef<'a, T, Dim<'M>, ContiguousFwd>,
56 u1: ColRef<'a, T, Dim<'N>>,
57 u2: ColRef<'a, T, Dim<'N>>,
58 v2: ColRef<'a, T, Dim<'M>, ContiguousFwd>,
59 f: T,
60 align: usize,
61 }
62 impl<'a, 'M, 'N, T: ComplexField> pulp::WithSimd for Impl<'a, 'M, 'N, T> {
63 type Output = ();
64
65 #[inline(always)]
66 fn with_simd<S: pulp::Simd>(self, simd: S) -> Self::Output {
67 let Self {
68 mut A,
69 mut y2,
70 mut z2,
71 ry2,
72 rz2,
73 u0,
74 u1,
75 u2,
76 v2,
77 ref f,
78 mut align,
79 } = self;
80 let simd = T::simd_ctx(simd);
81 let (m, n) = A.shape();
82 {
83 let simd = SimdCtx::<T, S>::new_align(simd, m, align);
84 simd_iter!(for i in [simd.indices()] {
85 simd.write(z2.rb_mut(), i, simd.zero());
86 });
87 }
88 for j in n.indices() {
89 let i = m.idx_inc(*j);
90 with_dim!(m, *m - *j);
91 let simd = SimdCtx::<T, S>::new_align(simd, m, align);
92 align -= 1;
93 let mut A = A.rb_mut().col_mut(j).subrows_mut(i, m);
94 let mut z = z2.rb_mut().subrows_mut(i, m);
95 let rz = rz2.subrows(i, m);
96 let ua = u0.subrows(i, m);
97 let v = v2.subrows(i, m);
98 let y = y2.rb_mut().at_mut(j);
99 let ry = simd.splat(&(-&ry2[j]));
100 let ub = simd.splat(&(-&u1[j]));
101 let uc = simd.splat(&(f * &u2[j]));
102 let mut acc = [simd.zero(); 4];
103 simd_iter!(for (IDX, i) in [simd.batch_indices(); 4] {
104 let mut a = simd.read(A.rb(), i);
105 a = simd.conj_mul_add(ry, simd.read(ua, i), a);
106 a = simd.conj_mul_add(ub, simd.read(rz, i), a);
107 simd.write(A.rb_mut(), i, a);
108 let tmp = simd.read(z.rb(), i);
109 simd.write(z.rb_mut(), i, simd.mul_add(a, uc, tmp));
110 acc[IDX] = simd.conj_mul_add(a, simd.read(v, i), acc[IDX]);
111 });
112 let acc0 = simd.add(acc[0], acc[1]);
113 let acc2 = simd.add(acc[2], acc[3]);
114 let acc0 = simd.add(acc0, acc2);
115 let acc0 = simd.reduce_sum(acc0);
116 let i0 = m.idx(0);
117 *y = f * (acc0 - &A[i0] * &v[i0]);
118 }
119 }
120 }
121 with_dim!(M, A.nrows());
122 with_dim!(N, A.ncols());
123 dispatch!(
124 Impl {
125 A: A.as_shape_mut(M, N),
126 y2: y2.as_row_shape_mut(N),
127 z2: z2.as_row_shape_mut(M),
128 ry2: ry2.as_row_shape(N),
129 rz2: rz2.as_row_shape(M),
130 u0: u0.as_row_shape(M),
131 u1: u1.as_row_shape(N),
132 u2: u2.as_row_shape(N),
133 v2: v2.as_row_shape(M),
134 f,
135 align,
136 },
137 Impl,
138 T
139 )
140}
141fn tridiag_fused_op<T: ComplexField>(
142 A: MatMut<'_, T>,
143 y2: ColMut<'_, T>,
144 z2: ColMut<'_, T>,
145 ry2: ColRef<'_, T>,
146 rz2: ColRef<'_, T>,
147 u0: ColRef<'_, T>,
148 u1: ColRef<'_, T>,
149 u2: ColRef<'_, T>,
150 v2: ColRef<'_, T>,
151 f: T,
152 align: usize,
153) {
154 let mut A = A;
155 let mut z2 = z2;
156 if const { T::SIMD_CAPABILITIES.is_simd() } {
157 if let (Some(A), Some(z2), Some(rz2), Some(u0), Some(v2)) = (
158 A.rb_mut().try_as_col_major_mut(),
159 z2.rb_mut().try_as_col_major_mut(),
160 rz2.try_as_col_major(),
161 u0.try_as_col_major(),
162 v2.try_as_col_major(),
163 ) {
164 tridiag_fused_op_simd(
165 A, y2, z2, ry2, rz2, u0, u1, u2, v2, f, align,
166 );
167 } else {
168 tridiag_fused_op_fallback(A, y2, z2, ry2, rz2, u0, u1, u2, v2, f);
169 }
170 } else {
171 tridiag_fused_op_fallback(A, y2, z2, ry2, rz2, u0, u1, u2, v2, f);
172 }
173}
174fn tridiag_fused_op_fallback<T: ComplexField>(
175 A: MatMut<'_, T>,
176 y2: ColMut<'_, T>,
177 z2: ColMut<'_, T>,
178 ry2: ColRef<'_, T>,
179 rz2: ColRef<'_, T>,
180 u0: ColRef<'_, T>,
181 u1: ColRef<'_, T>,
182 u2: ColRef<'_, T>,
183 v2: ColRef<'_, T>,
184 f: T,
185) {
186 let par = Par::Seq;
187 let mut A = A;
188 let mut y2 = y2;
189 let n = A.ncols();
190 let (mut A0, mut A1) = A.rb_mut().split_at_row_mut(n);
191 let (u00, u01) = u0.split_at_row(n);
192 let (v20, v21) = v2.split_at_row(n);
193 let (mut z20, mut z21) = z2.split_at_row_mut(n);
194 let (rz20, rz21) = rz2.split_at_row(n);
195 matmul::triangular::matmul(
196 A0.rb_mut(),
197 BlockStructure::TriangularLower,
198 Accum::Add,
199 u00,
200 BlockStructure::Rectangular,
201 ry2.adjoint(),
202 BlockStructure::Rectangular,
203 -one::<T>(),
204 par,
205 );
206 matmul::triangular::matmul(
207 A0.rb_mut(),
208 BlockStructure::TriangularLower,
209 Accum::Add,
210 rz20,
211 BlockStructure::Rectangular,
212 u1.adjoint(),
213 BlockStructure::Rectangular,
214 -one::<T>(),
215 par,
216 );
217 matmul::matmul(
218 A1.rb_mut(),
219 Accum::Add,
220 u01,
221 ry2.adjoint(),
222 -one::<T>(),
223 par,
224 );
225 matmul::matmul(
226 A1.rb_mut(),
227 Accum::Add,
228 rz21,
229 u1.adjoint(),
230 -one::<T>(),
231 par,
232 );
233 matmul::triangular::matmul(
234 z20.rb_mut(),
235 BlockStructure::Rectangular,
236 Accum::Replace,
237 A0.rb(),
238 BlockStructure::TriangularLower,
239 u2,
240 BlockStructure::Rectangular,
241 f.clone(),
242 par,
243 );
244 matmul::triangular::matmul(
245 y2.rb_mut(),
246 BlockStructure::Rectangular,
247 Accum::Replace,
248 A0.rb().adjoint(),
249 BlockStructure::StrictTriangularUpper,
250 v20,
251 BlockStructure::Rectangular,
252 f.clone(),
253 par,
254 );
255 matmul::matmul(z21.rb_mut(), Accum::Replace, A1.rb(), u2, f.clone(), par);
256 matmul::matmul(
257 y2.rb_mut(),
258 Accum::Add,
259 A1.rb().adjoint(),
260 v21,
261 f.clone(),
262 par,
263 );
264}
265pub fn tridiag_in_place<T: ComplexField>(
275 A: MatMut<'_, T>,
276 householder: MatMut<'_, T>,
277 par: Par,
278 stack: &mut MemStack,
279 params: Spec<TridiagParams, T>,
280) {
281 let params = params.config;
282 let mut A = A;
283 let mut H = householder;
284 let mut par = par;
285 let n = A.nrows();
286 let b = H.nrows();
287 assert!(H.ncols() == n.saturating_sub(1));
288 if n == 0 {
289 return;
290 }
291 alloca!('stack: {
292 let mut y = unsafe { col![uninit::<T>, n] };
293 let mut w = unsafe { col![uninit::<T>, n] };
294 let mut z = unsafe { mat![uninit::<T>, n, par.degree()] };
295 });
296 {
297 let mut H = H.rb_mut().row_mut(0);
298 for k in 0..n {
299 let (_, A01, A10, A11) = A.rb_mut().split_at_mut(k, k);
300 let (_, _) = A01.split_first_col().unwrap();
301 let (_, A20) = A10.split_first_row_mut().unwrap();
302 let (mut A11, _, A21, mut A22) = A11.split_at_mut(1, 1);
303 let mut A21 = A21.col_mut(0);
304 let a11 = &mut A11[(0, 0)];
305 let (y1, mut y2) =
306 y.rb_mut().split_at_row_mut(k).1.split_at_row_mut(1);
307 if k > 0 {
308 let y1 = y1[0].copy();
309 let p = k - 1;
310 let u2 = A20.rb().col(p);
311 *a11 -= &y1 + y1.conj();
312 z!(A21.rb_mut(), u2, y2.rb()).for_each(
313 |uz!(a, u, y): Zip!(&mut T, &T, &T)| {
314 *a -= y1.conj() * u + y;
315 },
316 );
317 }
318 if k + 1 == n {
319 break;
320 }
321 let rem = n - k - 1;
322 if rem * rem / 2 < params.par_threshold {
323 par = Par::Seq;
324 }
325 let k1 = k + 1;
326 let tau_inv;
327 {
328 let (mut a11, mut x2) = A21.rb_mut().split_at_row_mut(1);
329 let a11 = &mut a11[0];
330 let householder::HouseholderInfo { tau, .. } =
331 householder::make_householder_in_place(a11, x2.rb_mut());
332 tau_inv = tau.real().recip();
333 H[k] = tau.to_cplx();
334 let mut z2 = z.rb_mut().split_at_row_mut(k + 2).1;
335 let mut w2 = w.rb_mut().split_at_row_mut(k + 2).1;
336 let (mut y1, mut y2) = y2.rb_mut().split_at_row_mut(1);
337 let y1 = &mut y1[0];
338 let (A1, A2) = A22.rb_mut().split_at_row_mut(1);
339 let A1 = A1.row_mut(0);
340 let (mut a11, _) = A1.split_at_col_mut(1);
341 let a11 = &mut a11[0];
342 let (A21, mut A22) = A2.split_at_col_mut(1);
343 let mut A21 = A21.col_mut(0);
344 if k > 0 {
345 let p = k - 1;
346 let (u1, u2) = A20.rb().col(p).split_at_row(1);
347 let ref u1 = u1[0].copy();
348 *a11 -= u1 * y1.conj() + &*y1 * u1.conj();
349 z!(A21.rb_mut(), u2.rb(), y2.rb()).for_each(
350 |uz!(a, u, y): Zip!(&mut T, &T, &T)| {
351 *a -= u * y1.conj() + &*y * u1.conj();
352 },
353 );
354 w2.copy_from(y2.rb());
355 match par {
356 Par::Seq => {
357 let mut z2 = z2.rb_mut().col_mut(0);
358 tridiag_fused_op(
359 A22.rb_mut(),
360 y2.rb_mut(),
361 z2.rb_mut(),
362 w2.rb(),
363 w2.rb(),
364 u2.rb(),
365 u2.rb(),
366 x2.rb(),
367 x2.rb(),
368 tau_inv.to_cplx(),
369 simd_align(k1 + 1),
370 );
371 z!(y2.rb_mut(), z2.rb()).for_each(
372 |uz!(y, z): Zip!(&mut T, &T)| *y += z,
373 );
374 },
375 #[cfg(feature = "rayon")]
376 Par::Rayon(nthreads) => {
377 use rayon::prelude::*;
378 let nthreads = nthreads.get();
379 let mut z2 = z2.rb_mut().subcols_mut(0, nthreads);
380 let n2 = A22.ncols();
381 assert!((n2 as u64) < (1u64 << 50));
382 let idx_to_col_start = |idx: usize| {
383 let idx_as_percent =
384 idx as f64 / nthreads as f64;
385 let col_start_percent = 1.0f64
386 - libm::sqrt(1.0f64 - idx_as_percent);
387 (col_start_percent * n2 as f64) as usize
388 };
389 {
390 let A22 = A22.rb();
391 let y2 = y2.rb();
392 let ref f = tau_inv.to_cplx::<T>();
393 spindle::for_each(
394 nthreads,
395 z2.rb_mut().par_col_iter_mut().enumerate(),
396 |(idx, mut z2)| {
397 let first = idx_to_col_start(idx);
398 let last_col =
399 idx_to_col_start(idx + 1);
400 let nrows = n2 - first;
401 let ncols = last_col - first;
402 let mut A = unsafe {
403 A22.rb()
404 .subcols(first, ncols)
405 .subrows(first, nrows)
406 .const_cast()
407 };
408 {
409 let y2 = unsafe {
410 y2.subrows(first, ncols)
411 .const_cast()
412 };
413 let mut z2 = z2
414 .rb_mut()
415 .subrows_mut(first, nrows);
416 let ry2 =
417 w2.rb().subrows(first, ncols);
418 let rz2 =
419 w2.rb().subrows(first, nrows);
420 let u0 = u2.subrows(first, nrows);
421 let u1 = u2.subrows(first, ncols);
422 let u2 =
423 x2.rb().subrows(first, ncols);
424 let v2 =
425 x2.rb().subrows(first, nrows);
426 tridiag_fused_op(
427 A.rb_mut(),
428 y2,
429 z2.rb_mut(),
430 ry2,
431 rz2,
432 u0,
433 u1,
434 u2,
435 v2,
436 f.copy(),
437 n.next_power_of_two()
438 - (k1 + 1) - first,
439 );
440 }
441 z2.rb_mut()
442 .subrows_mut(0, first)
443 .fill(zero());
444 },
445 );
446 }
447 for z2 in z2.rb().col_iter() {
448 z!(y2.rb_mut(), z2.rb())
449 .for_each(|uz!(y, z)| *y += z);
450 }
451 },
452 }
453 } else {
454 matmul::triangular::matmul(
455 y2.rb_mut(),
456 BlockStructure::Rectangular,
457 Accum::Replace,
458 A22.rb(),
459 BlockStructure::TriangularLower,
460 x2.rb(),
461 BlockStructure::Rectangular,
462 tau_inv.to_cplx(),
463 par,
464 );
465 matmul::triangular::matmul(
466 y2.rb_mut(),
467 BlockStructure::Rectangular,
468 Accum::Add,
469 A22.rb().adjoint(),
470 BlockStructure::StrictTriangularUpper,
471 x2.rb(),
472 BlockStructure::Rectangular,
473 tau_inv.to_cplx(),
474 par,
475 );
476 }
477 z!(y2.rb_mut(), A21.rb()).for_each(
478 |uz!(y, a): Zip!(&mut T, &T)| {
479 *y += a.mul_real(&tau_inv);
480 },
481 );
482 *y1 = (&*a11
483 + dot::inner_prod(
484 A21.rb().transpose(),
485 Conj::Yes,
486 x2.rb(),
487 Conj::No,
488 ))
489 .mul_real(&tau_inv);
490 let ref b = (&*y1
491 + dot::inner_prod(
492 x2.rb().transpose(),
493 Conj::Yes,
494 y2.rb(),
495 Conj::No,
496 ))
497 .mul_pow2(from_f64::<T::Real>(0.5))
498 .mul_real(&tau_inv);
499 *y1 -= b;
500 z!(y2.rb_mut(), x2.rb()).for_each(
501 |uz!(y, u): Zip!(&mut T, &T)| {
502 *y -= b * u;
503 },
504 );
505 }
506 }
507 }
508 if n > 0 {
509 let n = n - 1;
510 let A = A.rb().submatrix(1, 0, n, n);
511 let mut H = H.rb_mut().subcols_mut(0, n);
512 let mut j = 0;
513 while j < n {
514 let b = Ord::min(b, n - j);
515 let mut H = H.rb_mut().submatrix_mut(0, j, b, b);
516 for k in 0..b {
517 H[(k, k)] = H[(0, k)].copy();
518 }
519 householder::upgrade_householder_factor(
520 H.rb_mut(),
521 A.submatrix(j, j, n - j, b),
522 b,
523 1,
524 par,
525 );
526 j += b;
527 }
528 }
529}
530#[cfg(test)]
531mod tests {
532 use super::*;
533 use crate::stats::prelude::*;
534 use crate::utils::approx::*;
535 use crate::{Mat, assert, c64};
536 use dyn_stack::MemBuffer;
537 #[test]
538 fn test_tridiag_real() {
539 let rng = &mut StdRng::seed_from_u64(0);
540 for n in [2, 3, 4, 8, 16] {
541 let A = CwiseMatDistribution {
542 nrows: n,
543 ncols: n,
544 dist: StandardNormal,
545 }
546 .rand::<Mat<f64>>(rng);
547 let A = A.rb() + A.adjoint();
548 let b = 3;
549 let mut H = Mat::zeros(b, n - 1);
550 let mut V = A.clone();
551 let mut V = V.as_mut();
552 tridiag_in_place(
553 V.rb_mut(),
554 H.rb_mut(),
555 Par::Seq,
556 MemStack::new(&mut MemBuffer::new(StackReq::all_of(&[
557 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<f64>(n - 1, b, n),
558 tridiag_in_place_scratch::<f64>(n, Par::Seq, default()),
559 ]))),
560 default(),
561 );
562 let mut A = A.clone();
563 let mut A = A.as_mut();
564 for iter in 0..2 {
565 let mut A = if iter == 0 {
566 A.rb_mut()
567 } else {
568 A.rb_mut().transpose_mut()
569 };
570 let n = n - 1;
571 let V = V.rb().submatrix(1, 0, n, n);
572 let mut A = A.rb_mut().subrows_mut(1, n);
573 let H = H.as_ref();
574 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_with_conj(
575 V,
576 H.as_ref(),
577 if iter == 0 { Conj::Yes } else { Conj::No },
578 A.rb_mut(),
579 Par::Seq,
580 MemStack::new(&mut MemBuffer::new(
581 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<f64>(n, b, n + 1),
582 )),
583 );
584 }
585 let approx_eq = CwiseMat(ApproxEq::<f64>::eps());
586 for j in 0..n {
587 for i in 0..n {
588 if i > j + 1 || j > i + 1 {
589 V[(i, j)] = 0.0;
590 }
591 }
592 }
593 for i in 0..n {
594 if i + 1 < n {
595 V[(i, i + 1)] = V[(i + 1, i)];
596 }
597 }
598 assert!(V ~ A);
599 }
600 }
601 #[test]
602 fn test_tridiag_cplx() {
603 let rng = &mut StdRng::seed_from_u64(0);
604 for n in [2, 3, 4, 8, 16] {
605 let A = CwiseMatDistribution {
606 nrows: n,
607 ncols: n,
608 dist: ComplexDistribution::new(StandardNormal, StandardNormal),
609 }
610 .rand::<Mat<c64>>(rng);
611 let A = A.rb() + A.adjoint();
612 let b = 3;
613 let mut H = Mat::zeros(b, n - 1);
614 let mut V = A.clone();
615 let mut V = V.as_mut();
616 tridiag_in_place(
617 V.rb_mut(),
618 H.as_mut(),
619 Par::Seq,
620 MemStack::new(&mut MemBuffer::new(tridiag_in_place_scratch::<
621 c64,
622 >(
623 n, Par::Seq, default()
624 ))),
625 default(),
626 );
627 let mut A = A.clone();
628 let mut A = A.as_mut();
629 for iter in 0..2 {
630 let mut A = if iter == 0 {
631 A.rb_mut()
632 } else {
633 A.rb_mut().transpose_mut()
634 };
635 let n = n - 1;
636 let V = V.rb().submatrix(1, 0, n, n);
637 let mut A = A.rb_mut().subrows_mut(1, n);
638 let H = H.as_ref();
639 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_with_conj(
640 V,
641 H.as_ref(),
642 if iter == 0 { Conj::Yes } else { Conj::No },
643 A.rb_mut(),
644 Par::Seq,
645 MemStack::new(&mut MemBuffer::new(
646 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<c64>(n, b, n + 1),
647 )),
648 );
649 }
650 let approx_eq = CwiseMat(ApproxEq::eps());
651 for j in 0..n {
652 for i in 0..n {
653 if i > j + 1 || j > i + 1 {
654 V[(i, j)] = c64::ZERO;
655 }
656 }
657 }
658 for i in 0..n {
659 if i + 1 < n {
660 V[(i, i + 1)] = V[(i + 1, i)].conj();
661 }
662 }
663 assert!(V ~ A);
664 }
665 }
666}