sparse_ir_capi/
gemm.rs

1//! GEMM (Matrix Multiplication) C-API
2//!
3//! This module provides C-API functions for registering external BLAS implementations.
4//! These functions allow users to inject their own BLAS libraries (OpenBLAS, MKL, Accelerate, etc.)
5//! at runtime without recompiling.
6//!
7//! # API Functions
8//! - `spir_gemm_backend_new_from_fblas_lp64`: Create backend from LP64 BLAS (32-bit integers)
9//! - `spir_gemm_backend_new_from_fblas_ilp64`: Create backend from ILP64 BLAS (64-bit integers)
10//! - `spir_gemm_backend_release`: Release backend handle
11//!
12//! # Example (C)
13//! ```c
14//! // Link against BLAS library (e.g., OpenBLAS, MKL, Accelerate)
15//! // Fortran BLAS functions typically have trailing underscore
16//!
17//! // Create backend from Fortran BLAS (or pass NULL to use default backend)
18//! spir_gemm_backend* backend = spir_gemm_backend_new_from_fblas_lp64(
19//!     (void*)dgemm_,
20//!     (void*)zgemm_
21//! );
22//!
23//! // Use backend in evaluate/fit functions (pass NULL to use default backend)
24//! // ...
25//!
26//! // Release backend when done
27//! spir_gemm_backend_release(backend);
28//! ```
29
30use sparse_ir::gemm::{
31    Dgemm64FnPtr, DgemmFnPtr, ExternalBlas64Backend, ExternalBlasBackend, GemmBackendHandle,
32    Zgemm64FnPtr, ZgemmFnPtr,
33};
34
35//==============================================================================
36// Backend Handle C-API
37//==============================================================================
38
39/// Opaque pointer type for GEMM backend handle
40///
41/// This type wraps a `GemmBackendHandle` and provides a C-compatible interface.
42/// The handle can be created, cloned, and passed to evaluate/fit functions.
43///
44/// Note: The internal structure is hidden using a void pointer to prevent exposing GemmBackendHandle to C.
45#[repr(C)]
46pub struct spir_gemm_backend {
47    pub(crate) _private: *const std::ffi::c_void,
48}
49
50impl spir_gemm_backend {
51    /// Get a reference to the inner GemmBackendHandle
52    pub(crate) fn inner(&self) -> &GemmBackendHandle {
53        unsafe { &*(self._private as *const GemmBackendHandle) }
54    }
55
56    pub(crate) fn new(handle: GemmBackendHandle) -> Self {
57        Self {
58            _private: Box::into_raw(Box::new(handle)) as *const std::ffi::c_void,
59        }
60    }
61}
62
63impl Drop for spir_gemm_backend {
64    fn drop(&mut self) {
65        if !self._private.is_null() {
66            unsafe {
67                let _ = Box::from_raw(
68                    self._private as *const GemmBackendHandle as *mut GemmBackendHandle,
69                );
70            }
71        }
72    }
73}
74
75impl Clone for spir_gemm_backend {
76    fn clone(&self) -> Self {
77        // GemmBackendHandle is already Arc-based, so we can clone the inner handle
78        let inner = self.inner().clone();
79        Self::new(inner)
80    }
81}
82
83/// Create GEMM backend from Fortran BLAS function pointers (LP64)
84///
85/// Creates a new backend handle from Fortran BLAS function pointers.
86///
87/// # Arguments
88/// * `dgemm` - Function pointer to Fortran BLAS dgemm (double precision)
89/// * `zgemm` - Function pointer to Fortran BLAS zgemm (complex double precision)
90///
91/// # Returns
92/// * Pointer to `spir_gemm_backend` on success
93/// * `NULL` if function pointers are null
94///
95/// # Safety
96/// The provided function pointers must:
97/// - Be valid Fortran BLAS function pointers following the standard Fortran BLAS interface
98/// - Use 32-bit integers for all dimension parameters (LP64 interface)
99/// - Be thread-safe (will be called from multiple threads)
100/// - Remain valid for the entire lifetime of the backend handle
101///
102/// The returned pointer must be freed with `spir_gemm_backend_free` when no longer needed.
103#[unsafe(no_mangle)]
104pub extern "C" fn spir_gemm_backend_new_from_fblas_lp64(
105    dgemm: *const libc::c_void,
106    zgemm: *const libc::c_void,
107) -> *mut spir_gemm_backend {
108    // Validate input
109    if dgemm.is_null() || zgemm.is_null() {
110        return std::ptr::null_mut();
111    }
112
113    // Catch panics (e.g. OOM from Box::new) to prevent unwinding across FFI boundary
114    let result = std::panic::catch_unwind(|| {
115        // Cast to Fortran BLAS function pointer types
116        let dgemm_fn: DgemmFnPtr = unsafe { std::mem::transmute(dgemm) };
117        let zgemm_fn: ZgemmFnPtr = unsafe { std::mem::transmute(zgemm) };
118
119        // Create backend
120        let backend = ExternalBlasBackend::new(dgemm_fn, zgemm_fn);
121
122        // Wrap in handle
123        let handle = GemmBackendHandle::new(Box::new(backend));
124        Box::into_raw(Box::new(spir_gemm_backend::new(handle)))
125    });
126
127    result.unwrap_or(std::ptr::null_mut())
128}
129
130/// Create GEMM backend from Fortran BLAS function pointers (ILP64)
131///
132/// Creates a new backend handle from Fortran BLAS function pointers with 64-bit integers.
133///
134/// # Arguments
135/// * `dgemm64` - Function pointer to Fortran BLAS dgemm (double precision, 64-bit integers)
136/// * `zgemm64` - Function pointer to Fortran BLAS zgemm (complex double precision, 64-bit integers)
137///
138/// # Returns
139/// * Pointer to `spir_gemm_backend` on success
140/// * `NULL` if function pointers are null
141///
142/// # Safety
143/// The provided function pointers must:
144/// - Be valid Fortran BLAS function pointers following the standard Fortran BLAS interface
145/// - Use 64-bit integers for all dimension parameters (ILP64 interface)
146/// - Be thread-safe (will be called from multiple threads)
147/// - Remain valid for the entire lifetime of the backend handle
148///
149/// The returned pointer must be freed with `spir_gemm_backend_free` when no longer needed.
150#[unsafe(no_mangle)]
151pub extern "C" fn spir_gemm_backend_new_from_fblas_ilp64(
152    dgemm64: *const libc::c_void,
153    zgemm64: *const libc::c_void,
154) -> *mut spir_gemm_backend {
155    // Validate input
156    if dgemm64.is_null() || zgemm64.is_null() {
157        return std::ptr::null_mut();
158    }
159
160    // Catch panics (e.g. OOM from Box::new) to prevent unwinding across FFI boundary
161    let result = std::panic::catch_unwind(|| {
162        // Cast to Fortran BLAS function pointer types
163        let dgemm64_fn: Dgemm64FnPtr = unsafe { std::mem::transmute(dgemm64) };
164        let zgemm64_fn: Zgemm64FnPtr = unsafe { std::mem::transmute(zgemm64) };
165
166        // Create backend
167        let backend = ExternalBlas64Backend::new(dgemm64_fn, zgemm64_fn);
168
169        // Wrap in handle
170        let handle = GemmBackendHandle::new(Box::new(backend));
171        Box::into_raw(Box::new(spir_gemm_backend::new(handle)))
172    });
173
174    result.unwrap_or(std::ptr::null_mut())
175}
176
177/// Release GEMM backend handle
178///
179/// Releases the memory associated with a backend handle.
180///
181/// # Arguments
182/// * `backend` - Pointer to backend handle (can be NULL)
183///
184/// # Safety
185/// The pointer must have been created by `spir_gemm_backend_new_from_fblas_lp64` or
186/// `spir_gemm_backend_new_from_fblas_ilp64`.
187/// After calling this function, the pointer must not be used again.
188#[unsafe(no_mangle)]
189pub extern "C" fn spir_gemm_backend_release(backend: *mut spir_gemm_backend) {
190    if !backend.is_null() {
191        unsafe {
192            let _ = Box::from_raw(backend);
193        }
194    }
195}
196
197/// Get backend handle from opaque pointer (internal use)
198///
199/// # Safety
200/// The pointer must be valid and not null.
201/// The returned reference is only valid while the backend pointer is valid.
202/// The caller must ensure the backend pointer remains valid for the lifetime of the returned reference.
203pub(crate) unsafe fn get_backend_handle<'a>(
204    backend: *const spir_gemm_backend,
205) -> Option<&'a GemmBackendHandle> {
206    if backend.is_null() {
207        None
208    } else {
209        unsafe { Some((*backend).inner()) }
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216
217    // Mock Fortran BLAS functions for testing
218    unsafe extern "C" fn mock_dgemm(
219        _transa: *const libc::c_char,
220        _transb: *const libc::c_char,
221        _m: *const libc::c_int,
222        _n: *const libc::c_int,
223        _k: *const libc::c_int,
224        _alpha: *const libc::c_double,
225        _a: *const libc::c_double,
226        _lda: *const libc::c_int,
227        _b: *const libc::c_double,
228        _ldb: *const libc::c_int,
229        _beta: *const libc::c_double,
230        _c: *mut libc::c_double,
231        _ldc: *const libc::c_int,
232    ) {
233        // Mock implementation - does nothing
234    }
235
236    unsafe extern "C" fn mock_zgemm(
237        _transa: *const libc::c_char,
238        _transb: *const libc::c_char,
239        _m: *const libc::c_int,
240        _n: *const libc::c_int,
241        _k: *const libc::c_int,
242        _alpha: *const num_complex::Complex<f64>,
243        _a: *const num_complex::Complex<f64>,
244        _lda: *const libc::c_int,
245        _b: *const num_complex::Complex<f64>,
246        _ldb: *const libc::c_int,
247        _beta: *const num_complex::Complex<f64>,
248        _c: *mut num_complex::Complex<f64>,
249        _ldc: *const libc::c_int,
250    ) {
251        // Mock implementation - does nothing
252    }
253
254    unsafe extern "C" fn mock_dgemm64(
255        _transa: *const libc::c_char,
256        _transb: *const libc::c_char,
257        _m: *const i64,
258        _n: *const i64,
259        _k: *const i64,
260        _alpha: *const libc::c_double,
261        _a: *const libc::c_double,
262        _lda: *const i64,
263        _b: *const libc::c_double,
264        _ldb: *const i64,
265        _beta: *const libc::c_double,
266        _c: *mut libc::c_double,
267        _ldc: *const i64,
268    ) {
269        // Mock implementation - does nothing
270    }
271
272    unsafe extern "C" fn mock_zgemm64(
273        _transa: *const libc::c_char,
274        _transb: *const libc::c_char,
275        _m: *const i64,
276        _n: *const i64,
277        _k: *const i64,
278        _alpha: *const num_complex::Complex<f64>,
279        _a: *const num_complex::Complex<f64>,
280        _lda: *const i64,
281        _b: *const num_complex::Complex<f64>,
282        _ldb: *const i64,
283        _beta: *const num_complex::Complex<f64>,
284        _c: *mut num_complex::Complex<f64>,
285        _ldc: *const i64,
286    ) {
287        // Mock implementation - does nothing
288    }
289
290    #[test]
291    fn test_backend_new_from_fblas_lp64_success() {
292        unsafe {
293            let backend = spir_gemm_backend_new_from_fblas_lp64(
294                mock_dgemm as *const _,
295                mock_zgemm as *const _,
296            );
297            assert!(!backend.is_null(), "Backend should not be null");
298            spir_gemm_backend_release(backend);
299        }
300    }
301
302    #[test]
303    fn test_backend_new_from_fblas_ilp64_success() {
304        unsafe {
305            let backend = spir_gemm_backend_new_from_fblas_ilp64(
306                mock_dgemm64 as *const _,
307                mock_zgemm64 as *const _,
308            );
309            assert!(!backend.is_null(), "Backend should not be null");
310            spir_gemm_backend_release(backend);
311        }
312    }
313
314    #[test]
315    fn test_backend_new_from_fblas_lp64_null_dgemm() {
316        unsafe {
317            let backend =
318                spir_gemm_backend_new_from_fblas_lp64(std::ptr::null(), mock_zgemm as *const _);
319            assert!(
320                backend.is_null(),
321                "Backend should be null when dgemm is null"
322            );
323        }
324    }
325
326    #[test]
327    fn test_backend_new_from_fblas_lp64_null_zgemm() {
328        unsafe {
329            let backend =
330                spir_gemm_backend_new_from_fblas_lp64(mock_dgemm as *const _, std::ptr::null());
331            assert!(
332                backend.is_null(),
333                "Backend should be null when zgemm is null"
334            );
335        }
336    }
337
338    #[test]
339    fn test_backend_new_from_fblas_ilp64_null_pointers() {
340        unsafe {
341            let backend =
342                spir_gemm_backend_new_from_fblas_ilp64(std::ptr::null(), std::ptr::null());
343            assert!(
344                backend.is_null(),
345                "Backend should be null when pointers are null"
346            );
347        }
348    }
349
350    #[test]
351    fn test_backend_release_null() {
352        unsafe {
353            // Should not panic when releasing null pointer
354            spir_gemm_backend_release(std::ptr::null_mut());
355        }
356    }
357
358    // System BLAS integration tests (only when system-blas feature is enabled)
359    #[cfg(all(test, feature = "system-blas"))]
360    mod system_blas_tests {
361        use super::*;
362        use blas_sys::{dgemm_, zgemm_};
363        use mdarray::tensor;
364        use sparse_ir::gemm::matmul_par;
365
366        // Helper to create backend from blas-sys functions
367        unsafe fn create_blas_backend() -> *mut spir_gemm_backend {
368            unsafe {
369                spir_gemm_backend_new_from_fblas_lp64(
370                    dgemm_ as *const _,
371                    // Cast zgemm_ to match our function pointer type (memory layout is compatible)
372                    unsafe {
373                        std::mem::transmute::<
374                            unsafe extern "C" fn(
375                                *const libc::c_char,
376                                *const libc::c_char,
377                                *const libc::c_int,
378                                *const libc::c_int,
379                                *const libc::c_int,
380                                *const blas_sys::c_double_complex,
381                                *const blas_sys::c_double_complex,
382                                *const libc::c_int,
383                                *const blas_sys::c_double_complex,
384                                *const libc::c_int,
385                                *const blas_sys::c_double_complex,
386                                *mut blas_sys::c_double_complex,
387                                *const libc::c_int,
388                            ),
389                            sparse_ir::gemm::ZgemmFnPtr,
390                        >(zgemm_)
391                    } as *const _,
392                )
393            }
394        }
395
396        #[test]
397        fn test_default_backend_matrix_multiplication_f64() {
398            unsafe {
399                // Use default backend (NULL means use default)
400                let backend = std::ptr::null();
401
402                // Test matrix multiplication: C = A * B
403                // A = [[1.0, 2.0], [3.0, 4.0]]
404                // B = [[5.0, 6.0], [7.0, 8.0]]
405                // Expected: C = [[19.0, 22.0], [43.0, 50.0]]
406                let a: mdarray::DTensor<f64, 2> = tensor![[1.0, 2.0], [3.0, 4.0]];
407                let b: mdarray::DTensor<f64, 2> = tensor![[5.0, 6.0], [7.0, 8.0]];
408                let backend_handle = get_backend_handle(backend);
409                let c = matmul_par(&a, &b, backend_handle);
410
411                // Verify results
412                assert!(
413                    (c[[0, 0]] - 19.0).abs() < 1e-10,
414                    "c[0,0] should be 19.0, got {}",
415                    c[[0, 0]]
416                );
417                assert!((c[[0, 1]] - 22.0).abs() < 1e-10, "c[0,1] should be 22.0");
418                assert!((c[[1, 0]] - 43.0).abs() < 1e-10, "c[1,0] should be 43.0");
419                assert!((c[[1, 1]] - 50.0).abs() < 1e-10, "c[1,1] should be 50.0");
420            }
421        }
422
423        #[test]
424        fn test_lp64_backend_matrix_multiplication_f64() {
425            unsafe {
426                // Create backend from system BLAS (LP64)
427                let backend = create_blas_backend();
428                assert!(!backend.is_null());
429
430                // Test matrix multiplication: C = A * B
431                // A = [[1.0, 2.0], [3.0, 4.0]]
432                // B = [[5.0, 6.0], [7.0, 8.0]]
433                // Expected: C = [[19.0, 22.0], [43.0, 50.0]]
434                let a: mdarray::DTensor<f64, 2> = tensor![[1.0, 2.0], [3.0, 4.0]];
435                let b: mdarray::DTensor<f64, 2> = tensor![[5.0, 6.0], [7.0, 8.0]];
436                let backend_handle = get_backend_handle(backend);
437                let c = matmul_par(&a, &b, backend_handle);
438
439                // Verify results
440                assert!(
441                    (c[[0, 0]] - 19.0).abs() < 1e-10,
442                    "c[0,0] should be 19.0, got {}",
443                    c[[0, 0]]
444                );
445                assert!((c[[0, 1]] - 22.0).abs() < 1e-10, "c[0,1] should be 22.0");
446                assert!((c[[1, 0]] - 43.0).abs() < 1e-10, "c[1,0] should be 43.0");
447                assert!((c[[1, 1]] - 50.0).abs() < 1e-10, "c[1,1] should be 50.0");
448
449                // Clean up
450                spir_gemm_backend_release(backend);
451            }
452        }
453
454        #[test]
455        fn test_default_backend_matrix_multiplication_complex() {
456            unsafe {
457                // Use default backend (NULL means use default)
458                let backend = std::ptr::null();
459
460                // Test complex matrix multiplication
461                let a: mdarray::DTensor<num_complex::Complex<f64>, 2> = tensor![
462                    [
463                        num_complex::Complex::new(1.0, 0.0),
464                        num_complex::Complex::new(2.0, 0.0)
465                    ],
466                    [
467                        num_complex::Complex::new(3.0, 0.0),
468                        num_complex::Complex::new(4.0, 0.0)
469                    ]
470                ];
471                let b: mdarray::DTensor<num_complex::Complex<f64>, 2> = tensor![
472                    [
473                        num_complex::Complex::new(5.0, 0.0),
474                        num_complex::Complex::new(6.0, 0.0)
475                    ],
476                    [
477                        num_complex::Complex::new(7.0, 0.0),
478                        num_complex::Complex::new(8.0, 0.0)
479                    ]
480                ];
481                let backend_handle = get_backend_handle(backend);
482                let c = matmul_par(&a, &b, backend_handle);
483
484                // Verify results (same as real case)
485                assert!((c[[0, 0]].re - 19.0).abs() < 1e-10);
486                assert!((c[[0, 1]].re - 22.0).abs() < 1e-10);
487                assert!((c[[1, 0]].re - 43.0).abs() < 1e-10);
488                assert!((c[[1, 1]].re - 50.0).abs() < 1e-10);
489                assert!(c[[0, 0]].im.abs() < 1e-10);
490            }
491        }
492
493        #[test]
494        fn test_lp64_backend_matrix_multiplication_complex() {
495            unsafe {
496                // Create backend from system BLAS (LP64)
497                let backend = create_blas_backend();
498                assert!(!backend.is_null());
499
500                // Test complex matrix multiplication
501                let a: mdarray::DTensor<num_complex::Complex<f64>, 2> = tensor![
502                    [
503                        num_complex::Complex::new(1.0, 0.0),
504                        num_complex::Complex::new(2.0, 0.0)
505                    ],
506                    [
507                        num_complex::Complex::new(3.0, 0.0),
508                        num_complex::Complex::new(4.0, 0.0)
509                    ]
510                ];
511                let b: mdarray::DTensor<num_complex::Complex<f64>, 2> = tensor![
512                    [
513                        num_complex::Complex::new(5.0, 0.0),
514                        num_complex::Complex::new(6.0, 0.0)
515                    ],
516                    [
517                        num_complex::Complex::new(7.0, 0.0),
518                        num_complex::Complex::new(8.0, 0.0)
519                    ]
520                ];
521                let backend_handle = get_backend_handle(backend);
522                let c = matmul_par(&a, &b, backend_handle);
523
524                // Verify results (same as real case)
525                assert!((c[[0, 0]].re - 19.0).abs() < 1e-10);
526                assert!((c[[0, 1]].re - 22.0).abs() < 1e-10);
527                assert!((c[[1, 0]].re - 43.0).abs() < 1e-10);
528                assert!((c[[1, 1]].re - 50.0).abs() < 1e-10);
529                assert!(c[[0, 0]].im.abs() < 1e-10);
530
531                // Clean up
532                spir_gemm_backend_release(backend);
533            }
534        }
535
536        #[test]
537        fn test_default_backend_larger_matrix() {
538            unsafe {
539                // Use default backend (NULL means use default)
540                let backend = std::ptr::null();
541
542                // Test with larger matrices (3x2 * 2x4 = 3x4)
543                let a: mdarray::DTensor<f64, 2> = tensor![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
544                let b: mdarray::DTensor<f64, 2> =
545                    tensor![[7.0, 8.0, 9.0, 10.0], [11.0, 12.0, 13.0, 14.0]];
546                let backend_handle = get_backend_handle(backend);
547                let c = matmul_par(&a, &b, backend_handle);
548
549                // Verify some results
550                // First row: [1*7+2*11, 1*8+2*12, 1*9+2*13, 1*10+2*14] = [29, 32, 35, 38]
551                assert!((c[[0, 0]] - 29.0).abs() < 1e-10);
552                assert!((c[[0, 1]] - 32.0).abs() < 1e-10);
553                assert!((c[[0, 2]] - 35.0).abs() < 1e-10);
554                assert!((c[[0, 3]] - 38.0).abs() < 1e-10);
555            }
556        }
557
558        #[test]
559        fn test_lp64_backend_larger_matrix() {
560            unsafe {
561                // Create backend from system BLAS (LP64)
562                let backend = create_blas_backend();
563                assert!(!backend.is_null());
564
565                // Test with larger matrices (3x2 * 2x4 = 3x4)
566                let a: mdarray::DTensor<f64, 2> = tensor![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
567                let b: mdarray::DTensor<f64, 2> =
568                    tensor![[7.0, 8.0, 9.0, 10.0], [11.0, 12.0, 13.0, 14.0]];
569                let backend_handle = get_backend_handle(backend);
570                let c = matmul_par(&a, &b, backend_handle);
571
572                // Verify some results
573                // First row: [1*7+2*11, 1*8+2*12, 1*9+2*13, 1*10+2*14] = [29, 32, 35, 38]
574                assert!((c[[0, 0]] - 29.0).abs() < 1e-10);
575                assert!((c[[0, 1]] - 32.0).abs() < 1e-10);
576                assert!((c[[0, 2]] - 35.0).abs() < 1e-10);
577                assert!((c[[0, 3]] - 38.0).abs() < 1e-10);
578
579                // Clean up
580                spir_gemm_backend_release(backend);
581            }
582        }
583    }
584}
sparse_ir_capi/gemm.rs

sparse_ir_capi/
gemm.rs
