Skip to content

Commit ea5c96b

Browse files
termoshtttermoshtt
authored
Merge pull request #329 from rust-ndarray/maybe-uninit
`MaybeUninit<T>` for uninitialized memory
2 parents 791713f + b5b92e3 commit ea5c96b

File tree

14 files changed

+187
-75
lines changed

14 files changed

+187
-75
lines changed

‎lax/src/eig.rs

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,12 @@ macro_rules! impl_eig_complex {
4141
} else {
4242
(EigenVectorFlag::Not, EigenVectorFlag::Not)
4343
};
44-
let mut eigs = unsafe { vec_uninit(n as usize) };
45-
let mut rwork: Vec<Self::Real> = unsafe { vec_uninit(2 * n as usize) };
44+
let mut eigs: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(n as usize) };
45+
let mut rwork: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(2 * n as usize) };
4646

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

5252
// calc work size
@@ -74,7 +74,7 @@ macro_rules! impl_eig_complex {
7474

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

99+
let eigs = unsafe { eigs.assume_init() };
100+
let vr = unsafe { vr.map(|v| v.assume_init()) };
101+
let mut vl = unsafe { vl.map(|v| v.assume_init()) };
102+
99103
// Hermite conjugate
100104
if jobvl.is_calc() {
101105
for c in vl.as_mut().unwrap().iter_mut() {
102-
c.im = -c.im
106+
c.im = -c.im;
103107
}
104108
}
105109

@@ -145,12 +149,12 @@ macro_rules! impl_eig_real {
145149
} else {
146150
(EigenVectorFlag::Not, EigenVectorFlag::Not)
147151
};
148-
let mut eig_re: Vec<Self> = unsafe { vec_uninit(n as usize) };
149-
let mut eig_im: Vec<Self> = unsafe { vec_uninit(n as usize) };
152+
let mut eig_re: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(n as usize) };
153+
let mut eig_im: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(n as usize) };
150154

151-
let mut vl: Option<Vec<Self>> =
155+
let mut vl: Option<Vec<MaybeUninit<Self>>> =
152156
jobvl.then(|| unsafe { vec_uninit((n * n) as usize) });
153-
let mut vr: Option<Vec<Self>> =
157+
let mut vr: Option<Vec<MaybeUninit<Self>>> =
154158
jobvr.then(|| unsafe { vec_uninit((n * n) as usize) });
155159

156160
// calc work size
@@ -178,7 +182,7 @@ macro_rules! impl_eig_real {
178182

179183
// actual ev
180184
let lwork = work_size[0].to_usize().unwrap();
181-
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
185+
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
182186
let lwork = lwork as i32;
183187
unsafe {
184188
$ev(
@@ -200,6 +204,11 @@ macro_rules! impl_eig_real {
200204
};
201205
info.as_lapack_result()?;
202206

207+
let eig_re = unsafe { eig_re.assume_init() };
208+
let eig_im = unsafe { eig_im.assume_init() };
209+
let vl = unsafe { vl.map(|v| v.assume_init()) };
210+
let vr = unsafe { vr.map(|v| v.assume_init()) };
211+
203212
// reconstruct eigenvalues
204213
let eigs: Vec<Self::Complex> = eig_re
205214
.iter()
@@ -228,14 +237,14 @@ macro_rules! impl_eig_real {
228237

229238
let n = n as usize;
230239
let v = vr.or(vl).unwrap();
231-
let mut eigvecs = unsafe { vec_uninit(n * n) };
240+
let mut eigvecs: Vec<MaybeUninit<Self::Complex>> = unsafe { vec_uninit(n * n) };
232241
let mut col = 0;
233242
while col < n {
234243
if eig_im[col] == 0. {
235244
// The corresponding eigenvalue is real.
236245
for row in 0..n {
237246
let re = v[row + col * n];
238-
eigvecs[row + col * n] = Self::complex(re, 0.);
247+
eigvecs[row + col * n].write(Self::complex(re, 0.));
239248
}
240249
col += 1;
241250
} else {
@@ -247,12 +256,13 @@ macro_rules! impl_eig_real {
247256
if jobvl.is_calc() {
248257
im = -im;
249258
}
250-
eigvecs[row + col * n] = Self::complex(re, im);
251-
eigvecs[row + (col + 1) * n] = Self::complex(re, -im);
259+
eigvecs[row + col * n].write(Self::complex(re, im));
260+
eigvecs[row + (col + 1) * n].write(Self::complex(re, -im));
252261
}
253262
col += 2;
254263
}
255264
}
265+
let eigvecs = unsafe { eigvecs.assume_init() };
256266

257267
Ok((eigs, eigvecs))
258268
}

‎lax/src/eigh.rs

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,10 @@ macro_rules! impl_eigh {
4242
assert_eq!(layout.len(), layout.lda());
4343
let n = layout.len();
4444
let jobz = if calc_v { EigenVectorFlag::Calc } else { EigenVectorFlag::Not };
45-
let mut eigs = unsafe { vec_uninit(n as usize) };
45+
let mut eigs: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(n as usize) };
4646

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

5151
// calc work size
@@ -69,7 +69,7 @@ macro_rules! impl_eigh {
6969

7070
// actual ev
7171
let lwork = work_size[0].to_usize().unwrap();
72-
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
72+
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
7373
let lwork = lwork as i32;
7474
unsafe {
7575
$ev(
@@ -86,6 +86,8 @@ macro_rules! impl_eigh {
8686
);
8787
}
8888
info.as_lapack_result()?;
89+
90+
let eigs = unsafe { eigs.assume_init() };
8991
Ok(eigs)
9092
}
9193

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

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

108110
// calc work size
@@ -129,7 +131,7 @@ macro_rules! impl_eigh {
129131

130132
// actual evg
131133
let lwork = work_size[0].to_usize().unwrap();
132-
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
134+
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
133135
let lwork = lwork as i32;
134136
unsafe {
135137
$evg(
@@ -149,6 +151,7 @@ macro_rules! impl_eigh {
149151
);
150152
}
151153
info.as_lapack_result()?;
154+
let eigs = unsafe { eigs.assume_init() };
152155
Ok(eigs)
153156
}
154157
}

‎lax/src/layout.rs

Lines changed: 65 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
//! This `S` for a matrix `A` is called "leading dimension of the array A" in LAPACK document, and denoted by `lda`.
3838
//!
3939
40-
use cauchy::Scalar;
40+
use super::*;
4141

4242
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4343
pub enum MatrixLayout {
@@ -153,7 +153,7 @@ impl MatrixLayout {
153153
/// ------
154154
/// - If size of `a` and `layout` size mismatch
155155
///
156-
pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
156+
pub fn square_transpose<T: Copy>(layout: MatrixLayout, a: &mut [T]) {
157157
let (m, n) = layout.size();
158158
let n = n as usize;
159159
let m = m as usize;
@@ -162,23 +162,78 @@ pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
162162
for j in (i + 1)..n {
163163
let a_ij = a[i * n + j];
164164
let a_ji = a[j * m + i];
165-
a[i * n + j] = a_ji.conj();
166-
a[j * m + i] = a_ij.conj();
165+
a[i * n + j] = a_ji;
166+
a[j * m + i] = a_ij;
167167
}
168168
}
169169
}
170170

171171
/// Out-place transpose for general matrix
172172
///
173-
/// Inplace transpose of non-square matrices is hard.
174-
/// See also: https://en.wikipedia.org/wiki/In-place_matrix_transposition
173+
/// Examples
174+
/// ---------
175+
///
176+
/// ```rust
177+
/// # use lax::layout::*;
178+
/// let layout = MatrixLayout::C { row: 2, lda: 3 };
179+
/// let a = vec![1., 2., 3., 4., 5., 6.];
180+
/// let (l, b) = transpose(layout, &a);
181+
/// assert_eq!(l, MatrixLayout::F { col: 3, lda: 2 });
182+
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
183+
/// ```
184+
///
185+
/// ```rust
186+
/// # use lax::layout::*;
187+
/// let layout = MatrixLayout::F { col: 2, lda: 3 };
188+
/// let a = vec![1., 2., 3., 4., 5., 6.];
189+
/// let (l, b) = transpose(layout, &a);
190+
/// assert_eq!(l, MatrixLayout::C { row: 3, lda: 2 });
191+
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
192+
/// ```
193+
///
194+
/// Panics
195+
/// ------
196+
/// - If input array size and `layout` size mismatch
197+
///
198+
pub fn transpose<T: Copy>(layout: MatrixLayout, input: &[T]) -> (MatrixLayout, Vec<T>) {
199+
let (m, n) = layout.size();
200+
let transposed = layout.resized(n, m).t();
201+
let m = m as usize;
202+
let n = n as usize;
203+
assert_eq!(input.len(), m * n);
204+
205+
let mut out: Vec<MaybeUninit<T>> = unsafe { vec_uninit(m * n) };
206+
207+
match layout {
208+
MatrixLayout::C { .. } => {
209+
for i in 0..m {
210+
for j in 0..n {
211+
out[j * m + i].write(input[i * n + j]);
212+
}
213+
}
214+
}
215+
MatrixLayout::F { .. } => {
216+
for i in 0..m {
217+
for j in 0..n {
218+
out[i * n + j].write(input[j * m + i]);
219+
}
220+
}
221+
}
222+
}
223+
(transposed, unsafe { out.assume_init() })
224+
}
225+
226+
/// Out-place transpose for general matrix
227+
///
228+
/// Examples
229+
/// ---------
175230
///
176231
/// ```rust
177232
/// # use lax::layout::*;
178233
/// let layout = MatrixLayout::C { row: 2, lda: 3 };
179234
/// let a = vec![1., 2., 3., 4., 5., 6.];
180235
/// let mut b = vec![0.0; a.len()];
181-
/// let l = transpose(layout, &a, &mut b);
236+
/// let l = transpose_over(layout, &a, &mut b);
182237
/// assert_eq!(l, MatrixLayout::F { col: 3, lda: 2 });
183238
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
184239
/// ```
@@ -188,16 +243,16 @@ pub fn square_transpose<T: Scalar>(layout: MatrixLayout, a: &mut [T]) {
188243
/// let layout = MatrixLayout::F { col: 2, lda: 3 };
189244
/// let a = vec![1., 2., 3., 4., 5., 6.];
190245
/// let mut b = vec![0.0; a.len()];
191-
/// let l = transpose(layout, &a, &mut b);
246+
/// let l = transpose_over(layout, &a, &mut b);
192247
/// assert_eq!(l, MatrixLayout::C { row: 3, lda: 2 });
193248
/// assert_eq!(b, &[1., 4., 2., 5., 3., 6.]);
194249
/// ```
195250
///
196251
/// Panics
197252
/// ------
198-
/// - If size of `a` and `layout` size mismatch
253+
/// - If input array sizes and `layout` size mismatch
199254
///
200-
pub fn transpose<T: Scalar>(layout: MatrixLayout, from: &[T], to: &mut [T]) -> MatrixLayout {
255+
pub fn transpose_over<T: Copy>(layout: MatrixLayout, from: &[T], to: &mut [T]) -> MatrixLayout {
201256
let (m, n) = layout.size();
202257
let transposed = layout.resized(n, m).t();
203258
let m = m as usize;

‎lax/src/least_squares.rs

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,9 @@ macro_rules! impl_least_squares {
6868
let mut a_t = None;
6969
let a_layout = match a_layout {
7070
MatrixLayout::C { .. } => {
71-
a_t = Some(unsafe { vec_uninit( a.len()) });
72-
transpose(a_layout, a, a_t.as_mut().unwrap())
71+
let (layout, t) = transpose(a_layout, a);
72+
a_t = Some(t);
73+
layout
7374
}
7475
MatrixLayout::F { .. } => a_layout,
7576
};
@@ -78,14 +79,15 @@ macro_rules! impl_least_squares {
7879
let mut b_t = None;
7980
let b_layout = match b_layout {
8081
MatrixLayout::C { .. } => {
81-
b_t = Some(unsafe { vec_uninit( b.len()) });
82-
transpose(b_layout, b, b_t.as_mut().unwrap())
82+
let (layout, t) = transpose(b_layout, b);
83+
b_t = Some(t);
84+
layout
8385
}
8486
MatrixLayout::F { .. } => b_layout,
8587
};
8688

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

9193
// eval work size
@@ -118,12 +120,12 @@ macro_rules! impl_least_squares {
118120

119121
// calc
120122
let lwork = work_size[0].to_usize().unwrap();
121-
let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
123+
let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit(lwork) };
122124
let liwork = iwork_size[0].to_usize().unwrap();
123-
let mut iwork = unsafe { vec_uninit(liwork) };
125+
let mut iwork: Vec<MaybeUninit<i32>> = unsafe { vec_uninit(liwork) };
124126
$(
125127
let lrwork = $rwork[0].to_usize().unwrap();
126-
let mut $rwork: Vec<Self::Real> = unsafe { vec_uninit(lrwork) };
128+
let mut $rwork: Vec<MaybeUninit<Self::Real>> = unsafe { vec_uninit(lrwork) };
127129
)*
128130
unsafe {
129131
$gelsd(
@@ -140,16 +142,18 @@ macro_rules! impl_least_squares {
140142
AsPtr::as_mut_ptr(&mut work),
141143
&(lwork as i32),
142144
$(AsPtr::as_mut_ptr(&mut $rwork),)*
143-
iwork.as_mut_ptr(),
145+
AsPtr::as_mut_ptr(&mut iwork),
144146
&mut info,
145147
);
146148
}
147149
info.as_lapack_result()?;
148150

151+
let singular_values = unsafe { singular_values.assume_init() };
152+
149153
// Skip a_t -> a transpose because A has been destroyed
150154
// Re-transpose b
151155
if let Some(b_t) = b_t {
152-
transpose(b_layout, &b_t, b);
156+
transpose_over(b_layout, &b_t, b);
153157
}
154158

155159
Ok(LeastSquaresOutput {

0 commit comments

Comments
 (0)