oxiblas_ndarray/
parallel.rs

1//! Parallel BLAS operations on ndarray types.
2//!
3//! This module provides parallelized BLAS operations using Rayon,
4//! gated behind the `parallel` feature flag.
5//!
6//! # Example
7//!
8//! ```
9//! # #[cfg(feature = "parallel")] {
10//! use ndarray::array;
11//! use oxiblas_ndarray::parallel::matmul_par;
12//!
13//! // A (2x3) * B (3x2) = C (2x2)
14//! let a = array![[1.0f64, 2.0, 3.0], [4.0, 5.0, 6.0]];
15//! let b = array![[7.0f64, 8.0], [9.0, 10.0], [11.0, 12.0]];
16//!
17//! let c = matmul_par(&a, &b);
18//! assert_eq!(c, array![[58.0, 64.0], [139.0, 154.0]]);
19//! # }
20//! ```
21
22use crate::conversions::array2_to_mat;
23use ndarray::{Array2, ShapeBuilder};
24use oxiblas_blas::level3::{GemmKernel, gemm_with_par};
25use oxiblas_core::parallel::Par;
26use oxiblas_core::scalar::Field;
27use oxiblas_matrix::Mat;
28
29/// Parallel general matrix-matrix multiplication: C = alpha * A * B + beta * C
30///
31/// Uses Rayon to parallelize the GEMM computation across available threads.
32///
33/// # Arguments
34/// * `alpha` - Scalar multiplier for A * B
35/// * `a` - Left matrix (m x k)
36/// * `b` - Right matrix (k x n)
37/// * `beta` - Scalar multiplier for C
38/// * `c` - Output matrix (m x n), modified in place
39///
40/// # Panics
41/// Panics if matrix dimensions are incompatible.
42pub fn gemm_par_ndarray<T: Field + GemmKernel>(
43    alpha: T,
44    a: &Array2<T>,
45    b: &Array2<T>,
46    beta: T,
47    c: &mut Array2<T>,
48) where
49    T: bytemuck::Zeroable + Clone,
50{
51    let a_mat = array2_to_mat(a);
52    let b_mat = array2_to_mat(b);
53
54    let (m, n) = c.dim();
55    let mut c_mat: Mat<T> = Mat::zeros(m, n);
56
57    // Copy existing C values if beta != 0
58    if beta != T::zero() {
59        for i in 0..m {
60            for j in 0..n {
61                c_mat[(i, j)] = c[[i, j]];
62            }
63        }
64    }
65
66    gemm_with_par(
67        alpha,
68        a_mat.as_ref(),
69        b_mat.as_ref(),
70        beta,
71        c_mat.as_mut(),
72        Par::Rayon,
73    );
74
75    // Copy result back
76    for i in 0..m {
77        for j in 0..n {
78            c[[i, j]] = c_mat[(i, j)];
79        }
80    }
81}
82
83/// Parallel matrix multiplication: C = A * B
84///
85/// Simplified parallel version that allocates a new output matrix.
86/// Uses all available Rayon threads for computation.
87///
88/// # Arguments
89/// * `a` - Left matrix (m x k)
90/// * `b` - Right matrix (k x n)
91///
92/// # Returns
93/// New matrix C = A * B in column-major order
94///
95/// # Panics
96/// Panics if inner dimensions do not match.
97pub fn matmul_par<T: Field + GemmKernel>(a: &Array2<T>, b: &Array2<T>) -> Array2<T>
98where
99    T: bytemuck::Zeroable + Clone,
100{
101    let (m, k1) = a.dim();
102    let (k2, n) = b.dim();
103    assert_eq!(k1, k2, "Inner dimensions must match: {} vs {}", k1, k2);
104
105    let a_mat = array2_to_mat(a);
106    let b_mat = array2_to_mat(b);
107    let mut c_mat: Mat<T> = Mat::zeros(m, n);
108
109    gemm_with_par(
110        T::one(),
111        a_mat.as_ref(),
112        b_mat.as_ref(),
113        T::zero(),
114        c_mat.as_mut(),
115        Par::Rayon,
116    );
117
118    Array2::from_shape_fn((m, n).f(), |(i, j)| c_mat[(i, j)])
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use ndarray::array;
125
126    #[test]
127    fn test_matmul_par_basic() {
128        let a = Array2::from_shape_fn((2, 3), |_| 1.0f64);
129        let b = Array2::from_shape_fn((3, 2), |_| 2.0f64);
130        let c = matmul_par(&a, &b);
131
132        assert_eq!(c.dim(), (2, 2));
133        for i in 0..2 {
134            for j in 0..2 {
135                assert!((c[[i, j]] - 6.0).abs() < 1e-10);
136            }
137        }
138    }
139
140    #[test]
141    fn test_matmul_par_identity() {
142        let n = 50;
143        let a = Array2::from_shape_fn((n, n), |(i, j)| (i * n + j + 1) as f64);
144        let id = {
145            let mut m = Array2::<f64>::zeros((n, n));
146            for i in 0..n {
147                m[[i, i]] = 1.0;
148            }
149            m
150        };
151
152        let c = matmul_par(&a, &id);
153        for i in 0..n {
154            for j in 0..n {
155                assert!(
156                    (c[[i, j]] - a[[i, j]]).abs() < 1e-10,
157                    "Mismatch at ({}, {})",
158                    i,
159                    j
160                );
161            }
162        }
163    }
164
165    #[test]
166    fn test_gemm_par_ndarray_with_beta() {
167        let a = Array2::from_shape_fn((2, 3), |_| 1.0f64);
168        let b = Array2::from_shape_fn((3, 2), |_| 2.0f64);
169        let mut c = Array2::from_shape_fn((2, 2), |_| 1.0f64);
170
171        gemm_par_ndarray(1.0, &a, &b, 1.0, &mut c);
172
173        // C = 1 * A * B + 1 * C = 6 + 1 = 7
174        for i in 0..2 {
175            for j in 0..2 {
176                assert!((c[[i, j]] - 7.0).abs() < 1e-10);
177            }
178        }
179    }
180
181    #[test]
182    fn test_matmul_par_rectangular() {
183        let a = array![[1.0f64, 2.0, 3.0], [4.0, 5.0, 6.0]];
184        let b = array![[7.0f64, 8.0], [9.0, 10.0], [11.0, 12.0]];
185        let c = matmul_par(&a, &b);
186
187        assert_eq!(c.dim(), (2, 2));
188        // c[0,0] = 1*7 + 2*9 + 3*11 = 7+18+33 = 58
189        assert!((c[[0, 0]] - 58.0).abs() < 1e-10);
190        // c[0,1] = 1*8 + 2*10 + 3*12 = 8+20+36 = 64
191        assert!((c[[0, 1]] - 64.0).abs() < 1e-10);
192        // c[1,0] = 4*7 + 5*9 + 6*11 = 28+45+66 = 139
193        assert!((c[[1, 0]] - 139.0).abs() < 1e-10);
194        // c[1,1] = 4*8 + 5*10 + 6*12 = 32+50+72 = 154
195        assert!((c[[1, 1]] - 154.0).abs() < 1e-10);
196    }
197
198    #[test]
199    fn test_matmul_par_f32() {
200        let a = Array2::from_shape_fn((3, 3), |(i, j)| (i * 3 + j + 1) as f32);
201        let b = Array2::from_shape_fn((3, 3), |(i, j)| if i == j { 1.0f32 } else { 0.0f32 });
202        let c = matmul_par(&a, &b);
203
204        for i in 0..3 {
205            for j in 0..3 {
206                assert!((c[[i, j]] - a[[i, j]]).abs() < 1e-5);
207            }
208        }
209    }
210
211    #[test]
212    fn test_matmul_par_larger() {
213        let n = 100;
214        let a = Array2::from_shape_fn((n, n), |(i, j)| if i == j { 2.0f64 } else { 0.0 });
215        let b = Array2::from_shape_fn((n, n), |(i, j)| (i + j) as f64);
216        let c = matmul_par(&a, &b);
217
218        // C = 2*I * B = 2*B
219        for i in 0..n {
220            for j in 0..n {
221                let expected = 2.0 * (i + j) as f64;
222                assert!(
223                    (c[[i, j]] - expected).abs() < 1e-10,
224                    "Mismatch at ({}, {}): got {} expected {}",
225                    i,
226                    j,
227                    c[[i, j]],
228                    expected
229                );
230            }
231        }
232    }
233}
oxiblas_ndarray/parallel.rs

oxiblas_ndarray/
parallel.rs
