burn_std/
runtime_kind.rs

1//! Runtime kind of the host program.
2//!
3//! Some backend decisions depend not on the device but on *how the host program itself is
4//! being driven* — whether `main` runs on an asynchronous runtime, a synchronous
5//! thread-based runtime, or a restricted no-std environment. This module stores that kind
6//! in a process global so backends can read it and adapt their behavior.
7//!
8//! The canonical example is tensor readback (`into_data`): deferring the device→host copy
9//! lazily is fine under a sync/threaded runtime (a later blocking read just parks a thread),
10//! but under an async runtime the same blocking read parks an executor worker and starves
11//! the runtime, so the read must materialize eagerly instead.
12
13use core::sync::atomic::{AtomicU8, Ordering};
14
15/// How the host program is being driven.
16///
17/// Set once near program start with [`set_runtime_kind`]; read anywhere with
18/// [`runtime_kind`]. Defaults to [`RuntimeKind::Sync`].
19#[repr(u8)]
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
21pub enum RuntimeKind {
22    /// Synchronous, thread-based runtime (the default).
23    ///
24    /// Blocking readbacks are fine here, so tensor reads may defer the device→host copy
25    /// lazily and materialize it on first access.
26    #[default]
27    Sync = 0,
28    /// Asynchronous runtime (e.g. tokio).
29    ///
30    /// Blocking a runtime worker starves the executor, so tensor reads must materialize
31    /// eagerly inside the awaited future rather than deferring a blocking copy.
32    Async = 1,
33    /// Restricted no-std environment.
34    NoStd = 2,
35}
36
37impl RuntimeKind {
38    fn from_u8(value: u8) -> Self {
39        match value {
40            1 => RuntimeKind::Async,
41            2 => RuntimeKind::NoStd,
42            _ => RuntimeKind::Sync,
43        }
44    }
45}
46
47static RUNTIME_KIND: AtomicU8 = AtomicU8::new(RuntimeKind::Sync as u8);
48
49/// Declare the [kind](RuntimeKind) of runtime the host program is running on.
50///
51/// Intended to be called once near program start (e.g. when a remote server declares that
52/// it hosts the backend on an async runtime). Backends read it via [`runtime_kind`].
53pub fn set_runtime_kind(kind: RuntimeKind) {
54    RUNTIME_KIND.store(kind as u8, Ordering::Relaxed);
55}
56
57/// Return the currently declared [kind](RuntimeKind) of runtime the host program runs on.
58///
59/// Defaults to [`RuntimeKind::Sync`] until [`set_runtime_kind`] is called.
60pub fn runtime_kind() -> RuntimeKind {
61    RuntimeKind::from_u8(RUNTIME_KIND.load(Ordering::Relaxed))
62}
burn_std/runtime_kind.rs

burn_std/
runtime_kind.rs
