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

Explicit memory management of working memory #241

Closed
wants to merge 7 commits into from
Closed

Explicit memory management of working memory #241

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b8c35bd
Working memory for eigh (temporary disable eigh_generalized)
termoshtt Jul 29, 2020
a2bcae2
Store size info in lwork
termoshtt Jul 29, 2020
b9651e5
Add generalized eigenvalue problem for symmetric/Hermitian matrices
termoshtt Jul 29, 2020
2777558
vec_uninit
termoshtt Jul 31, 2020
91afa0c
Cholesky factorization
termoshtt Jul 31, 2020
d7eea76
Split Lapack trait definition into sub file
termoshtt Jul 31, 2020
ba45af3
LapackStrict trait for strict memory management API
termoshtt Jul 31, 2020
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
14 changes: 3 additions & 11 deletions lax/src/cholesky.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,24 +4,16 @@ use super::*;
use crate::{error::*, layout::*};
use cauchy::*;

pub trait Cholesky_: Sized {
/// Cholesky: wrapper of `*potrf`
///
/// **Warning: Only the portion of `a` corresponding to `UPLO` is written.**
/// Wrapper trait to switch triangular factorization `*{po,he}tr{f,i,s}`
pub(crate) trait Cholesky: Sized {
fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;

/// Wrapper of `*potri`
///
/// **Warning: Only the portion of `a` corresponding to `UPLO` is written.**
fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;

/// Wrapper of `*potrs`
fn solve_cholesky(l: MatrixLayout, uplo: UPLO, a: &[Self], b: &mut [Self]) -> Result<()>;
}

macro_rules! impl_cholesky {
($scalar:ty, $trf:path, $tri:path, $trs:path) => {
impl Cholesky_ for $scalar {
impl Cholesky for $scalar {
fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
let (n, _) = l.size();
if matches!(l, MatrixLayout::C { .. }) {
Expand Down
187 changes: 97 additions & 90 deletions lax/src/eigh.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,155 +5,162 @@ use crate::{error::*, layout::MatrixLayout};
use cauchy::*;
use num_traits::{ToPrimitive, Zero};

pub trait Eigh_: Scalar {
/// Wraps `*syev` for real and `*heev` for complex
fn eigh(
calc_eigenvec: bool,
layout: MatrixLayout,
uplo: UPLO,
a: &mut [Self],
) -> Result<Vec<Self::Real>>;
pub(crate) trait Eigh: Scalar {
/// Allocate working memory for eigenvalue problem
fn eigh_work(calc_eigenvec: bool, layout: MatrixLayout, uplo: UPLO) -> Result<EighWork<Self>>;

/// Wraps `*syegv` for real and `*heegv` for complex
fn eigh_generalized(
calc_eigenvec: bool,
layout: MatrixLayout,
uplo: UPLO,
/// Solve eigenvalue problem
fn eigh_calc<'work>(
work: &'work mut EighWork<Self>,
a: &mut [Self],
b: &mut [Self],
) -> Result<Vec<Self::Real>>;
) -> Result<&'work [Self::Real]>;
}

macro_rules! impl_eigh {
(@real, $scalar:ty, $ev:path, $evg:path) => {
impl_eigh!(@body, $scalar, $ev, $evg, );
};
(@complex, $scalar:ty, $ev:path, $evg:path) => {
impl_eigh!(@body, $scalar, $ev, $evg, rwork);
};
(@body, $scalar:ty, $ev:path, $evg:path, $($rwork_ident:ident),*) => {
impl Eigh_ for $scalar {
fn eigh(
calc_v: bool,
layout: MatrixLayout,
uplo: UPLO,
mut a: &mut [Self],
) -> Result<Vec<Self::Real>> {
/// Working memory for symmetric/Hermitian eigenvalue problem. See [LapackStrict trait](trait.LapackStrict.html)
pub struct EighWork<T: Scalar> {
jobz: u8,
uplo: UPLO,
n: i32,
eigs: Vec<T::Real>,
// This array is NOT initialized. Do not touch from Rust.
work: Vec<T>,
// Needs only for complex case
rwork: Option<Vec<T::Real>>,
}

macro_rules! impl_eigh_work_real {
($scalar:ty, $ev:path) => {
impl Eigh for $scalar {
fn eigh_work(calc_v: bool, layout: MatrixLayout, uplo: UPLO) -> Result<EighWork<Self>> {
assert_eq!(layout.len(), layout.lda());
let n = layout.len();
let jobz = if calc_v { b'V' } else { b'N' };
let mut eigs = unsafe { vec_uninit(n as usize) };

$(
let mut $rwork_ident = unsafe { vec_uninit(3 * n as usize - 2 as usize) };
)*

// calc work size
let mut info = 0;
let mut work_size = [Self::zero()];
unsafe {
$ev(
jobz,
uplo as u8,
n,
&mut a,
&mut [], // matrix A is not referenced in query mode
n,
&mut eigs,
&mut work_size,
-1,
$(&mut $rwork_ident,)*
&mut info,
);
}
info.as_lapack_result()?;

// actual ev
let lwork = work_size[0].to_usize().unwrap();
let mut work = unsafe { vec_uninit(lwork) };
let work = unsafe { vec_uninit(lwork) };
Ok(EighWork {
jobz,
uplo,
n,
eigs,
work,
rwork: None,
})
}

fn eigh_calc<'work>(
work: &'work mut EighWork<Self>,
a: &mut [Self],
) -> Result<&'work [Self::Real]> {
assert_eq!(a.len(), (work.n * work.n) as usize);
let mut info = 0;
let lwork = work.work.len() as i32;
unsafe {
$ev(
jobz,
uplo as u8,
n,
&mut a,
n,
&mut eigs,
&mut work,
lwork as i32,
$(&mut $rwork_ident,)*
work.jobz,
work.uplo as u8,
work.n,
a,
work.n,
&mut work.eigs,
&mut work.work,
lwork,
&mut info,
);
}
info.as_lapack_result()?;
Ok(eigs)
Ok(&work.eigs)
}
}
};
}

impl_eigh_work_real!(f32, lapack::ssyev);
impl_eigh_work_real!(f64, lapack::dsyev);

fn eigh_generalized(
calc_v: bool,
layout: MatrixLayout,
uplo: UPLO,
mut a: &mut [Self],
mut b: &mut [Self],
) -> Result<Vec<Self::Real>> {
macro_rules! impl_eigh_work_complex {
($scalar:ty, $ev:path) => {
impl Eigh for $scalar {
fn eigh_work(calc_v: bool, layout: MatrixLayout, uplo: UPLO) -> Result<EighWork<Self>> {
assert_eq!(layout.len(), layout.lda());
let n = layout.len();
let jobz = if calc_v { b'V' } else { b'N' };
let mut eigs = unsafe { vec_uninit(n as usize) };

$(
let mut $rwork_ident = unsafe { vec_uninit(3 * n as usize - 2) };
)*

// calc work size
let mut info = 0;
let mut work_size = [Self::zero()];
let mut rwork = unsafe { vec_uninit(3 * n as usize - 2) };
unsafe {
$evg(
&[1],
$ev(
jobz,
uplo as u8,
n,
&mut a,
n,
&mut b,
&mut [],
n,
&mut eigs,
&mut work_size,
-1,
$(&mut $rwork_ident,)*
&mut rwork,
&mut info,
);
}
info.as_lapack_result()?;

// actual evg
let lwork = work_size[0].to_usize().unwrap();
let mut work = unsafe { vec_uninit(lwork) };
let work = unsafe { vec_uninit(lwork) };
Ok(EighWork {
jobz,
uplo,
n,
eigs,
work,
rwork: Some(rwork),
})
}

fn eigh_calc<'work>(
work: &'work mut EighWork<Self>,
a: &mut [Self],
) -> Result<&'work [Self::Real]> {
assert_eq!(a.len(), (work.n * work.n) as usize);
let mut info = 0;
let lwork = work.work.len() as i32;
unsafe {
$evg(
&[1],
jobz,
uplo as u8,
n,
&mut a,
n,
&mut b,
n,
&mut eigs,
&mut work,
lwork as i32,
$(&mut $rwork_ident,)*
$ev(
work.jobz,
work.uplo as u8,
work.n,
a,
work.n,
&mut work.eigs,
&mut work.work,
lwork,
work.rwork.as_mut().unwrap(),
&mut info,
);
}
info.as_lapack_result()?;
Ok(eigs)
Ok(&work.eigs)
}
}
};
} // impl_eigh!
}

impl_eigh!(@real, f64, lapack::dsyev, lapack::dsygv);
impl_eigh!(@real, f32, lapack::ssyev, lapack::ssygv);
impl_eigh!(@complex, c64, lapack::zheev, lapack::zhegv);
impl_eigh!(@complex, c32, lapack::cheev, lapack::chegv);
impl_eigh_work_complex!(c32, lapack::cheev);
impl_eigh_work_complex!(c64, lapack::zheev);