qmbed/
runtime.rs

1//! Execution and vector-storage boundary.
2//!
3//! Physics-facing code depends on [`crate::operator::LinearOperator`]. A
4//! runtime owns vectors and coarse-grained vector primitives. The built-in CPU
5//! implementation is complete; accelerator and multi-rank profiles are
6//! explicit extension points rather than silently falling back to the host.
7
8use std::sync::mpsc;
9use std::thread;
10
11use num_complex::Complex64;
12
13use crate::operator::LinearOperator;
14use crate::{QmbedError, Result};
15
16/// Accelerator requested by an execution profile.
17#[derive(Clone, Copy, Debug, Eq, PartialEq)]
18#[non_exhaustive]
19pub enum Accelerator {
20    /// Host CPU execution.
21    Cpu,
22    /// GPU execution on a zero-based device index.
23    Gpu { device: usize },
24}
25
26/// Language-neutral execution request.
27///
28/// `ranks > 1` denotes a distributed run. Each rank may in turn use the
29/// requested accelerator and number of host threads.
30#[derive(Clone, Copy, Debug, Eq, PartialEq)]
31#[non_exhaustive]
32pub struct ExecutionProfile {
33    /// Compute device requested for each rank.
34    pub accelerator: Accelerator,
35    /// Number of cooperating process ranks.
36    pub ranks: usize,
37    /// Host threads available to each rank.
38    pub threads_per_rank: usize,
39}
40
41impl ExecutionProfile {
42    /// Request one CPU rank and one host thread.
43    pub const fn serial() -> Self {
44        Self::local_cpu(1)
45    }
46
47    /// Request bounded host parallelism for independent work items.
48    pub const fn throughput(threads: usize) -> Self {
49        Self::local_cpu(threads)
50    }
51
52    /// Request one CPU rank with `threads` host workers.
53    pub const fn local_cpu(threads: usize) -> Self {
54        Self {
55            accelerator: Accelerator::Cpu,
56            ranks: 1,
57            threads_per_rank: threads,
58        }
59    }
60
61    /// Request a distributed CPU run.
62    ///
63    /// The built-in runtime rejects this profile until an MPI implementation
64    /// of [`Runtime`] is installed.
65    pub const fn distributed_cpu(ranks: usize, threads_per_rank: usize) -> Self {
66        Self {
67            accelerator: Accelerator::Cpu,
68            ranks,
69            threads_per_rank,
70        }
71    }
72
73    /// Request one GPU device with the supplied host-thread allowance.
74    pub const fn local_gpu(device: usize, threads: usize) -> Self {
75        Self {
76            accelerator: Accelerator::Gpu { device },
77            ranks: 1,
78            threads_per_rank: threads,
79        }
80    }
81
82    /// Validate that rank and thread counts are positive.
83    pub fn validate(self) -> Result<Self> {
84        if self.ranks == 0 || self.threads_per_rank == 0 {
85            return Err(QmbedError::InvalidOptions(
86                "execution profiles require at least one rank and one thread per rank".into(),
87            ));
88        }
89        Ok(self)
90    }
91}
92
93/// Concrete resources supplied by one runtime implementation.
94#[derive(Clone, Copy, Debug, Eq, PartialEq)]
95pub struct RuntimeCapabilities {
96    /// Active accelerator.
97    pub accelerator: Accelerator,
98    /// Number of cooperating process ranks.
99    pub ranks: usize,
100    /// Host threads available to each rank.
101    pub threads_per_rank: usize,
102}
103
104/// Vector storage owned by a [`Runtime`].
105pub trait RuntimeBuffer: Send + Sync {
106    /// Return the number of complex elements.
107    fn len(&self) -> usize;
108
109    /// Return whether the buffer has no elements.
110    fn is_empty(&self) -> bool {
111        self.len() == 0
112    }
113}
114
115/// Coarse-grained storage and vector primitives required by iterative methods.
116///
117/// A future GPU or MPI backend implements this trait with its native buffer.
118/// No basis, model, or operator-string type crosses this boundary.
119pub trait Runtime: Send + Sync {
120    /// Backend-native vector type.
121    type Buffer: RuntimeBuffer;
122
123    /// Report the execution resources supplied by this runtime.
124    fn capabilities(&self) -> RuntimeCapabilities;
125    /// Allocate a zero-filled complex vector.
126    fn zeros(&self, length: usize) -> Result<Self::Buffer>;
127    /// Copy a host slice into backend-owned storage.
128    fn upload(&self, values: &[Complex64]) -> Result<Self::Buffer>;
129    /// Copy backend-owned storage to the host.
130    fn to_host(&self, buffer: &Self::Buffer) -> Result<Vec<Complex64>>;
131    /// Fill a backend buffer with one complex value.
132    fn fill(&self, buffer: &mut Self::Buffer, value: Complex64) -> Result<()>;
133    /// Compute `output += alpha * input`.
134    fn axpy(&self, alpha: Complex64, input: &Self::Buffer, output: &mut Self::Buffer)
135    -> Result<()>;
136    /// Scale a vector in place.
137    fn scale(&self, alpha: Complex64, buffer: &mut Self::Buffer) -> Result<()>;
138    /// Return the conjugating inner product of two buffers.
139    fn dotc(&self, left: &Self::Buffer, right: &Self::Buffer) -> Result<Complex64>;
140
141    /// Return the Euclidean norm derived from [`Runtime::dotc`].
142    fn norm(&self, buffer: &Self::Buffer) -> Result<f64> {
143        Ok(self.dotc(buffer, buffer)?.re.max(0.0).sqrt())
144    }
145}
146
147/// Runtime-aware action without changing the scientific operator contract.
148pub trait RuntimeLinearOperator<R>
149where
150    R: Runtime,
151{
152    /// Return `(rows, columns)` in runtime coordinates.
153    fn runtime_shape(&self) -> (usize, usize);
154    /// Apply this map to backend-owned buffers.
155    fn apply_on(&self, runtime: &R, input: &R::Buffer, output: &mut R::Buffer) -> Result<()>;
156}
157
158/// Runtime-aware adjoint action for reverse-mode scientific rules.
159///
160/// This stays separate from [`RuntimeLinearOperator`] so existing execution
161/// backends remain source-compatible. A GPU or distributed backend implements
162/// this trait with its native adjoint kernel; the CPU blanket implementation
163/// delegates to [`LinearOperator::apply_adjoint`].
164pub trait RuntimeAdjointLinearOperator<R>: RuntimeLinearOperator<R>
165where
166    R: Runtime,
167{
168    /// Compute `output = A† * input` in runtime-owned storage.
169    fn apply_adjoint_on(
170        &self,
171        runtime: &R,
172        input: &R::Buffer,
173        output: &mut R::Buffer,
174    ) -> Result<()>;
175}
176
177/// Host-resident complex vector used by [`CpuRuntime`].
178#[derive(Clone, Debug, PartialEq)]
179pub struct CpuBuffer {
180    values: Vec<Complex64>,
181}
182
183impl CpuBuffer {
184    /// Borrow the host values.
185    pub fn as_slice(&self) -> &[Complex64] {
186        &self.values
187    }
188
189    /// Mutably borrow the host values.
190    pub fn as_mut_slice(&mut self) -> &mut [Complex64] {
191        &mut self.values
192    }
193}
194
195impl RuntimeBuffer for CpuBuffer {
196    fn len(&self) -> usize {
197        self.values.len()
198    }
199}
200
201/// Built-in single-rank CPU runtime.
202#[derive(Clone, Copy, Debug, Eq, PartialEq)]
203pub struct CpuRuntime {
204    profile: ExecutionProfile,
205}
206
207impl CpuRuntime {
208    /// Construct a local CPU runtime with a positive worker count.
209    pub fn new(threads: usize) -> Result<Self> {
210        Self::from_profile(ExecutionProfile::local_cpu(threads))
211    }
212
213    /// Resolve a profile against the built-in runtime.
214    ///
215    /// Unsupported profiles fail explicitly so a requested GPU or MPI run
216    /// cannot be mistaken for a successful serial CPU calculation.
217    pub fn from_profile(profile: ExecutionProfile) -> Result<Self> {
218        let profile = profile.validate()?;
219        if profile.accelerator != Accelerator::Cpu {
220            return Err(QmbedError::UnsupportedBackend(
221                "the built-in runtime is CPU-only; install a GPU runtime implementation".into(),
222            ));
223        }
224        if profile.ranks != 1 {
225            return Err(QmbedError::UnsupportedBackend(
226                "the built-in runtime is single-rank; install an MPI runtime implementation".into(),
227            ));
228        }
229        Ok(Self { profile })
230    }
231
232    /// Return the validated profile used to construct this runtime.
233    pub const fn profile(self) -> ExecutionProfile {
234        self.profile
235    }
236
237    /// Apply an independent operation with bounded shared-memory parallelism.
238    ///
239    /// Results and errors are returned in input order regardless of worker
240    /// scheduling. A one-thread profile executes through the same contract.
241    pub fn map_ordered<T, U, F>(&self, items: &[T], operation: F) -> Result<Vec<U>>
242    where
243        T: Sync,
244        U: Send,
245        F: Fn(&T) -> Result<U> + Sync,
246    {
247        if items.is_empty() {
248            return Ok(Vec::new());
249        }
250        let workers = self.profile.threads_per_rank.min(items.len());
251        let mut ordered = (0..items.len()).map(|_| None).collect::<Vec<_>>();
252        thread::scope(|scope| {
253            let (sender, receiver) = mpsc::channel();
254            let mut handles = Vec::with_capacity(workers);
255            for worker in 0..workers {
256                let sender = sender.clone();
257                let operation = &operation;
258                handles.push(scope.spawn(move || {
259                    for index in (worker..items.len()).step_by(workers) {
260                        if sender.send((index, operation(&items[index]))).is_err() {
261                            break;
262                        }
263                    }
264                }));
265            }
266            drop(sender);
267            for (index, result) in receiver {
268                ordered[index] = Some(result);
269            }
270            if handles.into_iter().any(|handle| handle.join().is_err()) {
271                return Err(QmbedError::UnsupportedBackend(
272                    "a CPU runtime worker panicked".into(),
273                ));
274            }
275            let mut output = Vec::with_capacity(items.len());
276            for result in ordered {
277                let result = result.ok_or_else(|| {
278                    QmbedError::UnsupportedBackend(
279                        "a CPU runtime worker did not return a result".into(),
280                    )
281                })?;
282                output.push(result?);
283            }
284            Ok(output)
285        })
286    }
287}
288
289impl Runtime for CpuRuntime {
290    type Buffer = CpuBuffer;
291
292    fn capabilities(&self) -> RuntimeCapabilities {
293        RuntimeCapabilities {
294            accelerator: Accelerator::Cpu,
295            ranks: 1,
296            threads_per_rank: self.profile.threads_per_rank,
297        }
298    }
299
300    fn zeros(&self, length: usize) -> Result<Self::Buffer> {
301        Ok(CpuBuffer {
302            values: vec![Complex64::new(0.0, 0.0); length],
303        })
304    }
305
306    fn upload(&self, values: &[Complex64]) -> Result<Self::Buffer> {
307        Ok(CpuBuffer {
308            values: values.to_vec(),
309        })
310    }
311
312    fn to_host(&self, buffer: &Self::Buffer) -> Result<Vec<Complex64>> {
313        Ok(buffer.values.clone())
314    }
315
316    fn fill(&self, buffer: &mut Self::Buffer, value: Complex64) -> Result<()> {
317        buffer.values.fill(value);
318        Ok(())
319    }
320
321    fn axpy(
322        &self,
323        alpha: Complex64,
324        input: &Self::Buffer,
325        output: &mut Self::Buffer,
326    ) -> Result<()> {
327        equal_lengths(input, output)?;
328        for (target, source) in output.values.iter_mut().zip(&input.values) {
329            *target += alpha * *source;
330        }
331        Ok(())
332    }
333
334    fn scale(&self, alpha: Complex64, buffer: &mut Self::Buffer) -> Result<()> {
335        for value in &mut buffer.values {
336            *value *= alpha;
337        }
338        Ok(())
339    }
340
341    fn dotc(&self, left: &Self::Buffer, right: &Self::Buffer) -> Result<Complex64> {
342        equal_lengths(left, right)?;
343        Ok(left
344            .values
345            .iter()
346            .zip(&right.values)
347            .map(|(left, right)| left.conj() * *right)
348            .sum())
349    }
350}
351
352impl<O> RuntimeLinearOperator<CpuRuntime> for O
353where
354    O: LinearOperator + ?Sized,
355{
356    fn runtime_shape(&self) -> (usize, usize) {
357        self.shape()
358    }
359
360    fn apply_on(
361        &self,
362        _runtime: &CpuRuntime,
363        input: &CpuBuffer,
364        output: &mut CpuBuffer,
365    ) -> Result<()> {
366        self.apply(input.as_slice(), output.as_mut_slice())
367    }
368}
369
370impl<O> RuntimeAdjointLinearOperator<CpuRuntime> for O
371where
372    O: LinearOperator + ?Sized,
373{
374    fn apply_adjoint_on(
375        &self,
376        _runtime: &CpuRuntime,
377        input: &CpuBuffer,
378        output: &mut CpuBuffer,
379    ) -> Result<()> {
380        self.apply_adjoint(input.as_slice(), output.as_mut_slice())
381    }
382}
383
384fn equal_lengths(left: &impl RuntimeBuffer, right: &impl RuntimeBuffer) -> Result<()> {
385    if left.len() != right.len() {
386        return Err(QmbedError::DimensionMismatch(format!(
387            "runtime vector lengths differ: {} and {}",
388            left.len(),
389            right.len()
390        )));
391    }
392    Ok(())
393}
qmbed/runtime.rs

qmbed/
runtime.rs
