use crate::matrix::Matrix;
#[allow(unused_imports)]
use crate::prelude::*;
use num_rational::Rational64;
use num_traits::{One, Zero};
#[derive(Debug, Clone)]
pub struct BlasConfig {
pub gemm_threshold: usize,
pub gemv_threshold: usize,
pub trsm_threshold: usize,
pub block_size: usize,
}
impl Default for BlasConfig {
fn default() -> Self {
Self {
gemm_threshold: 64 * 64 * 64, gemv_threshold: 256 * 256, trsm_threshold: 128 * 128, block_size: 64,
}
}
}
impl BlasConfig {
pub fn new(gemm_threshold: usize, gemv_threshold: usize, trsm_threshold: usize) -> Self {
Self {
gemm_threshold,
gemv_threshold,
trsm_threshold,
block_size: 64,
}
}
pub fn small_matrix() -> Self {
Self {
gemm_threshold: usize::MAX,
gemv_threshold: usize::MAX,
trsm_threshold: usize::MAX,
block_size: 32,
}
}
pub fn large_matrix() -> Self {
Self {
gemm_threshold: 0,
gemv_threshold: 0,
trsm_threshold: 0,
block_size: 128,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Transpose {
NoTrans,
Trans,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Side {
Left,
Right,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UpLo {
Upper,
Lower,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Diag {
NonUnit,
Unit,
}
pub trait BlasOps {
#[allow(clippy::too_many_arguments)]
fn gemm_into(
&self,
trans_a: Transpose,
trans_b: Transpose,
alpha: &Rational64,
a: &Matrix,
b: &Matrix,
beta: &Rational64,
c: &mut Matrix,
);
fn gemm(&self, a: &Matrix, b: &Matrix) -> Matrix {
let m = a.nrows();
let n = b.ncols();
let mut c = Matrix::zeros(m, n);
self.gemm_into(
Transpose::NoTrans,
Transpose::NoTrans,
&Rational64::from(1),
a,
b,
&Rational64::zero(),
&mut c,
);
c
}
fn gemv_into(
&self,
trans: Transpose,
alpha: &Rational64,
a: &Matrix,
x: &[Rational64],
beta: &Rational64,
y: &mut [Rational64],
);
fn gemv(&self, a: &Matrix, x: &[Rational64]) -> Vec<Rational64> {
let m = a.nrows();
let mut y = vec![Rational64::zero(); m];
self.gemv_into(
Transpose::NoTrans,
&Rational64::from(1),
a,
x,
&Rational64::zero(),
&mut y,
);
y
}
#[allow(clippy::too_many_arguments)]
fn trsm(
&self,
side: Side,
uplo: UpLo,
trans: Transpose,
diag: Diag,
alpha: &Rational64,
a: &Matrix,
b: &mut Matrix,
);
fn trsv(&self, uplo: UpLo, trans: Transpose, diag: Diag, a: &Matrix, b: &mut [Rational64]);
fn config(&self) -> &BlasConfig;
fn should_use_blas_gemm(&self, m: usize, n: usize, k: usize) -> bool {
m * n * k >= self.config().gemm_threshold
}
fn should_use_blas_gemv(&self, m: usize, n: usize) -> bool {
m * n >= self.config().gemv_threshold
}
fn should_use_blas_trsm(&self, n: usize, nrhs: usize) -> bool {
n * nrhs >= self.config().trsm_threshold
}
}
#[derive(Debug, Clone)]
pub struct RustBlas {
config: BlasConfig,
}
impl Default for RustBlas {
fn default() -> Self {
Self::new()
}
}
impl RustBlas {
pub fn new() -> Self {
Self {
config: BlasConfig::default(),
}
}
pub fn with_config(config: BlasConfig) -> Self {
Self { config }
}
fn get_dims(m: &Matrix, trans: Transpose) -> (usize, usize) {
match trans {
Transpose::NoTrans => (m.nrows(), m.ncols()),
Transpose::Trans => (m.ncols(), m.nrows()),
}
}
fn get_elem(m: &Matrix, i: usize, j: usize, trans: Transpose) -> Rational64 {
match trans {
Transpose::NoTrans => m.get(i, j),
Transpose::Trans => m.get(j, i),
}
}
}
#[allow(clippy::needless_range_loop)]
impl BlasOps for RustBlas {
fn gemm_into(
&self,
trans_a: Transpose,
trans_b: Transpose,
alpha: &Rational64,
a: &Matrix,
b: &Matrix,
beta: &Rational64,
c: &mut Matrix,
) {
let (m, k_a) = Self::get_dims(a, trans_a);
let (k_b, n) = Self::get_dims(b, trans_b);
assert_eq!(
k_a, k_b,
"Inner dimensions must match for matrix multiplication"
);
assert_eq!(c.nrows(), m, "Output matrix row dimension mismatch");
assert_eq!(c.ncols(), n, "Output matrix column dimension mismatch");
let k = k_a;
for i in 0..m {
for j in 0..n {
let mut sum = Rational64::zero();
for l in 0..k {
let a_il = Self::get_elem(a, i, l, trans_a);
let b_lj = Self::get_elem(b, l, j, trans_b);
sum += a_il * b_lj;
}
let c_ij = c.get(i, j);
c.set(i, j, alpha * sum + beta * c_ij);
}
}
}
fn gemv_into(
&self,
trans: Transpose,
alpha: &Rational64,
a: &Matrix,
x: &[Rational64],
beta: &Rational64,
y: &mut [Rational64],
) {
let (m, n) = Self::get_dims(a, trans);
assert_eq!(x.len(), n, "Vector x dimension must match matrix columns");
assert_eq!(y.len(), m, "Vector y dimension must match matrix rows");
for i in 0..m {
let mut sum = Rational64::zero();
for j in 0..n {
let a_ij = Self::get_elem(a, i, j, trans);
sum += a_ij * x[j];
}
y[i] = alpha * sum + beta * y[i];
}
}
fn trsm(
&self,
side: Side,
uplo: UpLo,
trans: Transpose,
diag: Diag,
alpha: &Rational64,
a: &Matrix,
b: &mut Matrix,
) {
let n = a.nrows();
assert_eq!(a.nrows(), a.ncols(), "Triangular matrix must be square");
match side {
Side::Left => {
assert_eq!(b.nrows(), n, "B row dimension must match A");
self.trsm_left(uplo, trans, diag, alpha, a, b);
}
Side::Right => {
assert_eq!(b.ncols(), n, "B column dimension must match A");
self.trsm_right(uplo, trans, diag, alpha, a, b);
}
}
}
fn trsv(&self, uplo: UpLo, trans: Transpose, diag: Diag, a: &Matrix, b: &mut [Rational64]) {
let n = a.nrows();
assert_eq!(a.nrows(), a.ncols(), "Triangular matrix must be square");
assert_eq!(b.len(), n, "Vector dimension must match matrix");
match (uplo, trans) {
(UpLo::Lower, Transpose::NoTrans) | (UpLo::Upper, Transpose::Trans) => {
for i in 0..n {
let mut sum = b[i];
for j in 0..i {
let a_ij = Self::get_elem(a, i, j, trans);
sum -= a_ij * b[j];
}
if diag == Diag::NonUnit {
let a_ii = Self::get_elem(a, i, i, trans);
b[i] = sum / a_ii;
} else {
b[i] = sum;
}
}
}
(UpLo::Upper, Transpose::NoTrans) | (UpLo::Lower, Transpose::Trans) => {
for i in (0..n).rev() {
let mut sum = b[i];
for j in (i + 1)..n {
let a_ij = Self::get_elem(a, i, j, trans);
sum -= a_ij * b[j];
}
if diag == Diag::NonUnit {
let a_ii = Self::get_elem(a, i, i, trans);
b[i] = sum / a_ii;
} else {
b[i] = sum;
}
}
}
}
}
fn config(&self) -> &BlasConfig {
&self.config
}
}
impl RustBlas {
fn trsm_left(
&self,
uplo: UpLo,
trans: Transpose,
diag: Diag,
alpha: &Rational64,
a: &Matrix,
b: &mut Matrix,
) {
let n = a.nrows();
let nrhs = b.ncols();
if !alpha.is_one() {
for i in 0..n {
for j in 0..nrhs {
let val = b.get(i, j) * alpha;
b.set(i, j, val);
}
}
}
for col in 0..nrhs {
match (uplo, trans) {
(UpLo::Lower, Transpose::NoTrans) | (UpLo::Upper, Transpose::Trans) => {
for i in 0..n {
let mut sum = b.get(i, col);
for j in 0..i {
let a_ij = Self::get_elem(a, i, j, trans);
sum -= a_ij * b.get(j, col);
}
if diag == Diag::NonUnit {
let a_ii = Self::get_elem(a, i, i, trans);
b.set(i, col, sum / a_ii);
} else {
b.set(i, col, sum);
}
}
}
(UpLo::Upper, Transpose::NoTrans) | (UpLo::Lower, Transpose::Trans) => {
for i in (0..n).rev() {
let mut sum = b.get(i, col);
for j in (i + 1)..n {
let a_ij = Self::get_elem(a, i, j, trans);
sum -= a_ij * b.get(j, col);
}
if diag == Diag::NonUnit {
let a_ii = Self::get_elem(a, i, i, trans);
b.set(i, col, sum / a_ii);
} else {
b.set(i, col, sum);
}
}
}
}
}
}
fn trsm_right(
&self,
uplo: UpLo,
trans: Transpose,
diag: Diag,
alpha: &Rational64,
a: &Matrix,
b: &mut Matrix,
) {
let n = a.nrows();
let m = b.nrows();
if !alpha.is_one() {
for i in 0..m {
for j in 0..n {
let val = b.get(i, j) * alpha;
b.set(i, j, val);
}
}
}
for row in 0..m {
match (uplo, trans) {
(UpLo::Upper, Transpose::NoTrans) | (UpLo::Lower, Transpose::Trans) => {
for j in 0..n {
let mut sum = b.get(row, j);
for k in 0..j {
let a_kj = Self::get_elem(a, k, j, trans);
sum -= b.get(row, k) * a_kj;
}
if diag == Diag::NonUnit {
let a_jj = Self::get_elem(a, j, j, trans);
b.set(row, j, sum / a_jj);
} else {
b.set(row, j, sum);
}
}
}
(UpLo::Lower, Transpose::NoTrans) | (UpLo::Upper, Transpose::Trans) => {
for j in (0..n).rev() {
let mut sum = b.get(row, j);
for k in (j + 1)..n {
let a_kj = Self::get_elem(a, k, j, trans);
sum -= b.get(row, k) * a_kj;
}
if diag == Diag::NonUnit {
let a_jj = Self::get_elem(a, j, j, trans);
b.set(row, j, sum / a_jj);
} else {
b.set(row, j, sum);
}
}
}
}
}
}
}
pub fn get_blas() -> RustBlas {
RustBlas::new()
}
pub fn get_blas_with_config(config: BlasConfig) -> RustBlas {
RustBlas::with_config(config)
}
#[cfg(test)]
mod tests {
use super::*;
fn rat(n: i64) -> Rational64 {
Rational64::from(n)
}
#[test]
fn test_gemm_basic() {
let blas = RustBlas::new();
let a = Matrix::from_rows(vec![vec![rat(1), rat(2)], vec![rat(3), rat(4)]]);
let b = Matrix::from_rows(vec![vec![rat(5), rat(6)], vec![rat(7), rat(8)]]);
let c = blas.gemm(&a, &b);
assert_eq!(c.get(0, 0), rat(19));
assert_eq!(c.get(0, 1), rat(22));
assert_eq!(c.get(1, 0), rat(43));
assert_eq!(c.get(1, 1), rat(50));
}
#[test]
fn test_gemm_with_transpose() {
let blas = RustBlas::new();
let a = Matrix::from_rows(vec![vec![rat(1), rat(3)], vec![rat(2), rat(4)]]);
let b = Matrix::from_rows(vec![vec![rat(5), rat(7)], vec![rat(6), rat(8)]]);
let mut c = Matrix::zeros(2, 2);
blas.gemm_into(
Transpose::Trans,
Transpose::Trans,
&rat(1),
&a,
&b,
&rat(0),
&mut c,
);
assert_eq!(c.get(0, 0), rat(19));
assert_eq!(c.get(0, 1), rat(22));
assert_eq!(c.get(1, 0), rat(43));
assert_eq!(c.get(1, 1), rat(50));
}
#[test]
fn test_gemm_with_alpha_beta() {
let blas = RustBlas::new();
let a = Matrix::from_rows(vec![vec![rat(1), rat(2)], vec![rat(3), rat(4)]]);
let b = Matrix::from_rows(vec![vec![rat(1), rat(0)], vec![rat(0), rat(1)]]);
let mut c = Matrix::from_rows(vec![vec![rat(10), rat(20)], vec![rat(30), rat(40)]]);
blas.gemm_into(
Transpose::NoTrans,
Transpose::NoTrans,
&rat(2),
&a,
&b,
&rat(3),
&mut c,
);
assert_eq!(c.get(0, 0), rat(32));
assert_eq!(c.get(0, 1), rat(64));
assert_eq!(c.get(1, 0), rat(96));
assert_eq!(c.get(1, 1), rat(128));
}
#[test]
fn test_gemv_basic() {
let blas = RustBlas::new();
let a = Matrix::from_rows(vec![vec![rat(1), rat(2)], vec![rat(3), rat(4)]]);
let x = vec![rat(5), rat(6)];
let y = blas.gemv(&a, &x);
assert_eq!(y[0], rat(17));
assert_eq!(y[1], rat(39));
}
#[test]
fn test_gemv_with_transpose() {
let blas = RustBlas::new();
let a = Matrix::from_rows(vec![vec![rat(1), rat(3)], vec![rat(2), rat(4)]]);
let x = vec![rat(5), rat(6)];
let mut y = vec![rat(0), rat(0)];
blas.gemv_into(Transpose::Trans, &rat(1), &a, &x, &rat(0), &mut y);
assert_eq!(y[0], rat(17));
assert_eq!(y[1], rat(39));
}
#[test]
fn test_trsv_lower() {
let blas = RustBlas::new();
let l = Matrix::from_rows(vec![
vec![rat(2), rat(0), rat(0)],
vec![rat(1), rat(3), rat(0)],
vec![rat(1), rat(2), rat(4)],
]);
let mut b = vec![rat(4), rat(10), rat(21)];
blas.trsv(UpLo::Lower, Transpose::NoTrans, Diag::NonUnit, &l, &mut b);
assert_eq!(b[0], rat(2));
assert_eq!(b[1], Rational64::new(8, 3));
assert_eq!(b[2], Rational64::new(41, 12));
}
#[test]
fn test_trsv_upper() {
let blas = RustBlas::new();
let u = Matrix::from_rows(vec![
vec![rat(2), rat(1), rat(1)],
vec![rat(0), rat(3), rat(2)],
vec![rat(0), rat(0), rat(4)],
]);
let mut b = vec![rat(8), rat(14), rat(12)];
blas.trsv(UpLo::Upper, Transpose::NoTrans, Diag::NonUnit, &u, &mut b);
assert_eq!(b[2], rat(3));
assert_eq!(b[1], Rational64::new(8, 3));
assert_eq!(b[0], Rational64::new(7, 6));
}
#[test]
fn test_trsv_unit_diagonal() {
let blas = RustBlas::new();
let l = Matrix::from_rows(vec![
vec![rat(999), rat(0)], vec![rat(2), rat(888)], ]);
let mut b = vec![rat(4), rat(10)];
blas.trsv(UpLo::Lower, Transpose::NoTrans, Diag::Unit, &l, &mut b);
assert_eq!(b[0], rat(4));
assert_eq!(b[1], rat(2));
}
#[test]
fn test_trsm_left_lower() {
let blas = RustBlas::new();
let l = Matrix::from_rows(vec![vec![rat(2), rat(0)], vec![rat(1), rat(3)]]);
let mut b = Matrix::from_rows(vec![vec![rat(4), rat(6)], vec![rat(7), rat(12)]]);
blas.trsm(
Side::Left,
UpLo::Lower,
Transpose::NoTrans,
Diag::NonUnit,
&rat(1),
&l,
&mut b,
);
assert_eq!(b.get(0, 0), rat(2));
assert_eq!(b.get(1, 0), Rational64::new(5, 3));
assert_eq!(b.get(0, 1), rat(3));
assert_eq!(b.get(1, 1), rat(3));
}
#[test]
fn test_config_thresholds() {
let config = BlasConfig::new(100, 50, 25);
let blas = RustBlas::with_config(config);
assert!(!blas.should_use_blas_gemm(3, 3, 3)); assert!(blas.should_use_blas_gemm(5, 5, 5));
assert!(!blas.should_use_blas_gemv(5, 5)); assert!(blas.should_use_blas_gemv(8, 8));
assert!(!blas.should_use_blas_trsm(4, 4)); assert!(blas.should_use_blas_trsm(5, 5)); }
#[test]
fn test_config_presets() {
let small = BlasConfig::small_matrix();
assert_eq!(small.gemm_threshold, usize::MAX);
let large = BlasConfig::large_matrix();
assert_eq!(large.gemm_threshold, 0);
}
#[test]
fn test_identity_multiplication() {
let blas = RustBlas::new();
let a = Matrix::from_rows(vec![
vec![rat(1), rat(2), rat(3)],
vec![rat(4), rat(5), rat(6)],
vec![rat(7), rat(8), rat(9)],
]);
let i = Matrix::identity(3);
let result = blas.gemm(&a, &i);
for row in 0..3 {
for col in 0..3 {
assert_eq!(result.get(row, col), a.get(row, col));
}
}
let result2 = blas.gemm(&i, &a);
for row in 0..3 {
for col in 0..3 {
assert_eq!(result2.get(row, col), a.get(row, col));
}
}
}
#[test]
fn test_gemm_non_square() {
let blas = RustBlas::new();
let a = Matrix::from_rows(vec![
vec![rat(1), rat(2), rat(3)],
vec![rat(4), rat(5), rat(6)],
]);
let b = Matrix::from_rows(vec![
vec![rat(7), rat(8)],
vec![rat(9), rat(10)],
vec![rat(11), rat(12)],
]);
let c = blas.gemm(&a, &b);
assert_eq!(c.nrows(), 2);
assert_eq!(c.ncols(), 2);
assert_eq!(c.get(0, 0), rat(58));
assert_eq!(c.get(0, 1), rat(64));
assert_eq!(c.get(1, 0), rat(139));
assert_eq!(c.get(1, 1), rat(154));
}
#[test]
fn test_get_blas_functions() {
let blas1 = get_blas();
assert_eq!(
blas1.config().gemm_threshold,
BlasConfig::default().gemm_threshold
);
let config = BlasConfig::small_matrix();
let blas2 = get_blas_with_config(config);
assert_eq!(blas2.config().gemm_threshold, usize::MAX);
}
}