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