Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

MaybeUninit<T> for uninitialized memory #329

Merged
merged 10 commits into from Sep 1, 2022
Merged

MaybeUninit<T> for uninitialized memory #329

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
6e5b4b0
impl AsPtr for Vec<MaybeUninit<T>>
termoshtt Aug 31, 2022
e35bdbb
Introduce VecAssumeInit to convert Vec<MaybeUninit<T>> to Vec<T>
termoshtt Aug 31, 2022
e6d03a5
Introduce alternative `vec_uninit2`, and replace `vec_uninit` in eig.rs
termoshtt Aug 31, 2022
8c63bc7
Use vec_uninit2 in eigh.rs
termoshtt Aug 31, 2022
ee00a0d
Impl VecAssumeInit for any Vec<T>
termoshtt Sep 1, 2022
152e3d5
Rename `transpose` to `transpose_over`, create new `transpose` to wri…
termoshtt Sep 1, 2022
0669391
Use new `transpose` and `transpose_over`
termoshtt Sep 1, 2022
41a3247
Use vec_uninit2 in qr.rs, opnorm.rs
termoshtt Sep 1, 2022
dd425f4
Use vec_uninit2 in all
termoshtt Sep 1, 2022
b5b92e3
Replace vec_uninit by vec_uninit2
termoshtt Sep 1, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
40 changes: 25 additions & 15 deletions lax/src/eig.rs
View file
Open in desktop
Expand Up @@ -41,12 +41,12 @@ macro_rules! impl_eig_complex {
} else {
(EigenVectorFlag::Not, EigenVectorFlag::Not)
};
let mut eigs = unsafe { vec_uninit(n as usize) };
let mut rwork: Vec<Self::Real> = unsafe { vec_uninit(2 * n as usize) };
let mut eigs: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(n as usize) };
let mut rwork: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(2 * n as usize) };

let mut vl: Option<Vec<Self>> =
let mut vl: Option<Vec<MaybeUninit<Self>>> =
jobvl.then(|| unsafe { vec_uninit((n * n) as usize) });
let mut vr: Option<Vec<Self>> =
let mut vr: Option<Vec<MaybeUninit<Self>>> =
jobvr.then(|| unsafe { vec_uninit((n * n) as usize) });

// calc work size
Expand Down Expand Up @@ -74,7 +74,7 @@ macro_rules! impl_eig_complex {

// actal ev
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let lwork = lwork as i32;
unsafe {
$ev(
Expand All @@ -96,10 +96,14 @@ macro_rules! impl_eig_complex {
};
info.as_lapack_result()?;

let eigs = unsafe { eigs.assume_init() };
let vr = unsafe { vr.map(|v| v.assume_init()) };
let mut vl = unsafe { vl.map(|v| v.assume_init()) };

// Hermite conjugate
if jobvl.is_calc() {
for c in vl.as_mut().unwrap().iter_mut() {
c.im = -c.im
c.im = -c.im;
}
}

Expand Down Expand Up @@ -145,12 +149,12 @@ macro_rules! impl_eig_real {
} else {
(EigenVectorFlag::Not, EigenVectorFlag::Not)
};
let mut eig_re: Vec<Self> = unsafe { vec_uninit(n as usize) };
let mut eig_im: Vec<Self> = unsafe { vec_uninit(n as usize) };
let mut eig_re: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(n as usize) };
let mut eig_im: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(n as usize) };

let mut vl: Option<Vec<Self>> =
let mut vl: Option<Vec<MaybeUninit<Self>>> =
jobvl.then(|| unsafe { vec_uninit((n * n) as usize) });
let mut vr: Option<Vec<Self>> =
let mut vr: Option<Vec<MaybeUninit<Self>>> =
jobvr.then(|| unsafe { vec_uninit((n * n) as usize) });

// calc work size
Expand Down Expand Up @@ -178,7 +182,7 @@ macro_rules! impl_eig_real {

// actual ev
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let lwork = lwork as i32;
unsafe {
$ev(
Expand All @@ -200,6 +204,11 @@ macro_rules! impl_eig_real {
};
info.as_lapack_result()?;

let eig_re = unsafe { eig_re.assume_init() };
let eig_im = unsafe { eig_im.assume_init() };
let vl = unsafe { vl.map(|v| v.assume_init()) };
let vr = unsafe { vr.map(|v| v.assume_init()) };

// reconstruct eigenvalues
let eigs: Vec<Self::Complex> = eig_re
.iter()
Expand Down Expand Up @@ -228,14 +237,14 @@ macro_rules! impl_eig_real {

let n = n as usize;
let v = vr.or(vl).unwrap();
let mut eigvecs = unsafe { vec_uninit(n * n) };
let mut eigvecs: Vec<MaybeUninit<Self::Complex>> = unsafe { vec_uninit(n * n) };
let mut col = 0;
while col < n {
if eig_im[col] == 0. {
// The corresponding eigenvalue is real.
for row in 0..n {
let re = v[row + col * n];
eigvecs[row + col * n] = Self::complex(re, 0.);
eigvecs[row + col * n].write(Self::complex(re, 0.));
}
col += 1;
} else {
Expand All @@ -247,12 +256,13 @@ macro_rules! impl_eig_real {
if jobvl.is_calc() {
im = -im;
}
eigvecs[row + col * n] = Self::complex(re, im);
eigvecs[row + (col + 1) * n] = Self::complex(re, -im);
eigvecs[row + col * n].write(Self::complex(re, im));
eigvecs[row + (col + 1) * n].write(Self::complex(re, -im));
}
col += 2;
}
}
let eigvecs = unsafe { eigvecs.assume_init() };

Ok((eigs, eigvecs))
}
Expand Down
15 changes: 9 additions & 6 deletions lax/src/eigh.rs
View file
Open in desktop
Expand Up @@ -42,10 +42,10 @@ macro_rules! impl_eigh {
assert_eq!(layout.len(), layout.lda());
let n = layout.len();
let jobz = if calc_v { EigenVectorFlag::Calc } else { EigenVectorFlag::Not };
let mut eigs = unsafe { vec_uninit(n as usize) };
let mut eigs: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(n as usize) };

$(
let mut $rwork_ident: Vec<Self::Real> = unsafe { vec_uninit(3 * n as usize - 2 as usize) };
let mut $rwork_ident: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(3 * n as usize - 2 as usize) };
)*

// calc work size
Expand All @@ -69,7 +69,7 @@ macro_rules! impl_eigh {

// actual ev
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let lwork = lwork as i32;
unsafe {
$ev(
Expand All @@ -86,6 +86,8 @@ macro_rules! impl_eigh {
);
}
info.as_lapack_result()?;

let eigs = unsafe { eigs.assume_init() };
Ok(eigs)
}

Expand All @@ -99,10 +101,10 @@ macro_rules! impl_eigh {
assert_eq!(layout.len(), layout.lda());
let n = layout.len();
let jobz = if calc_v { EigenVectorFlag::Calc } else { EigenVectorFlag::Not };
let mut eigs = unsafe { vec_uninit(n as usize) };
let mut eigs: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(n as usize) };

$(
let mut $rwork_ident: Vec<Self::Real> = unsafe { vec_uninit(3 * n as usize - 2) };
let mut $rwork_ident: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(3 * n as usize - 2) };
)*

// calc work size
Expand All @@ -129,7 +131,7 @@ macro_rules! impl_eigh {

// actual evg
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let lwork = lwork as i32;
unsafe {
$evg(
Expand All @@ -149,6 +151,7 @@ macro_rules! impl_eigh {
);
}
info.as_lapack_result()?;
let eigs = unsafe { eigs.assume_init() };
Ok(eigs)
}
}
Expand Down
75 changes: 65 additions & 10 deletions lax/src/layout.rs
View file
Open in desktop
Expand Up @@ -37,7 +37,7 @@
//! This `S` for a matrix `A` is called "leading dimension of the array A" in LAPACK document, and denoted by `lda`.
//!

use cauchy::Scalar;
use super::*;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MatrixLayout {
Expand Down Expand Up @@ -153,7 +153,7 @@ impl MatrixLayout {
/// ------
/// - If size of `a` and `layout` size mismatch
///
pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
pub fn square_transpose<T: Copy>(layout: MatrixLayout, a: &mut [T]) {
let (m, n) = layout.size();
let n = n as usize;
let m = m as usize;
Expand All @@ -162,23 +162,78 @@ pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
for j in (i + 1)..n {
let a_ij = a[i * n + j];
let a_ji = a[j * m + i];
a[i * n + j] = a_ji.conj();
a[j * m + i] = a_ij.conj();
a[i * n + j] = a_ji;
a[j * m + i] = a_ij;
}
}
}

/// Out-place transpose for general matrix
///
/// Inplace transpose of non-square matrices is hard.
/// See also: https://en.wikipedia.org/wiki/In-place_matrix_transposition
/// Examples
/// ---------
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::C { row: 2, lda: 3 };
/// let a = vec![1., 2., 3., 4., 5., 6.];
/// let (l, b) = transpose(layout, &a);
/// assert_eq!(l, MatrixLayout::F { col: 3, lda: 2 });
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
/// ```
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::F { col: 2, lda: 3 };
/// let a = vec![1., 2., 3., 4., 5., 6.];
/// let (l, b) = transpose(layout, &a);
/// assert_eq!(l, MatrixLayout::C { row: 3, lda: 2 });
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
/// ```
///
/// Panics
/// ------
/// - If input array size and `layout` size mismatch
///
pub fn transpose<T: Copy>(layout: MatrixLayout, input: &[T]) -> (MatrixLayout, Vec<T>) {
let (m, n) = layout.size();
let transposed = layout.resized(n, m).t();
let m = m as usize;
let n = n as usize;
assert_eq!(input.len(), m * n);

let mut out: Vec<MaybeUninit<T>> = unsafe { vec_uninit(m * n) };

match layout {
MatrixLayout::C { .. } => {
for i in 0..m {
for j in 0..n {
out[j * m + i].write(input[i * n + j]);
}
}
}
MatrixLayout::F { .. } => {
for i in 0..m {
for j in 0..n {
out[i * n + j].write(input[j * m + i]);
}
}
}
}
(transposed, unsafe { out.assume_init() })
}

/// Out-place transpose for general matrix
///
/// Examples
/// ---------
///
/// ```rust
/// # use lax::layout::*;
/// let layout = MatrixLayout::C { row: 2, lda: 3 };
/// let a = vec![1., 2., 3., 4., 5., 6.];
/// let mut b = vec![0.0; a.len()];
/// let l = transpose(layout, &a, &mut b);
/// let l = transpose_over(layout, &a, &mut b);
/// assert_eq!(l, MatrixLayout::F { col: 3, lda: 2 });
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
/// ```
Expand All @@ -188,16 +243,16 @@ pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
/// let layout = MatrixLayout::F { col: 2, lda: 3 };
/// let a = vec![1., 2., 3., 4., 5., 6.];
/// let mut b = vec![0.0; a.len()];
/// let l = transpose(layout, &a, &mut b);
/// let l = transpose_over(layout, &a, &mut b);
/// assert_eq!(l, MatrixLayout::C { row: 3, lda: 2 });
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
/// ```
///
/// Panics
/// ------
/// - If size of `a` and `layout` size mismatch
/// - If input array sizes and `layout` size mismatch
///
pub fn transpose<T: Scalar>(layout: MatrixLayout, from: &[T], to: &mut [T]) -> MatrixLayout {
pub fn transpose_over<T: Copy>(layout: MatrixLayout, from: &[T], to: &mut [T]) -> MatrixLayout {
let (m, n) = layout.size();
let transposed = layout.resized(n, m).t();
let m = m as usize;
Expand Down
24 changes: 14 additions & 10 deletions lax/src/least_squares.rs
View file
Open in desktop
Expand Up @@ -68,8 +68,9 @@ macro_rules! impl_least_squares {
let mut a_t = None;
let a_layout = match a_layout {
MatrixLayout::C { .. } => {
a_t = Some(unsafe { vec_uninit( a.len()) });
transpose(a_layout, a, a_t.as_mut().unwrap())
let (layout, t) = transpose(a_layout, a);
a_t = Some(t);
layout
}
MatrixLayout::F { .. } => a_layout,
};
Expand All @@ -78,14 +79,15 @@ macro_rules! impl_least_squares {
let mut b_t = None;
let b_layout = match b_layout {
MatrixLayout::C { .. } => {
b_t = Some(unsafe { vec_uninit( b.len()) });
transpose(b_layout, b, b_t.as_mut().unwrap())
let (layout, t) = transpose(b_layout, b);
b_t = Some(t);
layout
}
MatrixLayout::F { .. } => b_layout,
};

let rcond: Self::Real = -1.;
let mut singular_values: Vec<Self::Real> = unsafe { vec_uninit( k as usize) };
let mut singular_values: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit( k as usize) };
let mut rank: i32 = 0;

// eval work size
Expand Down Expand Up @@ -118,12 +120,12 @@ macro_rules! impl_least_squares {

// calc
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let liwork = iwork_size[0].to_usize().unwrap();
let mut iwork = unsafe { vec_uninit(liwork) };
let mut iwork: Vec<MaybeUninit<i32>> = unsafe { vec_uninit(liwork) };
$(
let lrwork = $rwork[0].to_usize().unwrap();
let mut $rwork: Vec<Self::Real> = unsafe { vec_uninit(lrwork) };
let mut $rwork: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(lrwork) };
)*
unsafe {
$gelsd(
Expand All @@ -140,16 +142,18 @@ macro_rules! impl_least_squares {
AsPtr::as_mut_ptr(&mut work),
&(lwork as i32),
$(AsPtr::as_mut_ptr(&mut $rwork),)*
iwork.as_mut_ptr(),
AsPtr::as_mut_ptr(&mut iwork),
&mut info,
);
}
info.as_lapack_result()?;

let singular_values = unsafe { singular_values.assume_init() };

// Skip a_t -> a transpose because A has been destroyed
// Re-transpose b
if let Some(b_t) = b_t {
transpose(b_layout, &b_t, b);
transpose_over(b_layout, &b_t, b);
}

Ok(LeastSquaresOutput {
Expand Down