faer/linalg/svd/
bidiag.rs

1use crate::assert;
2use crate::internal_prelude::*;
3use linalg::householder::*;
4use linalg::matmul::{dot, matmul};
5/// computes the layout of the workspace required to compute a matrix's
6/// bidiagonalization
7pub fn bidiag_in_place_scratch<T: ComplexField>(
8	nrows: usize,
9	ncols: usize,
10	par: Par,
11	params: Spec<BidiagParams, T>,
12) -> StackReq {
13	_ = par;
14	_ = params;
15	StackReq::all_of(&[
16		temp_mat_scratch::<T>(nrows, 1),
17		temp_mat_scratch::<T>(ncols, 1),
18	])
19}
20/// bidiagonalization tuning parameters.
21#[derive(Debug, Copy, Clone)]
22pub struct BidiagParams {
23	/// threshold at which parallelism should be disabled
24	pub par_threshold: usize,
25	#[doc(hidden)]
26	pub non_exhaustive: NonExhaustive,
27}
28impl<T: ComplexField> Auto<T> for BidiagParams {
29	fn auto() -> Self {
30		Self {
31			par_threshold: 192 * 256,
32			non_exhaustive: NonExhaustive(()),
33		}
34	}
35}
36/// computes a matrix $A$'s bidiagonalization such that $A = U B V^H$
37///
38/// $B$ is a bidiagonal matrix stored in $A$'s diagonal and superdiagonal
39///
40/// $U$ is a sequence of householder reflections stored in the unit lower
41/// triangular half of $A$, with the householder coefficients being stored in
42/// `H_left`
43///
44/// $V$ is a sequence of householder reflections stored in the unit upper
45/// triangular half of $A$ (excluding the diagonal), with the householder
46/// coefficients being stored in `H_right`
47pub fn bidiag_in_place<T: ComplexField>(
48	A: MatMut<'_, T>,
49	H_left: MatMut<'_, T>,
50	H_right: MatMut<'_, T>,
51	par: Par,
52	stack: &mut MemStack,
53	params: Spec<BidiagParams, T>,
54) {
55	let params = params.config;
56	let m = A.nrows();
57	let n = A.ncols();
58	let size = Ord::min(m, n);
59	let bl = H_left.nrows();
60	let br = H_right.nrows();
61	assert!(H_left.ncols() == size);
62	assert!(H_right.ncols() == size.saturating_sub(1));
63	alloca!('stack: {
64		let mut y = unsafe { row![uninit::<T>, n] };
65		let mut z = unsafe { col![uninit::<T>, m] };
66	});
67	let mut A = A;
68	let mut Hl = H_left;
69	let mut Hr = H_right;
70	let mut par = par;
71	{
72		let mut Hl = Hl.rb_mut().row_mut(0);
73		let mut Hr = Hr.rb_mut().row_mut(0);
74		for k in 0..size {
75			let mut A = A.rb_mut();
76			let (_, A01, A10, A11) = A.rb_mut().split_at_mut(k, k);
77			let (_, A02) = A01.split_first_col().unwrap();
78			let (A10, A20) = A10.split_first_row_mut().unwrap();
79			let (mut A11, A12, A21, mut A22) = A11.split_at_mut(1, 1);
80			let mut A12 = A12.row_mut(0);
81			let mut A21 = A21.col_mut(0);
82			let a11 = &mut A11[(0, 0)];
83			let (y1, mut y2) =
84				y.rb_mut().split_at_col_mut(k).1.split_at_col_mut(1);
85			let (z1, mut z2) =
86				z.rb_mut().split_at_row_mut(k).1.split_at_row_mut(1);
87			if k > 0 {
88				let ref y1 = y1[0].copy();
89				let ref z1 = z1[0].copy();
90				let k1 = k - 1;
91				let ref up0 = A10[k1];
92				let up = A20.rb().col(k1);
93				let vp = A02.rb().row(k1);
94				*a11 -= up0 * y1 + z1;
95				z!(A21.rb_mut(), up.rb(), z2.rb()).for_each(
96					|uz!(a, u, z): Zip!(&mut T, &T, &T)| *a -= u * y1 + z,
97				);
98				z!(A12.rb_mut(), y2.rb(), vp.rb()).for_each(
99					|uz!(a, y, v): Zip!(&mut T, &T, &T)| {
100						*a -= up0 * y + z1 * v;
101					},
102				);
103			}
104			let HouseholderInfo { tau: tl, .. } =
105				make_householder_in_place(a11, A21.rb_mut());
106			let ref tl_inv = tl.recip();
107			Hl[k] = tl.to_cplx();
108			if (m - k - 1) * (n - k - 1) < params.par_threshold {
109				par = Par::Seq;
110			}
111			if k > 0 {
112				let k1 = k - 1;
113				let up = A20.rb().col(k1);
114				let vp = A02.row(k1);
115				match par {
116					Par::Seq => bidiag_fused_op(
117						A22.rb_mut(),
118						A21.rb(),
119						up.rb(),
120						z2.rb(),
121						y2.rb_mut(),
122						vp.rb(),
123						simd_align(k + 1),
124					),
125					#[cfg(feature = "rayon")]
126					Par::Rayon(nthreads) => {
127						use rayon::prelude::*;
128						let nthreads = nthreads.get();
129						spindle::for_each(
130							nthreads,
131							A22.rb_mut()
132								.par_col_partition_mut(nthreads)
133								.zip_eq(y2.rb_mut().par_partition_mut(nthreads))
134								.zip_eq(vp.par_partition(nthreads)),
135							|((A22, y2), vp)| {
136								bidiag_fused_op(
137									A22,
138									A21.rb(),
139									up.rb(),
140									z2.rb(),
141									y2,
142									vp.rb(),
143									simd_align(k + 1),
144								);
145							},
146						);
147					},
148				}
149			} else {
150				matmul(
151					y2.rb_mut(),
152					Accum::Replace,
153					A21.rb().adjoint(),
154					A22.rb(),
155					one(),
156					par,
157				);
158			}
159			z!(y2.rb_mut(), A12.rb_mut()).for_each(|uz!(y, a)| {
160				*y = (&*y + &*a).mul_real(tl_inv);
161				*a -= &*y;
162			});
163			let norm = A12.rb().norm_l2();
164			let ref norm_inv = norm.recip();
165			if norm != zero() {
166				z!(A12.rb_mut()).for_each(|uz!(a)| *a = a.mul_real(norm_inv));
167			}
168			matmul(
169				z2.rb_mut(),
170				Accum::Replace,
171				A22.rb(),
172				A12.rb().adjoint(),
173				one(),
174				par,
175			);
176			if k + 1 == size {
177				break;
178			}
179			let (mut A12_a, mut A12_b) = A12.rb_mut().split_at_col_mut(1);
180			let A22_a = A22.rb().col(0);
181			let (y2_a, y2_b) = y2.rb().split_at_col(1);
182			let y2_a = &y2_a[0];
183			let a12_a = &mut A12_a[0];
184			let HouseholderInfo {
185				tau: tr,
186				head_with_beta_inv: m,
187				..
188			} = make_householder_in_place(a12_a, A12_b.rb_mut().transpose_mut());
189			let ref tr_inv = tr.recip();
190			Hr[k] = tr.to_cplx();
191			let beta = a12_a.copy();
192			*a12_a = a12_a.mul_real(norm);
193			let ref b = y2_a
194				+ dot::inner_prod(
195					y2_b,
196					Conj::No,
197					A12_b.rb().transpose(),
198					Conj::Yes,
199				);
200			if m != infinity() {
201				z!(z2.rb_mut(), A21.rb(), A22_a.rb()).for_each(
202					|uz!(z, u, a): Zip!(&mut T, &T, &T)| {
203						let w = &*z - a * beta.conj();
204						let w = w * m.conj();
205						let w = w - u * b;
206						*z = w.mul_real(tr_inv);
207					},
208				);
209			} else {
210				z!(z2.rb_mut(), A21.rb(), A22_a.rb()).for_each(
211					|uz!(z, u, a): Zip!(&mut T, &T, &T)| {
212						let w = a - u * b;
213						*z = w.mul_real(tr_inv);
214					},
215				);
216			}
217		}
218	}
219	let mut j = 0;
220	while j < size {
221		let bl = Ord::min(bl, size - j);
222		let mut Hl = Hl.rb_mut().get_mut(..bl, j..j + bl);
223		for k in 0..bl {
224			Hl[(k, k)] = Hl[(0, k)].copy();
225		}
226		upgrade_householder_factor(
227			Hl.rb_mut(),
228			A.rb().get(j.., j..j + bl),
229			bl,
230			1,
231			par,
232		);
233		j += bl;
234	}
235	if size > 0 {
236		let size = size - 1;
237		let A = A.rb().get(..size, 1..);
238		let mut Hr = Hr.rb_mut().get_mut(.., ..size);
239		let mut j = 0;
240		while j < size {
241			let br = Ord::min(br, size - j);
242			let mut Hr = Hr.rb_mut().get_mut(..br, j..j + br);
243			for k in 0..br {
244				Hr[(k, k)] = Hr[(0, k)].copy();
245			}
246			upgrade_householder_factor(
247				Hr.rb_mut(),
248				A.transpose().get(j.., j..j + br),
249				br,
250				1,
251				par,
252			);
253			j += br;
254		}
255	}
256}
257fn bidiag_fused_op<T: ComplexField>(
258	A22: MatMut<'_, T>,
259	u: ColRef<'_, T>,
260	up: ColRef<'_, T>,
261	z: ColRef<'_, T>,
262	y: RowMut<'_, T>,
263	vp: RowRef<'_, T>,
264	align: usize,
265) {
266	let mut A22 = A22;
267	if const { T::SIMD_CAPABILITIES.is_simd() } {
268		if let (Some(A22), Some(u), Some(up), Some(z)) = (
269			A22.rb_mut().try_as_col_major_mut(),
270			u.try_as_col_major(),
271			up.try_as_col_major(),
272			z.try_as_col_major(),
273		) {
274			bidiag_fused_op_simd(A22, u, up, z, y, vp, align);
275		} else {
276			bidiag_fused_op_fallback(A22, u, up, z, y, vp);
277		}
278	} else {
279		bidiag_fused_op_fallback(A22, u, up, z, y, vp);
280	}
281}
282fn bidiag_fused_op_fallback<T: ComplexField>(
283	A22: MatMut<'_, T>,
284	u: ColRef<'_, T>,
285	up: ColRef<'_, T>,
286	z: ColRef<'_, T>,
287	y: RowMut<'_, T>,
288	vp: RowRef<'_, T>,
289) {
290	let mut A22 = A22;
291	let mut y = y;
292	matmul(A22.rb_mut(), Accum::Add, up, y.rb(), -one::<T>(), Par::Seq);
293	matmul(A22.rb_mut(), Accum::Add, z, vp, -one::<T>(), Par::Seq);
294	matmul(
295		y.rb_mut(),
296		Accum::Replace,
297		u.adjoint(),
298		A22.rb(),
299		one(),
300		Par::Seq,
301	);
302}
303fn bidiag_fused_op_simd<'M, 'N, T: ComplexField>(
304	A22: MatMut<'_, T, usize, usize, ContiguousFwd>,
305	u: ColRef<'_, T, usize, ContiguousFwd>,
306	up: ColRef<'_, T, usize, ContiguousFwd>,
307	z: ColRef<'_, T, usize, ContiguousFwd>,
308	y: RowMut<'_, T, usize>,
309	vp: RowRef<'_, T, usize>,
310	align: usize,
311) {
312	struct Impl<'a, 'M, 'N, T: ComplexField> {
313		A22: MatMut<'a, T, Dim<'M>, Dim<'N>, ContiguousFwd>,
314		u: ColRef<'a, T, Dim<'M>, ContiguousFwd>,
315		up: ColRef<'a, T, Dim<'M>, ContiguousFwd>,
316		z: ColRef<'a, T, Dim<'M>, ContiguousFwd>,
317		y: RowMut<'a, T, Dim<'N>>,
318		vp: RowRef<'a, T, Dim<'N>>,
319		align: usize,
320	}
321	impl<'a, 'M, 'N, T: ComplexField> pulp::WithSimd for Impl<'a, 'M, 'N, T> {
322		type Output = ();
323
324		#[inline(always)]
325		fn with_simd<S: pulp::Simd>(self, simd: S) -> Self::Output {
326			let Self {
327				mut A22,
328				u,
329				up,
330				z,
331				mut y,
332				vp,
333				align,
334			} = self;
335			let m = A22.nrows();
336			let n = A22.ncols();
337			let simd = SimdCtx::<T, S>::new_align(T::simd_ctx(simd), m, align);
338			let indices = simd.batch_indices::<4>();
339			for j in n.indices() {
340				let mut A = A22.rb_mut().col_mut(j);
341				let mut acc = [simd.zero(); 4];
342				let yj = simd.splat(-&y[j]);
343				let vj = simd.splat(-&vp[j]);
344				simd_iter!(for (IDX, i) in [indices; 4] {
345					let mut a = simd.read(A.rb(), i);
346					a = simd.mul_add(simd.read(up, i), yj, a);
347					a = simd.mul_add(simd.read(z, i), vj, a);
348					simd.write(A.rb_mut(), i, a);
349					acc[IDX] = simd.conj_mul_add(simd.read(u, i), a, acc[IDX]);
350				});
351				let acc0 = simd.add(acc[0], acc[1]);
352				let acc2 = simd.add(acc[2], acc[3]);
353				let acc0 = simd.add(acc0, acc2);
354				y[j] = simd.reduce_sum(acc0);
355			}
356		}
357	}
358	with_dim!(M, A22.nrows());
359	with_dim!(N, A22.ncols());
360	dispatch!(
361		Impl {
362			A22: A22.as_shape_mut(M, N),
363			u: u.as_row_shape(M),
364			up: up.as_row_shape(M),
365			z: z.as_row_shape(M),
366			y: y.as_col_shape_mut(N),
367			vp: vp.as_col_shape(N),
368			align,
369		},
370		Impl,
371		T
372	)
373}
374#[cfg(test)]
375mod tests {
376	use super::*;
377	use crate::stats::prelude::*;
378	use crate::utils::approx::*;
379	use crate::{Mat, assert, c64};
380	use dyn_stack::MemBuffer;
381	use std::mem::MaybeUninit;
382	#[test]
383	fn test_bidiag_real() {
384		let rng = &mut StdRng::seed_from_u64(0);
385		for (m, n) in [(8, 4), (8, 8)] {
386			let size = Ord::min(m, n);
387			let A = CwiseMatDistribution {
388				nrows: m,
389				ncols: n,
390				dist: StandardNormal,
391			}
392			.rand::<Mat<f64>>(rng);
393			let bl = 4;
394			let br = 3;
395			let mut Hl = Mat::zeros(bl, size);
396			let mut Hr = Mat::zeros(br, size - 1);
397			let mut UV = A.clone();
398			bidiag_in_place(
399				UV.rb_mut(),
400				Hl.rb_mut(),
401				Hr.rb_mut(),
402				Par::Seq,
403				MemStack::new(&mut [MaybeUninit::uninit(); 1024]),
404				default(),
405			);
406			let mut A = A.clone();
407			let mut A = A.as_mut();
408			apply_block_householder_sequence_transpose_on_the_left_in_place_with_conj(
409				UV.rb().get(.., ..size),
410				Hl.rb(),
411				Conj::Yes,
412				A.rb_mut(),
413				Par::Seq,
414				MemStack::new(&mut MemBuffer::new(
415					apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<f64>(n - 1, 1, m),
416				)),
417			);
418			let V = UV.rb().get(..size - 1, 1..size);
419			let A1 = A.rb_mut().get_mut(.., 1..size);
420			let Hr = Hr.as_ref();
421			apply_block_householder_sequence_on_the_right_in_place_with_conj(
422				V.transpose(),
423				Hr.as_ref(),
424				Conj::Yes,
425				A1,
426				Par::Seq,
427				MemStack::new(&mut MemBuffer::new(
428					apply_block_householder_sequence_on_the_right_in_place_scratch::<f64>(n - 1, 1, m),
429				)),
430			);
431			let approx_eq = CwiseMat(ApproxEq::<f64>::eps());
432			for j in 0..n {
433				for i in 0..m {
434					if i > j || j > i + 1 {
435						UV[(i, j)] = 0.0;
436					}
437				}
438			}
439			assert!(UV ~ A);
440		}
441	}
442	#[test]
443	fn test_bidiag_cplx() {
444		let rng = &mut StdRng::seed_from_u64(0);
445		for (m, n) in [(8, 4), (8, 8)] {
446			let size = Ord::min(m, n);
447			let A = CwiseMatDistribution {
448				nrows: m,
449				ncols: n,
450				dist: ComplexDistribution::new(StandardNormal, StandardNormal),
451			}
452			.rand::<Mat<c64>>(rng);
453			let bl = 4;
454			let br = 3;
455			let mut Hl = Mat::zeros(bl, size);
456			let mut Hr = Mat::zeros(br, size - 1);
457			let mut UV = A.clone();
458			let mut UV = UV.as_mut();
459			bidiag_in_place(
460				UV.rb_mut(),
461				Hl.rb_mut(),
462				Hr.rb_mut(),
463				Par::Seq,
464				MemStack::new(&mut [MaybeUninit::uninit(); 1024]),
465				default(),
466			);
467			let mut A = A.clone();
468			let mut A = A.as_mut();
469			apply_block_householder_sequence_transpose_on_the_left_in_place_with_conj(
470				UV.rb().subcols(0, size),
471				Hl.rb(),
472				Conj::Yes,
473				A.rb_mut(),
474				Par::Seq,
475				MemStack::new(&mut MemBuffer::new(
476					apply_block_householder_sequence_transpose_on_the_left_in_place_scratch::<c64>(n - 1, 1, m),
477				)),
478			);
479			let V = UV.rb().get(..size - 1, 1..size);
480			let A1 = A.rb_mut().get_mut(.., 1..size);
481			let Hr = Hr.rb();
482			apply_block_householder_sequence_on_the_right_in_place_with_conj(
483				V.transpose(),
484				Hr,
485				Conj::Yes,
486				A1,
487				Par::Seq,
488				MemStack::new(&mut MemBuffer::new(
489					apply_block_householder_sequence_on_the_right_in_place_scratch::<c64>(n - 1, 1, m),
490				)),
491			);
492			let approx_eq = CwiseMat(ApproxEq::eps());
493			for j in 0..n {
494				for i in 0..m {
495					if i > j || j > i + 1 {
496						UV[(i, j)] = c64::ZERO;
497					}
498				}
499			}
500			assert!(UV ~ A);
501		}
502	}
503}
faer/linalg/svd/bidiag.rs

faer/linalg/svd/
bidiag.rs
