diffsol_la/context/
mod.rs

1use crate::{error::LaError, DefaultDenseMatrix, Matrix, Vector};
2
3#[cfg(feature = "cuda")]
4pub mod cuda;
5
6#[cfg(feature = "cuda-oxide")]
7pub mod cuda_oxide;
8
9#[cfg(feature = "nalgebra")]
10pub mod nalgebra;
11
12#[cfg(feature = "faer")]
13pub mod faer;
14
15/// The panic raised when the broadcast-into rule is violated.  Kept out of line so that
16/// [`Context::assert_broadcastable_into`] inlines as the comparisons and the division and
17/// nothing else -- in particular not the panic's formatting.
18#[cold]
19#[inline(never)]
20fn panic_incompatible_nbatch(src_nbatch: usize, dest_nbatch: usize, op: &str) -> ! {
21    panic!(
22        "incompatible nbatch in {}: lhs={}, rhs={}",
23        op, dest_nbatch, src_nbatch
24    );
25}
26
27/// Source batch feeding destination batch `dest` of `dest_nbatch`, for an operand holding
28/// `src_nbatch` batches. Broadcast is grouped: each source batch is repeated across `dest_nbatch / src_nbatch`
29/// *contiguous* destination batches
30#[inline]
31pub(crate) fn broadcast_batch(dest: usize, src_nbatch: usize, dest_nbatch: usize) -> usize {
32    // matching batch counts are by far the common case, and returning `dest` unchanged keeps
33    // the caller's index provably in range -- without this branch the division loses that and
34    // the elementwise loops downstream deoptimise (~6% on lin_alg_ops axpy/faer/100)
35    if src_nbatch == dest_nbatch {
36        return dest;
37    }
38    dest * src_nbatch / dest_nbatch
39}
40
41/// defines the current execution and allocation context of an operator / vector / matrix
42/// for example:
43/// - threading model (e.g. single-threaded, multi-threaded, GPU)
44/// - custom allocators, host/device memory
45/// - etc.
46///
47/// It will generally be the case that all the operators / vectors / matrices for the current ode problem
48/// share the same context
49pub trait Context: Clone + Default {
50    /// Returns the batch count for this context.
51    ///
52    /// When `nbatch > 1`, vectors and matrices store data for `nbatch`
53    /// independent ODE systems simultaneously.  Operations between operands
54    /// with different batch counts use broadcast semantics: an operand with
55    /// `nbatch == B` is applied to `B * P` batches of the other operand, repeating each
56    /// of its batches over `P` contiguous batches (`nbatch == 1` is the `B == 1` case,
57    /// applied to every batch of the other operand).
58    #[inline]
59    fn nbatch(&self) -> usize {
60        1
61    }
62    /// Waits for any queued work to finish.
63    ///
64    /// A no-op on the CPU backends, where every operation has already run by the time it
65    /// returns. On a device backend the operations only enqueue, so anything timing them has to
66    /// call this before reading the clock.
67    #[inline]
68    fn synchronize(&self) {}
69    /// Creates a new context with the given batch count.
70    ///
71    /// Other properties of the context (e.g. CUDA stream, faer parallelism)
72    /// are preserved.
73    ///
74    /// Returns an error when a backend cannot create a requested batch context.
75    fn clone_with_nbatch(&self, nbatch: usize) -> Result<Self, LaError>;
76    /// Panics unless an operand holding `src_nbatch` batches broadcasts *into* `self`'s
77    /// batches, i.e. unless `self.nbatch()` is an exact multiple of `src_nbatch`.
78    ///
79    /// Every batched operation picks one operand whose batch count governs the result, and
80    /// `self` is always that operand:
81    ///
82    /// - an **owned** operand has priority, because the result can be written into it --
83    ///   `A + &B` and `A + B` are governed by `A`, `&A + B` by `B`;
84    /// - otherwise the **left-hand side** governs, and the result is a fresh allocation at its
85    ///   batch count -- `&A + &B` is governed by `A`;
86    /// - operations with a fixed destination are governed by it: assignments, `copy_from`,
87    ///   `axpy`, `gemv` into `y`, a linear solve in place.
88    ///
89    /// The other operand is then broadcast into the governing count: each of its batches is
90    /// repeated across `dest_nbatch / src_nbatch` *contiguous* destination batches.
91    #[inline]
92    fn assert_broadcastable_into(&self, src_nbatch: usize, op: &str) {
93        let nbatch = self.nbatch();
94        if src_nbatch == nbatch || src_nbatch == 1 {
95            return;
96        }
97        // grouped broadcast: the division stays inline, only the panic is out of line
98        if src_nbatch == 0 || !nbatch.is_multiple_of(src_nbatch) {
99            panic_incompatible_nbatch(src_nbatch, nbatch, op);
100        }
101    }
102    fn vector_from_element<V: Vector<C = Self>>(&self, len: usize, value: V::T) -> V {
103        V::from_element(len, value, self.clone())
104    }
105    fn vector_from_vec<V: Vector<C = Self>>(&self, vec: Vec<V::T>) -> V {
106        V::from_vec(vec, self.clone())
107    }
108    fn vector_zeros<V: Vector<C = Self>>(&self, len: usize) -> V {
109        V::zeros(len, self.clone())
110    }
111    fn dense_mat_zeros<V: Vector<C = Self> + DefaultDenseMatrix>(
112        &self,
113        rows: usize,
114        cols: usize,
115    ) -> <V as DefaultDenseMatrix>::M {
116        <<V as DefaultDenseMatrix>::M as Matrix>::zeros(rows, cols, self.clone())
117    }
118}
diffsol_la/context/mod.rs

diffsol_la/context/
mod.rs
