use std::{cell::RefCell, mem::MaybeUninit};
use crate::context::broadcast_batch;
use crate::{
error::LaError, linear_solver_error, Context, CudaContext, CudaMat, CudaVec, LinearOp,
LinearSolver, Matrix, ScalarCuda,
};
use cudarc::{
cusolver::sys::{
cublasOperation_t, cusolverDnCreate, cusolverDnDestroy, cusolverDnDgetrf,
cusolverDnDgetrf_bufferSize, cusolverDnDgetrs, cusolverDnHandle_t, cusolverStatus_t,
},
driver::{CudaSlice, DevicePtr, DevicePtrMut},
};
#[deprecated(
since = "0.2.1",
note = "the `cuda` backend is deprecated, use the `cuda-oxide` backend (`OxideContext`, `OxideVec`, `OxideMat`, `OxideLU`) instead"
)]
pub struct CudaLU<T>
where
T: ScalarCuda,
{
work: Option<CudaSlice<T>>,
pivots: Option<CudaSlice<i32>>,
nfo: Option<RefCell<CudaSlice<i32>>>,
matrix: Option<CudaMat<T>>,
handle: cusolverDnHandle_t,
linearisation_set: bool,
}
impl<T> Default for CudaLU<T>
where
T: ScalarCuda,
{
fn default() -> Self {
let handle = {
let mut handle = MaybeUninit::uninit();
unsafe {
let stat = cusolverDnCreate(handle.as_mut_ptr());
assert_eq!(stat, cusolverStatus_t::CUSOLVER_STATUS_SUCCESS);
handle.assume_init()
}
};
Self {
matrix: None,
work: None,
pivots: None,
nfo: None,
handle,
linearisation_set: false,
}
}
}
impl<T: ScalarCuda> Drop for CudaLU<T> {
fn drop(&mut self) {
unsafe {
cusolverDnDestroy(self.handle);
}
}
}
impl<T: ScalarCuda> CudaLU<T> {
fn solve_impl(&self, x: &mut CudaVec<T>, transpose: bool) -> Result<(), LaError> {
let matrix = if let Some(ref matrix) = self.matrix {
if matrix.nrows() != matrix.ncols() {
Err(linear_solver_error!(LinearSolverMatrixNotSquare))?;
}
matrix
} else {
Err(linear_solver_error!(LinearSolverNotSetup))?
};
if !self.linearisation_set {
Err(linear_solver_error!(LinearSolverNotSetup))?;
}
let nbatch = x.context.nbatch();
let lu_nbatch = matrix.context().nbatch();
x.context.assert_broadcastable_into(lu_nbatch, "lu_solve");
let nrows = matrix.nrows();
let ncols = matrix.ncols();
let x_nstates = x.data.len() / nbatch;
if x_nstates != nrows {
Err(linear_solver_error!(LinearSolverMatrixVectorNotCompatible))?;
}
let mut nfo = self.nfo.as_ref().expect("NFO not set").borrow_mut();
let stream = &x.context.stream;
let lda = i32::try_from(nrows).unwrap();
let n = i32::try_from(nrows).unwrap();
let nrhs = 1i32;
let (a_ptr, _) = matrix.data.device_ptr(stream);
let (p_ptr, _) = self.pivots.as_ref().unwrap().device_ptr(stream);
let (x_ptr, _) = x.data.device_ptr_mut(stream);
let (n_ptr, _) = nfo.device_ptr_mut(stream);
for b in 0..nbatch {
let lu_b = broadcast_batch(b, lu_nbatch, nbatch);
let a_offset = lu_b * nrows * ncols;
let p_offset = lu_b * nrows;
let x_offset = b * x_nstates;
unsafe {
cusolverDnDgetrs(
self.handle,
if transpose {
cublasOperation_t::CUBLAS_OP_T
} else {
cublasOperation_t::CUBLAS_OP_N
},
n,
nrhs,
(a_ptr as *const f64).add(a_offset),
lda,
(p_ptr as *const i32).add(p_offset),
(x_ptr as *mut f64).add(x_offset),
n,
(n_ptr as *mut i32).add(lu_b),
)
};
}
Ok(())
}
}
impl<T: ScalarCuda> LinearSolver<CudaMat<T>> for CudaLU<T> {
fn set_linearisation<C: LinearOp<T = T, V = CudaVec<T>, M = CudaMat<T>, C = CudaContext>>(
&mut self,
op: &C,
) {
let matrix = self.matrix.as_mut().expect("Matrix not set");
let work = self.work.as_mut().expect("Work space not set");
let pivots = self.pivots.as_mut().expect("Pivots not set");
let nfo = self.nfo.as_mut().expect("NFO not set").get_mut();
op.matrix_inplace(matrix);
let nbatch = op.context().nbatch();
let nrows = matrix.nrows();
let ncols = matrix.ncols();
let stream = &op.context().stream;
let m = i32::try_from(nrows).unwrap();
let n = i32::try_from(ncols).unwrap();
let lda = i32::try_from(nrows).unwrap();
let (a_ptr, _) = matrix.data.device_ptr_mut(stream);
let (ws_ptr, _) = work.device_ptr_mut(stream);
let (p_ptr, _) = pivots.device_ptr_mut(stream);
let (n_ptr, _) = nfo.device_ptr_mut(stream);
for b in 0..nbatch {
let a_offset = b * nrows * ncols;
let p_offset = b * nrows;
unsafe {
cusolverDnDgetrf(
self.handle,
m,
n,
(a_ptr as *mut f64).add(a_offset),
lda,
ws_ptr as *mut f64,
(p_ptr as *mut i32).add(p_offset),
(n_ptr as *mut i32).add(b),
)
};
}
self.linearisation_set = true;
}
fn solve_in_place(&self, x: &mut CudaVec<T>) -> Result<(), LaError> {
self.solve_impl(x, false)
}
fn solve_transpose_in_place(&self, x: &mut CudaVec<T>) -> Result<(), LaError> {
self.solve_impl(x, true)
}
fn set_sparsity<C: LinearOp<T = T, V = CudaVec<T>, M = CudaMat<T>, C = CudaContext>>(
&mut self,
op: &C,
) {
crate::context::cuda::assert_f64::<T>();
let ncols = op.ncols();
let nrows = op.nrows();
let nbatch = op.context().nbatch();
let matrix = C::M::new_from_sparsity(nrows, ncols, op.sparsity(), op.context().clone());
self.matrix = Some(matrix);
let stream = &op.context().stream;
let lwork = {
let mut lwork = 0;
let (a, _syn) = self.matrix.as_mut().unwrap().data.device_ptr_mut(stream);
let m = i32::try_from(nrows).unwrap();
let n = i32::try_from(ncols).unwrap();
let lda = i32::try_from(nrows).unwrap();
unsafe {
cusolverDnDgetrf_bufferSize(self.handle, m, n, a as *mut f64, lda, &mut lwork);
}
lwork
};
unsafe {
self.work = Some(
stream
.alloc(lwork as usize)
.expect("Failed to allocate work space"),
);
self.pivots = Some(
stream
.alloc(nrows * nbatch)
.expect("Failed to allocate pivots"),
);
self.nfo = Some(RefCell::new(
stream.alloc(nbatch).expect("Failed to allocate NFO"),
));
}
self.linearisation_set = false;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
linear_solver::tests::{
diagonal_op, test_grouped_lu_solve, test_narrow_state_lu_solve, test_transpose_lu_solve,
},
Vector,
};
#[test]
fn test_lu() {
let mut s = CudaLU::<f64>::default();
let op = diagonal_op::<CudaMat<f64>>(2.0);
s.set_sparsity(&op);
s.set_linearisation(&op);
let b = CudaVec::from_vec(vec![2.0, 4.0], Default::default());
let x = s.solve(&b).unwrap();
x.assert_eq_st(
&CudaVec::from_vec(vec![1.0, 2.0], Default::default()),
1e-10,
);
}
#[test]
#[should_panic(expected = "incompatible nbatch")]
fn test_narrow_state_lu() {
test_narrow_state_lu_solve::<CudaMat<f64>, CudaLU<f64>>(
CudaContext::default().with_nbatch(2),
);
}
#[test]
fn test_grouped_lu() {
test_grouped_lu_solve::<CudaMat<f64>, CudaLU<f64>>(CudaContext::default().with_nbatch(2));
}
#[test]
fn test_transpose_lu() {
test_transpose_lu_solve::<CudaMat<f64>, CudaLU<f64>>(CudaContext::default().with_nbatch(2));
}
}