diffsol_la/context/
mod.rs

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

diffsol_la/context/
mod.rs
