Impl tridiagonal by LAPACK

Author: termoshtt

lax/src/tridiagonal.rs

@@ -143,7 +143,13 @@ pub trait Tridiagonal_: Scalar + Sized {
 }
 
 macro_rules! impl_tridiagonal {
-    ($scalar:ty, $gttrf:path, $gtcon:path, $gttrs:path) => {
+    (@real, $scalar:ty, $gttrf:path, $gtcon:path, $gttrs:path) => {
+        impl_tridiagonal!(@body, $scalar, $gttrf, $gtcon, $gttrs, iwork);
+    };
+    (@complex, $scalar:ty, $gttrf:path, $gtcon:path, $gttrs:path) => {
+        impl_tridiagonal!(@body, $scalar, $gttrf, $gtcon, $gttrs, );
+    };
+    (@body, $scalar:ty, $gttrf:path, $gtcon:path, $gttrs:path, $($iwork:ident)*) => {
         impl Tridiagonal_ for $scalar {
             unsafe fn lu_tridiagonal(
                 mut a: Tridiagonal<Self>,
@@ -153,8 +159,11 @@ macro_rules! impl_tridiagonal {
                 let mut ipiv = vec![0; n as usize];
                 // We have to calc one-norm before LU factorization
                 let a_opnorm_one = a.opnorm_one();
-                $gttrf(n, &mut a.dl, &mut a.d, &mut a.du, &mut du2, &mut ipiv)
-                    .as_lapack_result()?;
+                let mut info = 0;
+                $gttrf(
+                    n, &mut a.dl, &mut a.d, &mut a.du, &mut du2, &mut ipiv, &mut info,
+                );
+                info.as_lapack_result()?;
                 Ok(LUFactorizedTridiagonal {
                     a,
                     du2,
@@ -166,7 +175,12 @@ macro_rules! impl_tridiagonal {
             unsafe fn rcond_tridiagonal(lu: &LUFactorizedTridiagonal<Self>) -> Result<Self::Real> {
                 let (n, _) = lu.a.l.size();
                 let ipiv = &lu.ipiv;
+                let mut work = vec![Self::zero(); 2 * n as usize];
+                $(
+                let mut $iwork = vec![0; n as usize];
+                )*
                 let mut rcond = Self::Real::zero();
+                let mut info = 0;
                 $gtcon(
                     NormType::One as u8,
                     n,
@@ -177,8 +191,11 @@ macro_rules! impl_tridiagonal {
                     ipiv,
                     lu.a_opnorm_one,
                     &mut rcond,
-                )
-                .as_lapack_result()?;
+                    &mut work,
+                    $(&mut $iwork,)*
+                    &mut info,
+                );
+                info.as_lapack_result()?;
                 Ok(rcond)
             }
 
@@ -192,27 +209,18 @@ macro_rules! impl_tridiagonal {
                 let (_, nrhs) = bl.size();
                 let ipiv = &lu.ipiv;
                 let ldb = bl.lda();
+                let mut info = 0;
                 $gttrs(
-                    lu.a.l.lapacke_layout(),
-                    t as u8,
-                    n,
-                    nrhs,
-                    &lu.a.dl,
-                    &lu.a.d,
-                    &lu.a.du,
-                    &lu.du2,
-                    ipiv,
-                    b,
-                    ldb,
-                )
-                .as_lapack_result()?;
+                    t as u8, n, nrhs, &lu.a.dl, &lu.a.d, &lu.a.du, &lu.du2, ipiv, b, ldb, &mut info,
+                );
+                info.as_lapack_result()?;
                 Ok(())
             }
         }
     };
 } // impl_tridiagonal!
 
-impl_tridiagonal!(f64, lapacke::dgttrf, lapacke::dgtcon, lapacke::dgttrs);
-impl_tridiagonal!(f32, lapacke::sgttrf, lapacke::sgtcon, lapacke::sgttrs);
-impl_tridiagonal!(c64, lapacke::zgttrf, lapacke::zgtcon, lapacke::zgttrs);
-impl_tridiagonal!(c32, lapacke::cgttrf, lapacke::cgtcon, lapacke::cgttrs);
+impl_tridiagonal!(@real, f64, lapack::dgttrf, lapack::dgtcon, lapack::dgttrs);
+impl_tridiagonal!(@real, f32, lapack::sgttrf, lapack::sgtcon, lapack::sgttrs);
+impl_tridiagonal!(@complex, c64, lapack::zgttrf, lapack::zgtcon, lapack::zgttrs);
+impl_tridiagonal!(@complex, c32, lapack::cgttrf, lapack::cgtcon, lapack::cgttrs);