Skip to main content

diffsol_la/context/
mod.rs

1use crate::{error::LaError, DefaultDenseMatrix, Matrix, Vector};
2
3#[cfg(feature = "cuda")]
4pub mod cuda;
5
6#[cfg(feature = "nalgebra")]
7pub mod nalgebra;
8
9#[cfg(feature = "faer")]
10pub mod faer;
11
12/// defines the current execution and allocation context of an operator / vector / matrix
13/// for example:
14/// - threading model (e.g. single-threaded, multi-threaded, GPU)
15/// - custom allocators, host/device memory
16/// - etc.
17///
18/// It will generally be the case that all the operators / vectors / matrices for the current ode problem
19/// share the same context
20pub trait Context: Clone + Default {
21    /// Returns the batch count for this context.
22    ///
23    /// When `nbatch > 1`, vectors and matrices store data for `nbatch`
24    /// independent ODE systems simultaneously.  Operations between operands
25    /// with different batch counts use broadcast semantics: an operand with
26    /// `nbatch == 1` is applied to all batches of the other operand.
27    fn nbatch(&self) -> usize {
28        1
29    }
30    /// Creates a new context with the given batch count.
31    ///
32    /// Other properties of the context (e.g. CUDA stream, faer parallelism)
33    /// are preserved.
34    ///
35    /// Returns an error if the backend does not support batching (i.e.
36    /// `nbatch > 1` for CPU backends such as faer and nalgebra).
37    fn clone_with_nbatch(&self, nbatch: usize) -> Result<Self, LaError>;
38    /// Panics if the two batch counts are incompatible.
39    ///
40    /// Compatibility rule: two batches are compatible if they are equal, or
41    /// if either one is `1` (broadcast).  Only panics when both are `> 1`
42    /// and differ.
43    fn assert_compatible_nbatch(&self, other_nbatch: usize, op: &str) {
44        let self_nbatch = self.nbatch();
45        if self_nbatch != other_nbatch && self_nbatch != 1 && other_nbatch != 1 {
46            panic!(
47                "incompatible nbatch in {}: lhs={}, rhs={}",
48                op, self_nbatch, other_nbatch
49            );
50        }
51    }
52    fn vector_from_element<V: Vector<C = Self>>(&self, len: usize, value: V::T) -> V {
53        V::from_element(len, value, self.clone())
54    }
55    fn vector_from_vec<V: Vector<C = Self>>(&self, vec: Vec<V::T>) -> V {
56        V::from_vec(vec, self.clone())
57    }
58    fn vector_zeros<V: Vector<C = Self>>(&self, len: usize) -> V {
59        V::zeros(len, self.clone())
60    }
61    fn dense_mat_zeros<V: Vector<C = Self> + DefaultDenseMatrix>(
62        &self,
63        rows: usize,
64        cols: usize,
65    ) -> <V as DefaultDenseMatrix>::M {
66        <<V as DefaultDenseMatrix>::M as Matrix>::zeros(rows, cols, self.clone())
67    }
68}