Skip to content

Commit

Permalink
Merge pull request #334 from rust-ndarray/safe-lax-vec-uninit
Browse files Browse the repository at this point in the history 
Safe `lax::vec_uninit`
  • Loading branch information
@termoshtt
termoshtt committed Sep 8, 2022
2 parents f7356c7 + 4737b9c commit 87107be
Show file tree
Hide file tree
Showing 14 changed files with 116 additions and 110 deletions.
62 changes: 62 additions & 0 deletions lax/src/alloc.rs
View file
@@ -0,0 +1,62 @@
use cauchy::*;
use std::mem::MaybeUninit;

/// Helper for getting pointer of slice
pub(crate) trait AsPtr: Sized {
type Elem;
fn as_ptr(vec: &[Self]) -> *const Self::Elem;
fn as_mut_ptr(vec: &mut [Self]) -> *mut Self::Elem;
}

macro_rules! impl_as_ptr {
($target:ty, $elem:ty) => {
impl AsPtr for $target {
type Elem = $elem;
fn as_ptr(vec: &[Self]) -> *const Self::Elem {
vec.as_ptr() as *const _
}
fn as_mut_ptr(vec: &mut [Self]) -> *mut Self::Elem {
vec.as_mut_ptr() as *mut _
}
}
};
}
impl_as_ptr!(i32, i32);
impl_as_ptr!(f32, f32);
impl_as_ptr!(f64, f64);
impl_as_ptr!(c32, lapack_sys::__BindgenComplex<f32>);
impl_as_ptr!(c64, lapack_sys::__BindgenComplex<f64>);
impl_as_ptr!(MaybeUninit<i32>, i32);
impl_as_ptr!(MaybeUninit<f32>, f32);
impl_as_ptr!(MaybeUninit<f64>, f64);
impl_as_ptr!(MaybeUninit<c32>, lapack_sys::__BindgenComplex<f32>);
impl_as_ptr!(MaybeUninit<c64>, lapack_sys::__BindgenComplex<f64>);

pub(crate) trait VecAssumeInit {
type Target;
unsafe fn assume_init(self) -> Self::Target;
}

impl<T> VecAssumeInit for Vec<MaybeUninit<T>> {
type Target = Vec<T>;
unsafe fn assume_init(self) -> Self::Target {
// FIXME use Vec::into_raw_parts instead after stablized
// https://doc.rust-lang.org/std/vec/struct.Vec.html#method.into_raw_parts
let mut me = std::mem::ManuallyDrop::new(self);
Vec::from_raw_parts(me.as_mut_ptr() as *mut T, me.len(), me.capacity())
}
}

/// Create a vector without initialization
///
/// Safety
/// ------
/// - Memory is not initialized. Do not read the memory before write.
///
pub(crate) fn vec_uninit<T: Sized>(n: usize) -> Vec<MaybeUninit<T>> {
let mut v = Vec::with_capacity(n);
unsafe {
v.set_len(n);
}
v
}
22 changes: 11 additions & 11 deletions lax/src/eig.rs
View file
Expand Up @@ -48,13 +48,13 @@ macro_rules! impl_eig_complex {
} else {
(JobEv::None, JobEv::None)
};
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 eigs: Vec<MaybeUninit<Self>> = vec_uninit(n as usize);
let mut rwork: Vec<MaybeUninit<Self::Real>> = vec_uninit(2 * n as usize);

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

// calc work size
let mut info = 0;
Expand All @@ -81,7 +81,7 @@ macro_rules! impl_eig_complex {

// actal ev
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
let lwork = lwork as i32;
unsafe {
$ev(
Expand Down Expand Up @@ -156,13 +156,13 @@ macro_rules! impl_eig_real {
} else {
(JobEv::None, JobEv::None)
};
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 eig_re: Vec<MaybeUninit<Self>> = vec_uninit(n as usize);
let mut eig_im: Vec<MaybeUninit<Self>> = vec_uninit(n as usize);

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

// calc work size
let mut info = 0;
Expand All @@ -189,7 +189,7 @@ macro_rules! impl_eig_real {

// actual ev
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
let lwork = lwork as i32;
unsafe {
$ev(
Expand Down Expand Up @@ -244,7 +244,7 @@ macro_rules! impl_eig_real {

let n = n as usize;
let v = vr.or(vl).unwrap();
let mut eigvecs: Vec<MaybeUninit<Self::Complex>> = unsafe { vec_uninit(n * n) };
let mut eigvecs: Vec<MaybeUninit<Self::Complex>> = vec_uninit(n * n);
let mut col = 0;
while col < n {
if eig_im[col] == 0. {
Expand Down
12 changes: 6 additions & 6 deletions lax/src/eigh.rs
View file
Expand Up @@ -58,10 +58,10 @@ macro_rules! impl_eigh {
assert_eq!(layout.len(), layout.lda());
let n = layout.len();
let jobz = if calc_v { JobEv::All } else { JobEv::None };
let mut eigs: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(n as usize) };
let mut eigs: Vec<MaybeUninit<Self::Real>> = vec_uninit(n as usize);

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

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

// actual ev
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
let lwork = lwork as i32;
unsafe {
$ev(
Expand Down Expand Up @@ -117,10 +117,10 @@ macro_rules! impl_eigh {
assert_eq!(layout.len(), layout.lda());
let n = layout.len();
let jobz = if calc_v { JobEv::All } else { JobEv::None };
let mut eigs: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(n as usize) };
let mut eigs: Vec<MaybeUninit<Self::Real>> = vec_uninit(n as usize);

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

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

// actual evg
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
let lwork = lwork as i32;
unsafe {
$evg(
Expand Down
2 changes: 1 addition & 1 deletion lax/src/layout.rs
View file
Expand Up @@ -202,7 +202,7 @@ pub fn transpose<T: Copy>(layout: MatrixLayout, input: &[T]) -> (MatrixLayout, V
let n = n as usize;
assert_eq!(input.len(), m * n);

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

match layout {
MatrixLayout::C { .. } => {
Expand Down
8 changes: 4 additions & 4 deletions lax/src/least_squares.rs
View file
Expand Up @@ -91,7 +91,7 @@ macro_rules! impl_least_squares {
};

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

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

// calc
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
let liwork = iwork_size[0].to_usize().unwrap();
let mut iwork: Vec<MaybeUninit<i32>> = unsafe { vec_uninit(liwork) };
let mut iwork: Vec<MaybeUninit<i32>> = vec_uninit(liwork);
$(
let lrwork = $rwork[0].to_usize().unwrap();
let mut $rwork: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(lrwork) };
let mut $rwork: Vec<MaybeUninit<Self::Real>> = vec_uninit(lrwork);
)*
unsafe {
$gelsd(
Expand Down
60 changes: 2 additions & 58 deletions lax/src/lib.rs
View file
Expand Up @@ -84,6 +84,7 @@ pub mod error;
pub mod flags;
pub mod layout;

mod alloc;
mod cholesky;
mod eig;
mod eigh;
Expand Down Expand Up @@ -113,6 +114,7 @@ pub use self::svddc::*;
pub use self::triangular::*;
pub use self::tridiagonal::*;

use self::alloc::*;
use cauchy::*;
use std::mem::MaybeUninit;

Expand Down Expand Up @@ -140,61 +142,3 @@ impl Lapack for f32 {}
impl Lapack for f64 {}
impl Lapack for c32 {}
impl Lapack for c64 {}

/// Helper for getting pointer of slice
pub(crate) trait AsPtr: Sized {
type Elem;
fn as_ptr(vec: &[Self]) -> *const Self::Elem;
fn as_mut_ptr(vec: &mut [Self]) -> *mut Self::Elem;
}

macro_rules! impl_as_ptr {
($target:ty, $elem:ty) => {
impl AsPtr for $target {
type Elem = $elem;
fn as_ptr(vec: &[Self]) -> *const Self::Elem {
vec.as_ptr() as *const _
}
fn as_mut_ptr(vec: &mut [Self]) -> *mut Self::Elem {
vec.as_mut_ptr() as *mut _
}
}
};
}
impl_as_ptr!(i32, i32);
impl_as_ptr!(f32, f32);
impl_as_ptr!(f64, f64);
impl_as_ptr!(c32, lapack_sys::__BindgenComplex<f32>);
impl_as_ptr!(c64, lapack_sys::__BindgenComplex<f64>);
impl_as_ptr!(MaybeUninit<i32>, i32);
impl_as_ptr!(MaybeUninit<f32>, f32);
impl_as_ptr!(MaybeUninit<f64>, f64);
impl_as_ptr!(MaybeUninit<c32>, lapack_sys::__BindgenComplex<f32>);
impl_as_ptr!(MaybeUninit<c64>, lapack_sys::__BindgenComplex<f64>);

pub(crate) trait VecAssumeInit {
type Target;
unsafe fn assume_init(self) -> Self::Target;
}

impl<T> VecAssumeInit for Vec<MaybeUninit<T>> {
type Target = Vec<T>;
unsafe fn assume_init(self) -> Self::Target {
// FIXME use Vec::into_raw_parts instead after stablized
// https://doc.rust-lang.org/std/vec/struct.Vec.html#method.into_raw_parts
let mut me = std::mem::ManuallyDrop::new(self);
Vec::from_raw_parts(me.as_mut_ptr() as *mut T, me.len(), me.capacity())
}
}

/// Create a vector without initialization
///
/// Safety
/// ------
/// - Memory is not initialized. Do not read the memory before write.
///
unsafe fn vec_uninit<T: Sized>(n: usize) -> Vec<MaybeUninit<T>> {
let mut v = Vec::with_capacity(n);
v.set_len(n);
v
}
2 changes: 1 addition & 1 deletion lax/src/opnorm.rs
View file
Expand Up @@ -19,7 +19,7 @@ macro_rules! impl_opnorm {
MatrixLayout::C { .. } => t.transpose(),
};
let mut work: Vec<MaybeUninit<Self::Real>> = if matches!(t, NormType::Infinity) {
unsafe { vec_uninit(m as usize) }
vec_uninit(m as usize)
} else {
Vec::new()
};
Expand Down
6 changes: 3 additions & 3 deletions lax/src/qr.rs
View file
Expand Up @@ -25,7 +25,7 @@ macro_rules! impl_qr {
let m = l.lda();
let n = l.len();
let k = m.min(n);
let mut tau = unsafe { vec_uninit(k as usize) };
let mut tau = vec_uninit(k as usize);

// eval work size
let mut info = 0;
Expand Down Expand Up @@ -62,7 +62,7 @@ macro_rules! impl_qr {

// calc
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
unsafe {
match l {
MatrixLayout::F { .. } => {
Expand Down Expand Up @@ -136,7 +136,7 @@ macro_rules! impl_qr {

// calc
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
unsafe {
match l {
MatrixLayout::F { .. } => $gqr(
Expand Down
8 changes: 4 additions & 4 deletions lax/src/rcond.rs
View file
Expand Up @@ -17,8 +17,8 @@ macro_rules! impl_rcond_real {
let mut rcond = Self::Real::zero();
let mut info = 0;

let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(4 * n as usize) };
let mut iwork: Vec<MaybeUninit<i32>> = unsafe { vec_uninit(n as usize) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(4 * n as usize);
let mut iwork: Vec<MaybeUninit<i32>> = vec_uninit(n as usize);
let norm_type = match l {
MatrixLayout::C { .. } => NormType::Infinity,
MatrixLayout::F { .. } => NormType::One,
Expand Down Expand Up @@ -54,8 +54,8 @@ macro_rules! impl_rcond_complex {
let (n, _) = l.size();
let mut rcond = Self::Real::zero();
let mut info = 0;
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(2 * n as usize) };
let mut rwork: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(2 * n as usize) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(2 * n as usize);
let mut rwork: Vec<MaybeUninit<Self::Real>> = vec_uninit(2 * n as usize);
let norm_type = match l {
MatrixLayout::C { .. } => NormType::Infinity,
MatrixLayout::F { .. } => NormType::One,
Expand Down
4 changes: 2 additions & 2 deletions lax/src/solve.rs
View file
Expand Up @@ -75,7 +75,7 @@ macro_rules! impl_solve {
return Ok(Vec::new());
}
let k = ::std::cmp::min(row, col);
let mut ipiv = unsafe { vec_uninit(k as usize) };
let mut ipiv = vec_uninit(k as usize);
let mut info = 0;
unsafe {
$getrf(
Expand Down Expand Up @@ -117,7 +117,7 @@ macro_rules! impl_solve {

// actual
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
unsafe {
$getri(
&l.len(),
Expand Down
6 changes: 3 additions & 3 deletions lax/src/solveh.rs
View file
Expand Up @@ -58,7 +58,7 @@ macro_rules! impl_solveh {
impl Solveh_ for $scalar {
fn bk(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<Pivot> {
let (n, _) = l.size();
let mut ipiv = unsafe { vec_uninit(n as usize) };
let mut ipiv = vec_uninit(n as usize);
if n == 0 {
return Ok(Vec::new());
}
Expand All @@ -82,7 +82,7 @@ macro_rules! impl_solveh {

// actual
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
unsafe {
$trf(
uplo.as_ptr(),
Expand All @@ -103,7 +103,7 @@ macro_rules! impl_solveh {
fn invh(l: MatrixLayout, uplo: UPLO, a: &mut [Self], ipiv: &Pivot) -> Result<()> {
let (n, _) = l.size();
let mut info = 0;
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(n as usize) };
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(n as usize);
unsafe {
$tri(
uplo.as_ptr(),
Expand Down

0 comments on commit 87107be

Please sign in to comment.