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}