strided-einsum2 0.1.2

Binary einsum (pairwise tensor contraction) on strided views.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

//! Backend abstraction for batched GEMM dispatch.
//!
//! This module defines the [`Backend`] trait, marker structs for each backend,
//! and the `ActiveBackend` type alias that serves as the single point of
//! backend selection based on Cargo features.

/// Trait for backends that can execute batched GEMM on contiguous operands.
///
/// Each backend declares its configuration (conjugation materialization,
/// stride requirements) and provides a GEMM implementation.
///
/// Implementations are provided by each backend module (faer, blas).
/// External crates can implement this trait for custom scalar types
/// (e.g., tropical semiring) and pass the backend to [`einsum2_with_backend_into`].
///
/// [`einsum2_with_backend_into`]: crate::einsum2_with_backend_into
pub trait Backend<T: crate::ScalarBase> {
    /// Whether the backend needs conjugation materialized into the data
    /// before GEMM (e.g., CBLAS has no conjugation flag for `?gemm`).
    const MATERIALIZES_CONJ: bool;

    /// Whether the backend requires at least one unit stride per matrix
    /// dimension (row or column stride must be 1). CBLAS `?gemm` requires
    /// this; faer does not.
    const REQUIRES_UNIT_STRIDE: bool;

    /// Execute batched GEMM: `C = alpha * A * B + beta * C` for each batch.
    ///
    /// - `c`: mutable output operand (batch x m x n)
    /// - `a`: input operand (batch x m x k)
    /// - `b`: input operand (batch x k x n)
    /// - `batch_dims`: sizes of the batch dimensions
    /// - `m`, `n`, `k`: fused matrix dimensions
    /// - `alpha`, `beta`: scaling factors
    fn bgemm_contiguous_into(
        c: &mut crate::contiguous::ContiguousOperandMut<T>,
        a: &crate::contiguous::ContiguousOperand<T>,
        b: &crate::contiguous::ContiguousOperand<T>,
        batch_dims: &[usize],
        m: usize,
        n: usize,
        k: usize,
        alpha: T,
        beta: T,
    ) -> strided_view::Result<()>;
}

// ---------------------------------------------------------------------------
// Marker structs
// ---------------------------------------------------------------------------

/// Batched GEMM backend using the [`faer`] library.
///
/// `Backend<T>` is implemented in `bgemm_faer.rs`.
#[cfg(feature = "faer")]
pub struct FaerBackend;

/// Batched GEMM backend using CBLAS (via `cblas-sys` or `cblas-inject`).
///
/// `Backend<T>` is implemented in `bgemm_blas.rs`.
#[cfg(any(feature = "blas", feature = "blas-inject"))]
pub struct BlasBackend;

/// Fallback batched GEMM backend using explicit loops (no external library).
///
/// This backend is used as `ActiveBackend` when no GEMM feature is enabled.
/// The GEMM dispatch in `einsum2_into` calls `bgemm_naive` directly rather
/// than going through the `Backend` trait, so `bgemm_contiguous_into` is
/// unreachable.
#[allow(dead_code)]
pub struct NaiveBackend;

impl<T: crate::ScalarBase> Backend<T> for NaiveBackend {
    const MATERIALIZES_CONJ: bool = false;
    const REQUIRES_UNIT_STRIDE: bool = false;

    fn bgemm_contiguous_into(
        _c: &mut crate::contiguous::ContiguousOperandMut<T>,
        _a: &crate::contiguous::ContiguousOperand<T>,
        _b: &crate::contiguous::ContiguousOperand<T>,
        _batch_dims: &[usize],
        _m: usize,
        _n: usize,
        _k: usize,
        _alpha: T,
        _beta: T,
    ) -> strided_view::Result<()> {
        unreachable!("NaiveBackend GEMM is dispatched directly, not through Backend trait")
    }
}

// ---------------------------------------------------------------------------
// ActiveBackend type alias -- the SINGLE point of backend selection
// ---------------------------------------------------------------------------

/// The active GEMM backend, selected by Cargo features.
///
/// - `blas` or `blas-inject` -> `BlasBackend`
/// - `faer` without BLAS -> `FaerBackend`
/// - no backend feature -> `NaiveBackend`
/// - invalid combos -> `NaiveBackend` (placeholder; `compile_error!` fires first)
#[cfg(any(
    all(feature = "blas", not(feature = "blas-inject")),
    all(feature = "blas-inject", not(feature = "blas"))
))]
pub type ActiveBackend = BlasBackend;

#[cfg(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject"))))]
pub type ActiveBackend = FaerBackend;

#[cfg(not(any(feature = "faer", feature = "blas", feature = "blas-inject")))]
pub type ActiveBackend = NaiveBackend;

/// Placeholder for invalid mutually-exclusive feature combinations.
///
/// The crate emits `compile_error!` for these combinations (in `lib.rs`), so this
/// alias only suppresses cascading type-resolution errors.
#[cfg(all(feature = "blas", feature = "blas-inject"))]
pub type ActiveBackend = NaiveBackend;