1use crate::assert;
2use crate::internal_prelude::*;
3use crate::utils::thread::with_width;
4use crate::utils::cost::{ParThreshold, rate};
5use linalg::householder::{self, HouseholderInfo};
6use linalg::matmul::triangular::BlockStructure;
7use linalg::matmul::{self, dot, matmul};
8use linalg::triangular_solve;
9#[derive(Copy, Clone, Debug)]
11pub struct HessenbergParams {
12 pub par_threshold: ParThreshold,
16 pub blocking_threshold: usize,
18 #[doc(hidden)]
19 pub non_exhaustive: NonExhaustive,
20}
21impl<T: ComplexField> Auto<T> for HessenbergParams {
22 fn auto() -> Self {
23 Self {
24 par_threshold: ParThreshold::Derived,
25 blocking_threshold: 256 * 256,
26 non_exhaustive: NonExhaustive(()),
27 }
28 }
29}
30pub fn hessenberg_in_place_scratch<T: ComplexField>(
33 dim: usize,
34 block_size: usize,
35 par: ParShape,
36 params: Spec<HessenbergParams, T>,
37) -> StackReq {
38 let params = params.config;
39 let _ = par;
40 let n = dim;
41 if n * n < params.blocking_threshold {
42 StackReq::any_of(&[StackReq::all_of(&[
43 temp_mat_scratch::<T>(n, 1).array(3),
44 temp_mat_scratch::<T>(n, par.degree()),
45 ])])
46 } else {
47 StackReq::all_of(&[
48 temp_mat_scratch::<T>(n, block_size),
49 temp_mat_scratch::<T>(block_size, 1),
50 StackReq::any_of(&[
51 StackReq::all_of(&[
52 temp_mat_scratch::<T>(n, 1),
53 temp_mat_scratch::<T>(n, par.degree()),
54 ]),
55 temp_mat_scratch::<T>(n, block_size),
56 ]),
57 ])
58 }
59}
60fn hessenberg_fused_op_simd<T: ComplexField>(
61 A: MatMut<'_, T, usize, usize, ContiguousFwd>,
62 l_out: RowMut<'_, T, usize>,
63 r_out: ColMut<'_, T, usize, ContiguousFwd>,
64 l_in: RowRef<'_, T, usize, ContiguousFwd>,
65 r_in: ColRef<'_, T, usize>,
66 l0: ColRef<'_, T, usize, ContiguousFwd>,
67 l1: ColRef<'_, T, usize, ContiguousFwd>,
68 r0: RowRef<'_, T, usize>,
69 r1: RowRef<'_, T, usize>,
70 align: usize,
71) {
72 struct Impl<'a, 'M, 'N, T: ComplexField> {
73 A: MatMut<'a, T, Dim<'M>, Dim<'N>, ContiguousFwd>,
74 l_out: RowMut<'a, T, Dim<'N>>,
75 r_out: ColMut<'a, T, Dim<'M>, ContiguousFwd>,
76 l_in: RowRef<'a, T, Dim<'M>, ContiguousFwd>,
77 r_in: ColRef<'a, T, Dim<'N>>,
78 l0: ColRef<'a, T, Dim<'M>, ContiguousFwd>,
79 l1: ColRef<'a, T, Dim<'M>, ContiguousFwd>,
80 r0: RowRef<'a, T, Dim<'N>>,
81 r1: RowRef<'a, T, Dim<'N>>,
82 align: usize,
83 }
84 impl<'a, 'M, 'N, T: ComplexField> pulp::WithSimd for Impl<'a, 'M, 'N, T> {
85 type Output = ();
86
87 #[inline(always)]
88 fn with_simd<S: pulp::Simd>(self, simd: S) -> Self::Output {
89 let Self {
90 mut A,
91 mut l_out,
92 mut r_out,
93 l_in,
94 r_in,
95 l0,
96 l1,
97 r0,
98 r1,
99 align,
100 } = self;
101 let (m, n) = A.shape();
102 let simd = SimdCtx::<T, S>::new_align(T::simd_ctx(simd), m, align);
103
104 simd_iter!(for i in [simd.indices()] {
105 simd.write(r_out.rb_mut(), i, simd.zero());
106 });
107
108 let indices = simd.batch_indices::<4>();
109 let l_in = l_in.transpose();
110 for j in n.indices() {
111 let mut A = A.rb_mut().col_mut(j);
112 let r_in = simd.splat(r_in.at(j));
113 let r0 = simd.splat(-&r0[j]);
114 let r1 = simd.splat(-&r1[j]);
115 let mut acc = [simd.zero(); 4];
116 simd_iter!(for (IDX, i) in [indices; 4] {
117 let mut a = simd.read(A.rb(), i);
118 a = simd.mul_add(simd.read(l0, i), r0, a);
119 a = simd.conj_mul_add(r1, simd.read(l1, i), a);
120 simd.write(A.rb_mut(), i, a);
121 acc[IDX] =
122 simd.conj_mul_add(simd.read(l_in, i), a, acc[IDX]);
123 let tmp = simd.read(r_out.rb(), i);
124 simd.write(r_out.rb_mut(), i, simd.mul_add(a, r_in, tmp));
125 });
126 let acc0 = simd.add(acc[0], acc[1]);
127 let acc2 = simd.add(acc[2], acc[3]);
128 let acc0 = simd.add(acc0, acc2);
129 let l_out = l_out.rb_mut().at_mut(j);
130 *l_out = simd.reduce_sum(acc0);
131 }
132 }
133 }
134 with_dim!(M, A.nrows());
135 with_dim!(N, A.ncols());
136 dispatch!(
137 Impl {
138 A: A.as_shape_mut(M, N),
139 l_out: l_out.as_col_shape_mut(N),
140 r_out: r_out.as_row_shape_mut(M),
141 l_in: l_in.as_col_shape(M),
142 r_in: r_in.as_row_shape(N),
143 l0: l0.as_row_shape(M),
144 l1: l1.as_row_shape(M),
145 r0: r0.as_col_shape(N),
146 r1: r1.as_col_shape(N),
147 align,
148 },
149 Impl,
150 T
151 )
152}
153fn hessenberg_fused_op_fallback<T: ComplexField>(
154 A: MatMut<'_, T>,
155 l_out: RowMut<'_, T>,
156 r_out: ColMut<'_, T>,
157 l_in: RowRef<'_, T>,
158 r_in: ColRef<'_, T>,
159 l0: ColRef<'_, T>,
160 l1: ColRef<'_, T>,
161 r0: RowRef<'_, T>,
162 r1: RowRef<'_, T>,
163) {
164 let mut A = A;
165 matmul(
166 A.rb_mut(),
167 Accum::Add,
168 l0.as_mat(),
169 r0.as_mat(),
170 -one::<T>(),
171 Par::Seq,
172 );
173 matmul(
174 A.rb_mut(),
175 Accum::Add,
176 l1.as_mat(),
177 r1.as_mat().conjugate(),
178 -one::<T>(),
179 Par::Seq,
180 );
181 matmul(
182 r_out.as_mat_mut(),
183 Accum::Replace,
184 A.rb(),
185 r_in.as_mat(),
186 one(),
187 Par::Seq,
188 );
189 matmul(
190 l_out.as_mat_mut(),
191 Accum::Replace,
192 l_in.as_mat().conjugate(),
193 A.rb(),
194 one(),
195 Par::Seq,
196 );
197}
198fn hessenberg_fused_op<T: ComplexField>(
199 A: MatMut<'_, T>,
200 l_out: RowMut<'_, T>,
201 r_out: ColMut<'_, T>,
202 l_in: RowRef<'_, T>,
203 r_in: ColRef<'_, T>,
204 l0: ColRef<'_, T>,
205 l1: ColRef<'_, T>,
206 r0: RowRef<'_, T>,
207 r1: RowRef<'_, T>,
208 align: usize,
209) {
210 let mut A = A;
211 let mut r_out = r_out;
212 if const { T::SIMD_CAPABILITIES.is_simd() } {
213 if let (Some(A), Some(r_out), Some(l_in), Some(l0), Some(l1)) = (
214 A.rb_mut().try_as_col_major_mut(),
215 r_out.rb_mut().try_as_col_major_mut(),
216 l_in.try_as_row_major(),
217 l0.try_as_col_major(),
218 l1.try_as_col_major(),
219 ) {
220 hessenberg_fused_op_simd(
221 A, l_out, r_out, l_in, r_in, l0, l1, r0, r1, align,
222 );
223 } else {
224 hessenberg_fused_op_fallback(
225 A, l_out, r_out, l_in, r_in, l0, l1, r0, r1,
226 );
227 }
228 } else {
229 hessenberg_fused_op_fallback(
230 A, l_out, r_out, l_in, r_in, l0, l1, r0, r1,
231 );
232 }
233}
234fn hessenberg_rearranged_unblocked<T: ComplexField>(
235 A: MatMut<'_, T>,
236 H: MatMut<'_, T>,
237 par: Par,
238 stack: &mut MemStack,
239 params: HessenbergParams,
240) {
241 assert!(all(
242 A.nrows() == A.ncols(),
243 H.ncols() == A.ncols().saturating_sub(1)
244 ));
245 let n = A.nrows();
246 let b = H.nrows();
247 if n == 0 {
248 return;
249 }
250 let mut A = A;
251 let mut H = H;
252 let mut par = par;
253 if matches!(params.par_threshold, ParThreshold::Never) {
255 par = Par::Seq;
256 }
257 {
258 let mut H = H.rb_mut().row_mut(0);
259 alloca!('stack: {
260 let mut y = unsafe { row![uninit::<T>, n] };
261 let mut z = unsafe { col![uninit::<T>, n] };
262 let mut v = unsafe { row![uninit::<T>, n] };
263 let mut w = unsafe { mat![uninit::<T>, n, par.degree()] };
264 });
265 for k in 0..n {
266 let (_, A01, A10, A11) = A.rb_mut().split_at_mut(k, k);
267 let (_, mut A02) = A01.split_first_col_mut().unwrap();
268 let (_, A20) = A10.split_first_row_mut().unwrap();
269 let (mut A11, A12, A21, mut A22) = A11.split_at_mut(1, 1);
270 let mut A12 = A12.row_mut(0);
271 let mut A21 = A21.col_mut(0);
272 let A11 = &mut A11[(0, 0)];
273 let (y1, mut y2) =
274 y.rb_mut().split_at_col_mut(k).1.split_at_col_mut(1);
275 let (z1, mut z2) =
276 z.rb_mut().split_at_row_mut(k).1.split_at_row_mut(1);
277 let (_, mut v2) =
278 v.rb_mut().split_at_col_mut(k).1.split_at_col_mut(1);
279 let (mut w0, w12) = w.rb_mut().split_at_row_mut(k);
280 let (_, mut w2) = w12.split_at_row_mut(1);
281 if k > 0 {
282 let y1 = y1[0].copy();
283 let z1 = z1[0].copy();
284
285 let p = k - 1;
286 let u2 = A20.rb().col(p);
287 *A11 -= &y1 + &z1;
288 z!(&mut A12, &y2, u2.rb().transpose()).for_each(
289 |uz!(a, y, u): Zip!(&mut T, &T, &T)| {
290 *a -= y + &z1 * u.conj();
291 },
292 );
293 z!(&mut A21, &u2, &z2).for_each(
294 |uz!(a, u, z): Zip!(&mut T, &T, &T)| *a -= u * &y1 + z,
295 );
296 }
297 let width = {
302 let n = n - k - 1;
303 params.par_threshold.region_width::<T>(&par, n * n, rate::hessenberg(), 2)
304 };
305 if k + 1 == n {
306 break;
307 }
308 let beta;
309 let tau_inv;
310 {
311 let (mut A11, mut A21) = A21.rb_mut().split_at_row_mut(1);
312 let A11 = &mut A11[0];
313 let HouseholderInfo { tau, .. } =
314 householder::make_householder_in_place(A11, A21.rb_mut());
315 tau_inv = tau.recip();
316 beta = A11.copy();
317 *A11 = one();
318 H[k] = tau.to_cplx();
319 }
320 let x2 = A21.rb();
321 if k > 0 {
322 let p = k - 1;
323 let u2 = A20.rb().col(p);
324 hessenberg_fused_op(
325 A22.rb_mut(),
326 v2.rb_mut(),
327 w2.rb_mut().col_mut(0),
328 x2.transpose(),
329 x2,
330 u2,
331 z2.rb(),
332 y2.rb(),
333 u2.transpose(),
334 simd_align(k + 1),
335 );
336 y2.copy_from(v2.rb());
337 z2.copy_from(w2.rb().col(0));
338 } else {
339 with_width(par.reborrow(), width, |par| {
340 matmul(
341 z2.rb_mut().as_mat_mut(),
342 Accum::Replace,
343 A22.rb(),
344 x2.as_mat(),
345 one(),
346 par,
347 );
348 });
349 with_width(par.reborrow(), width, |par| {
350 matmul(
351 y2.rb_mut().as_mat_mut(),
352 Accum::Replace,
353 x2.adjoint().as_mat(),
354 A22.rb(),
355 one(),
356 par,
357 );
358 });
359 }
360 let u2 = x2;
361 let ref b = dot::inner_prod(
362 u2.rb().transpose(),
363 Conj::Yes,
364 z2.rb(),
365 Conj::No,
366 )
367 .mul_pow2(from_f64::<T::Real>(0.5))
368 .mul_real(&tau_inv);
369 z!(&mut y2, u2.transpose()).for_each(
370 |uz!(y, u): Zip!(&mut T, &T)| {
371 *y = (&*y - b * u.conj()).mul_real(&tau_inv);
372 },
373 );
374 z!(&mut z2, u2).for_each(|uz!(z, u): Zip!(&mut T, &T)| {
375 *z = (&*z - b * u).mul_real(&tau_inv);
376 });
377 let ref dot =
378 dot::inner_prod(A12.rb(), Conj::No, u2.rb(), Conj::No)
379 .mul_real(&tau_inv);
380 z!(&mut A12, u2.transpose())
381 .for_each(|uz!(a, u): Zip!(&mut T, &T)| *a -= dot * u.conj());
382 with_width(par.reborrow(), width, |par| {
383 matmul(
384 w0.rb_mut().col_mut(0).as_mat_mut(),
385 Accum::Replace,
386 A02.rb(),
387 u2.as_mat(),
388 one(),
389 par,
390 );
391 });
392 with_width(par.reborrow(), width, |par| {
393 matmul(
394 A02.rb_mut(),
395 Accum::Add,
396 w0.rb().col(0).as_mat(),
397 u2.adjoint().as_mat(),
398 (-&tau_inv).to_cplx(),
399 par,
400 );
401 });
402 A21[0] = beta;
403 }
404 }
405 if n > 0 {
406 let n = n - 1;
407 let A = A.rb().submatrix(1, 0, n, n);
408 let mut H = H.rb_mut().subcols_mut(0, n);
409 let mut j = 0;
410 while j < n {
411 let b = Ord::min(b, n - j);
412 let mut H = H.rb_mut().submatrix_mut(0, j, b, b);
413 for k in 0..b {
414 H[(k, k)] = H[(0, k)].copy();
415 }
416 householder::upgrade_householder_factor(
417 H.rb_mut(),
418 A.submatrix(j, j, n - j, b),
419 b,
420 1,
421 par.reborrow(),
422 );
423 j += b;
424 }
425 }
426}
427fn hessenberg_gqvdg_unblocked<T: ComplexField>(
428 A: MatMut<'_, T>,
429 Z: MatMut<'_, T>,
430 H: MatMut<'_, T>,
431 beta: ColMut<'_, T>,
432 mut par: Par,
433 stack: &mut MemStack,
434 params: HessenbergParams,
435) {
436 let n = A.nrows();
437 let b = H.nrows();
438 let mut A = A;
439 let mut H = H;
440 let mut Z = Z;
441 _ = params;
442 let (mut x, _) = unsafe { temp_mat_uninit(n, 1, stack) };
443 let mut x = x.as_mat_mut().col_mut(0);
444 let mut beta = beta;
445 for k in 0..b {
446 let mut x0 = x.rb_mut().subrows_mut(0, k);
447 let (T00, T01, _, T11) = H.rb_mut().split_at_mut(k, k);
448 let (mut T01, _) = T01.split_first_col_mut().unwrap();
449 let (mut T11, _, _, _) = T11.split_at_mut(1, 1);
450 let T11 = &mut T11[(0, 0)];
451 let (U0, A12) = A.rb_mut().split_at_col_mut(k);
452 let (mut A1, A2) = A12.split_first_col_mut().unwrap();
453 let (Z0, Z12) = Z.rb_mut().split_at_col_mut(k);
454 let (mut Z1, _) = Z12.split_first_col_mut().unwrap();
455 let U0 = U0.rb();
456 let Z0 = Z0.rb();
457 let T00 = T00.rb();
458 let (U00, U10) = U0.split_at_row(k);
459 let (U10, U20) = U10.split_first_row().unwrap();
460 x0.copy_from(U10.adjoint());
461 triangular_solve::solve_upper_triangular_in_place(
462 T00,
463 x0.rb_mut().as_mat_mut(),
464 par.reborrow(),
465 );
466 matmul::matmul(
467 A1.rb_mut().as_mat_mut(),
468 Accum::Add,
469 Z0,
470 x0.rb().as_mat(),
471 -one::<T>(),
472 par.reborrow(),
473 );
474 let (mut A01, A11) = A1.rb_mut().split_at_row_mut(k);
475 let (mut A11, mut A21) = A11.split_at_row_mut(1);
476 let A11 = &mut A11[0];
477 {
478 matmul::triangular::matmul(
479 x0.rb_mut().as_mat_mut(),
480 BlockStructure::Rectangular,
481 Accum::Replace,
482 U00.adjoint(),
483 BlockStructure::StrictTriangularUpper,
484 A01.rb().as_mat(),
485 BlockStructure::Rectangular,
486 one(),
487 par.reborrow(),
488 );
489 z!(x0.rb_mut(), U10.transpose())
490 .for_each(|uz!(x, u): Zip!(&mut T, &T)| *x += &*A11 * u.conj());
491 matmul::matmul(
492 x0.rb_mut().as_mat_mut(),
493 Accum::Add,
494 U20.adjoint(),
495 A21.rb().as_mat(),
496 one(),
497 par.reborrow(),
498 );
499 }
500 {
501 triangular_solve::solve_lower_triangular_in_place(
502 T00.adjoint(),
503 x0.rb_mut().as_mat_mut(),
504 par.reborrow(),
505 );
506 }
507 {
508 matmul::triangular::matmul(
509 A01.rb_mut().as_mat_mut(),
510 BlockStructure::Rectangular,
511 Accum::Add,
512 U00,
513 BlockStructure::StrictTriangularLower,
514 x0.rb().as_mat(),
515 BlockStructure::Rectangular,
516 -one::<T>(),
517 par.reborrow(),
518 );
519 *A11 -= dot::inner_prod(U10, Conj::No, x0.rb(), Conj::No);
520 matmul::matmul(
521 A21.rb_mut().as_mat_mut(),
522 Accum::Add,
523 U20,
524 x0.rb().as_mat(),
525 -one::<T>(),
526 par.reborrow(),
527 );
528 }
529 if k + 1 < n {
530 let (mut A11, mut A21) = A21.rb_mut().split_at_row_mut(1);
531 let A11 = &mut A11[0];
532 let HouseholderInfo { tau, .. } =
533 householder::make_householder_in_place(A11, A21.rb_mut());
534 beta[k] = A11.copy();
535 *A11 = one();
536 *T11 = tau.to_cplx();
537 } else {
538 *T11 = infinity();
539 }
540 matmul::matmul(
541 Z1.rb_mut().as_mat_mut(),
542 Accum::Replace,
543 A2.rb(),
544 A21.rb().as_mat(),
545 one(),
546 par.reborrow(),
547 );
548 matmul::matmul(
549 T01.rb_mut().as_mat_mut(),
550 Accum::Replace,
551 U20.adjoint(),
552 A21.rb().as_mat(),
553 one(),
554 par.reborrow(),
555 );
556 }
557}
558#[track_caller]
567pub fn hessenberg_in_place<T: ComplexField>(
568 A: MatMut<'_, T>,
569 householder: MatMut<'_, T>,
570 par: Par,
571 stack: &mut MemStack,
572 params: Spec<HessenbergParams, T>,
573) {
574 let params = params.config;
575 assert!(all(
576 A.nrows() == A.ncols(),
577 householder.ncols() == A.ncols().saturating_sub(1)
578 ));
579 let n = A.nrows().unbound();
580 if n * n < params.blocking_threshold {
581 hessenberg_rearranged_unblocked(A, householder, par, stack, params);
582 } else {
583 hessenberg_gqvdg_blocked(A, householder, par, stack, params);
584 }
585}
586fn hessenberg_gqvdg_blocked<T: ComplexField>(
587 A: MatMut<'_, T>,
588 H: MatMut<'_, T>,
589 mut par: Par,
590 stack: &mut MemStack,
591 params: HessenbergParams,
592) {
593 let n = A.nrows();
594 let b = H.nrows();
595 let mut A = A;
596 let mut H = H;
597 let (mut Z, stack) = unsafe { temp_mat_uninit(n, b, stack) };
598 let mut Z = Z.as_mat_mut();
599 let mut j = 0;
600 while j < n {
601 let bs = Ord::min(b, n - j);
602 let bs_u = Ord::min(bs, n - j - 1);
603 let (mut beta, stack) = unsafe { temp_mat_uninit(bs, 1, stack) };
604 let mut beta = beta.as_mat_mut().col_mut(0);
605 {
606 let mut T11 = H.rb_mut().submatrix_mut(0, j, bs_u, bs_u);
607 {
608 let A11 = A.rb_mut().submatrix_mut(j, j, n - j, n - j);
609 let Z1 = Z.rb_mut().submatrix_mut(j, 0, n - j, bs);
610 hessenberg_gqvdg_unblocked(
611 A11,
612 Z1,
613 T11.rb_mut(),
614 beta.rb_mut(),
615 par.reborrow(),
616 stack,
617 params,
618 );
619 }
620 let (mut X, _) = unsafe { temp_mat_uninit(n, bs_u, stack) };
621 let mut X = X.as_mat_mut();
622 let (mut X0, X12) = X.rb_mut().split_at_row_mut(j);
623 let (_, mut X2) = X12.split_at_row_mut(bs_u);
624 let (_, Z12) = Z.rb_mut().subcols_mut(0, bs_u).split_at_row_mut(j);
625 let (mut Z1, mut Z2) = Z12.split_at_row_mut(bs_u);
626 let (_, A01, _, A11) = A.rb_mut().split_at_mut(j, j);
627 let (mut A01, mut A02) = A01.split_at_col_mut(bs_u);
628 let (A11, mut A12, A21, mut A22) = A11.split_at_mut(bs_u, bs_u);
629 let U1 = A11.rb();
630 let U2 = A21.rb();
631 let T1 = T11.rb();
632 matmul::triangular::matmul(
633 X0.rb_mut(),
634 BlockStructure::Rectangular,
635 Accum::Replace,
636 A01.rb(),
637 BlockStructure::Rectangular,
638 U1,
639 BlockStructure::StrictTriangularLower,
640 one(),
641 par.reborrow(),
642 );
643 matmul::matmul(X0.rb_mut(), Accum::Add, A02.rb(), U2, one(), par.reborrow());
644 triangular_solve::solve_lower_triangular_in_place(
645 T1.transpose(),
646 X0.rb_mut().transpose_mut(),
647 par.reborrow(),
648 );
649 matmul::triangular::matmul(
650 A01.rb_mut(),
651 BlockStructure::Rectangular,
652 Accum::Add,
653 X0.rb(),
654 BlockStructure::Rectangular,
655 U1.adjoint(),
656 BlockStructure::StrictTriangularUpper,
657 -one::<T>(),
658 par.reborrow(),
659 );
660 matmul::matmul(
661 A02.rb_mut(),
662 Accum::Add,
663 X0.rb(),
664 U2.adjoint(),
665 -one::<T>(),
666 par.reborrow(),
667 );
668 triangular_solve::solve_lower_triangular_in_place(
669 T1.transpose(),
670 Z1.rb_mut().transpose_mut(),
671 par.reborrow(),
672 );
673 triangular_solve::solve_lower_triangular_in_place(
674 T1.transpose(),
675 Z2.rb_mut().transpose_mut(),
676 par.reborrow(),
677 );
678 matmul::matmul(
679 A12.rb_mut(),
680 Accum::Add,
681 Z1.rb(),
682 U2.adjoint(),
683 -one::<T>(),
684 par.reborrow(),
685 );
686 matmul::matmul(
687 A22.rb_mut(),
688 Accum::Add,
689 Z2.rb(),
690 U2.adjoint(),
691 -one::<T>(),
692 par.reborrow(),
693 );
694 let mut X = X2.rb_mut().transpose_mut();
695 matmul::triangular::matmul(
696 X.rb_mut(),
697 BlockStructure::Rectangular,
698 Accum::Replace,
699 U1.adjoint(),
700 BlockStructure::StrictTriangularUpper,
701 A12.rb(),
702 BlockStructure::Rectangular,
703 one(),
704 par.reborrow(),
705 );
706 matmul::matmul(
707 X.rb_mut(),
708 Accum::Add,
709 U2.adjoint(),
710 A22.rb(),
711 one(),
712 par.reborrow(),
713 );
714 triangular_solve::solve_lower_triangular_in_place(
715 T1.adjoint(),
716 X.rb_mut(),
717 par.reborrow(),
718 );
719 matmul::triangular::matmul(
720 A12.rb_mut(),
721 BlockStructure::Rectangular,
722 Accum::Add,
723 U1,
724 BlockStructure::StrictTriangularLower,
725 X.rb(),
726 BlockStructure::Rectangular,
727 -one::<T>(),
728 par.reborrow(),
729 );
730 matmul::matmul(
731 A22.rb_mut(),
732 Accum::Add,
733 U2,
734 X.rb(),
735 -one::<T>(),
736 par.reborrow(),
737 );
738 }
739 let n = n - j;
740 let mut A = A.rb_mut().submatrix_mut(j, j, n, bs);
741 for k in 0..bs {
742 if k + 1 < n {
743 A[(k + 1, k)] = beta[k].copy();
744 }
745 }
746 j += bs;
747 }
748}
749#[cfg(test)]
750mod tests {
751 use super::*;
752 use crate::stats::prelude::*;
753 use crate::utils::approx::*;
754 use crate::{Mat, assert, c64};
755 use dyn_stack::MemBuffer;
756 use std::mem::MaybeUninit;
757 #[test]
758 fn test_hessenberg_real() {
759 let rng = &mut StdRng::seed_from_u64(0);
760 for n in [3, 4, 8, 16] {
761 let A = CwiseMatDistribution {
762 nrows: n,
763 ncols: n,
764 dist: StandardNormal,
765 }
766 .rand::<Mat<f64>>(rng);
767 let b = 3;
768 let mut H = Mat::zeros(b, n - 1);
769 let mut V = A.clone();
770 let mut V = V.as_mut();
771 hessenberg_rearranged_unblocked(
772 V.rb_mut(),
773 H.as_mut(),
774 Par::Seq,
775 MemStack::new(&mut [MaybeUninit::uninit(); 1024]),
776 auto!(f64),
777 );
778 let mut A = A.clone();
779 let mut A = A.as_mut();
780 for iter in 0..2 {
781 let mut A = if iter == 0 {
782 A.rb_mut()
783 } else {
784 A.rb_mut().transpose_mut()
785 };
786 let n = n - 1;
787 let V = V.rb().submatrix(1, 0, n, n);
788 let mut A = A.rb_mut().subrows_mut(1, n);
789 let H = H.as_ref();
790 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_with_conj(
791 V,
792 H.as_ref(),
793 if iter == 0 { Conj::Yes } else { Conj::No },
794 A.rb_mut(),
795 Par::Seq,
796 MemStack::new(&mut MemBuffer::new(
797 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<f64>(n, b, n + 1),
798 )),
799 );
800 }
801 let approx_eq = CwiseMat(ApproxEq::<f64>::eps());
802 for j in 0..n {
803 for i in 0..n {
804 if i > j + 1 {
805 V[(i, j)] = 0.0;
806 }
807 }
808 }
809 assert!(V ~ A);
810 }
811 }
812 #[test]
813 fn test_hessenberg_cplx() {
814 let rng = &mut StdRng::seed_from_u64(0);
815 for n in [1, 2, 3, 4, 8, 16] {
816 let mut run = |par: Par<'_>| {
817 let A = CwiseMatDistribution {
818 nrows: n,
819 ncols: n,
820 dist: ComplexDistribution::new(
821 StandardNormal,
822 StandardNormal,
823 ),
824 }
825 .rand::<Mat<c64>>(rng);
826 let b = 3;
827 let mut H = Mat::zeros(b, n - 1);
828 let mut V = A.clone();
829 let mut V = V.as_mut();
830 hessenberg_rearranged_unblocked(
831 V.rb_mut(),
832 H.as_mut(),
833 par,
834 MemStack::new(&mut [MaybeUninit::uninit(); 8 * 1024]),
835 HessenbergParams {
836 par_threshold: ParThreshold::Fixed(1),
837 ..auto!(c64)
838 }
839 .into(),
840 );
841 let mut A = A.clone();
842 let mut A = A.as_mut();
843 for iter in 0..2 {
844 let mut A = if iter == 0 {
845 A.rb_mut()
846 } else {
847 A.rb_mut().transpose_mut()
848 };
849 let n = n - 1;
850 let V = V.rb().submatrix(1, 0, n, n);
851 let mut A = A.rb_mut().subrows_mut(1, n);
852 let H = H.as_ref();
853 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_with_conj(
854 V,
855 H.as_ref(),
856 if iter == 0 { Conj::Yes } else { Conj::No },
857 A.rb_mut(),
858 Par::Seq,
859 MemStack::new(&mut MemBuffer::new(
860 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<c64>(n, b, n + 1),
861 )),
862 );
863 }
864 let approx_eq = CwiseMat(ApproxEq::eps());
865 for j in 0..n {
866 for i in 0..n {
867 if i > j + 1 {
868 V[(i, j)] = c64::ZERO;
869 }
870 }
871 }
872 assert!(V ~ A);
873 };
874 run(Par::Seq);
875 crate::team::with_team(4, |par| run(par));
876 }
877 }
878 #[test]
884 fn test_hessenberg_cplx_team_reconstructs_blocked() {
885 fn check<T: crate::ComplexField<Real = f64>>(a: Mat<T>) {
886 let n = a.nrows();
887 let b = 8;
888 let mut h = Mat::<T>::zeros(b, n - 1);
889 let mut v = a.clone();
890 let params: HessenbergParams = HessenbergParams {
891 par_threshold: ParThreshold::Fixed(1),
892 blocking_threshold: usize::MAX,
893 ..auto!(T)
894 };
895 let mut mem = MemBuffer::new(hessenberg_in_place_scratch::<T>(
896 n,
897 b,
898 ParShape::team(4),
899 params.into(),
900 ));
901 crate::team::with_team(4, |par| {
902 let mut v_v = v.as_mut();
903 hessenberg_in_place(
904 v_v.rb_mut(),
905 h.as_mut(),
906 par,
907 MemStack::new(&mut mem),
908 params.into(),
909 );
910 });
911 let mut rebuilt = a.clone();
912 let mut rebuilt = rebuilt.as_mut();
913 for iter in 0..2 {
914 let mut dst = if iter == 0 {
915 rebuilt.rb_mut()
916 } else {
917 rebuilt.rb_mut().transpose_mut()
918 };
919 let m = n - 1;
920 let vv = v.rb().submatrix(1, 0, m, m);
921 let mut dst = dst.rb_mut().subrows_mut(1, m);
922 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_with_conj(
923 vv,
924 h.as_ref(),
925 if iter == 0 { Conj::Yes } else { Conj::No },
926 dst.rb_mut(),
927 Par::Seq,
928 MemStack::new(&mut MemBuffer::new(
929 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<T>(m, b, n + 1),
930 )),
931 );
932 }
933 for j in 0..n {
934 for i in 0..n {
935 if i > j + 1 {
936 v[(i, j)] = T::zero();
937 }
938 }
939 }
940 let mut err = 0.0f64;
941 let mut nonfinite = false;
942 let mut scale = 0.0f64;
943 for j in 0..n {
944 for i in 0..n {
945 let d = (v[(i, j)].clone() - rebuilt[(i, j)].clone()).abs();
946 nonfinite |= !d.is_finite() || !v[(i, j)].abs().is_finite();
947 if d.is_finite() {
948 err = err.max(d);
949 }
950 scale = scale.max(a[(i, j)].abs());
951 }
952 }
953 core::assert!(!nonfinite, "non-finite factor or difference at {n}x{n}");
954 core::assert!(err <= 1e-10 * scale.max(1.0), "{n}x{n}: {err} {scale}");
955 }
956 let rng = &mut StdRng::seed_from_u64(21);
957 for n in [16usize, 96] {
958 check::<c64>(Mat::from_fn(n, n, |_, _| {
959 c64::new(
960 rand::Rng::random::<f64>(rng),
961 rand::Rng::random::<f64>(rng),
962 )
963 }));
964 }
965 check::<c64>(Mat::from_fn(96, 96, |i, j| {
967 let k = (i * 2654435761usize).wrapping_add(j * 40503).wrapping_add(977) % 1009;
968 c64::new(k as f64 / 504.0 - 1.0, 0.0)
969 }));
970 }
971 #[test]
972 fn test_hessenberg_cplx_gqvdg() {
973 let rng = &mut StdRng::seed_from_u64(0);
974 for n in [2, 3, 4, 8, 16, 21] {
975 let mut run = |par: Par<'_>| {
976 let b = 4;
977 let A = CwiseMatDistribution {
978 nrows: n,
979 ncols: n,
980 dist: ComplexDistribution::new(
981 StandardNormal,
982 StandardNormal,
983 ),
984 }
985 .rand::<Mat<c64, _, _>>(rng);
986 let mut H = Mat::zeros(b, n - 1);
987 let mut V = A.clone();
988 let mut V = V.as_mut();
989 hessenberg_gqvdg_blocked(
990 V.rb_mut(),
991 H.as_mut(),
992 par,
993 MemStack::new(&mut [MaybeUninit::uninit(); 16 * 1024]),
994 HessenbergParams {
995 par_threshold: ParThreshold::Fixed(1),
996 ..auto!(c64)
997 }
998 .into(),
999 );
1000 let mut A = A.clone();
1001 let mut A = A.as_mut();
1002 for iter in 0..2 {
1003 let mut A = if iter == 0 {
1004 A.rb_mut()
1005 } else {
1006 A.rb_mut().transpose_mut()
1007 };
1008 let n = n - 1;
1009 let V = V.rb().submatrix(1, 0, n, n);
1010 let mut A = A.rb_mut().subrows_mut(1, n);
1011 let H = H.as_ref().subcols(0, n);
1012 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_with_conj(
1013 V,
1014 H.as_ref(),
1015 if iter == 0 { Conj::Yes } else { Conj::No },
1016 A.rb_mut(),
1017 Par::Seq,
1018 MemStack::new(&mut MemBuffer::new(
1019 householder::apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<c64>(n, b, n + 1),
1020 )),
1021 );
1022 }
1023 let approx_eq = CwiseMat(ApproxEq::eps());
1024 for j in 0..n {
1025 for i in 0..n {
1026 if i > j + 1 {
1027 V[(i, j)] = c64::ZERO;
1028 }
1029 }
1030 }
1031 assert!(V ~ A);
1032 };
1033 run(Par::Seq);
1034 crate::team::with_team(4, |par| run(par));
1035 }
1036 }
1037}