Skip to content

Commit

Permalink
Merge pull request #329 from rust-ndarray/maybe-uninit
Browse files Browse the repository at this point in the history 
`MaybeUninit<T>` for uninitialized memory
  • Loading branch information
@termoshtt
termoshtt committed Sep 1, 2022
2 parents 791713f + b5b92e3 commit ea5c96b
Show file tree
Hide file tree
Showing 14 changed files with 187 additions and 75 deletions.
40 changes: 25 additions & 15 deletions lax/src/eig.rs
View file
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
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
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
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

0 comments on commit ea5c96b

Please sign in to comment.