Skip to main content

transcribe_cpp/
backend.rs

1//! Backend module loading and compute-device discovery.
2//!
3//! In a static (or plain `shared`) build the compiled-in backends are already
4//! registered and [`init_backends`] is a harmless no-op. In a `dynamic-backends`
5//! build the compute backends are loadable modules; a host points the library at
6//! the provider directory ONCE, before the first model load — use
7//! [`init_backends_default`] when the modules sit next to libtranscribe, or
8//! [`init_backends`] with an explicit bundled provider directory. See the C
9//! header's backend-module section for the degradation contract.
10
11use std::ffi::CString;
12use std::path::Path;
13
14use transcribe_cpp_sys as sys;
15
16use crate::error::{check, Result};
17use crate::result::{owned_opt_str, owned_str};
18use crate::types::Backend;
19
20/// ggml's vendor-agnostic class for a compute device, orthogonal to
21/// [`Device::kind`] (which carries the vendor). Backends report this
22/// classification themselves, so use it as a runtime hint rather than a
23/// portable hardware-memory taxonomy.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
26pub enum DeviceType {
27    /// CPU using system memory.
28    Cpu,
29    /// Backend-reported GPU.
30    Gpu,
31    /// Backend-reported integrated GPU.
32    Igpu,
33    /// Host-memory accelerator (BLAS/AMX/...).
34    Accel,
35    /// A device-type value this binding does not recognize — reported by a
36    /// runtime newer than the header this binding was generated against.
37    /// Distinguish such devices by [`Device::device_id`] / [`Device::name`],
38    /// not by this axis.
39    Unknown,
40}
41
42impl DeviceType {
43    fn from_raw(raw: sys::transcribe_device_type) -> Self {
44        use sys::transcribe_device_type as T;
45        match raw {
46            T::TRANSCRIBE_DEVICE_TYPE_CPU => DeviceType::Cpu,
47            T::TRANSCRIBE_DEVICE_TYPE_GPU => DeviceType::Gpu,
48            T::TRANSCRIBE_DEVICE_TYPE_IGPU => DeviceType::Igpu,
49            T::TRANSCRIBE_DEVICE_TYPE_ACCEL => DeviceType::Accel,
50            _ => DeviceType::Unknown,
51        }
52    }
53}
54
55/// One registered compute device.
56#[derive(Debug, Clone)]
57pub struct Device {
58    /// ggml device name, e.g. "Metal".
59    pub name: String,
60    /// Human-readable description, e.g. "Apple M4 Max".
61    pub description: String,
62    /// Classified vendor kind: "cpu", "accel", "metal", "vulkan", "cuda",
63    /// "rocm", "sycl", "gpu", or "unknown".
64    pub kind: String,
65    /// The CPU/GPU/IGPU/ACCEL axis, orthogonal to [`Device::kind`].
66    pub device_type: DeviceType,
67    /// Stable hardware id when the backend reports one (PCI bus id for PCI
68    /// devices), or `None` (e.g. Metal).
69    pub device_id: Option<String>,
70    /// Reported device memory capacity in bytes, or 0 if unreported.
71    pub memory_total: u64,
72    /// Available device memory in bytes — a snapshot at the time this was
73    /// queried, or 0 if unreported. Re-query (via [`devices`] or
74    /// [`crate::Model::device`]) to refresh it; the value is backend-defined
75    /// and not comparable across device kinds.
76    pub memory_free: u64,
77    /// Registry index when this device came from [`devices`]. This is useful
78    /// for display only; pass the [`Device`] itself to [`ModelOptions`] for
79    /// exact selection. Registry indices are not stable across processes.
80    pub index: Option<usize>,
81    pub(crate) handle: sys::transcribe_device_t,
82}
83
84// Device handles are immutable process-lifetime registry entries. Backend
85// registration must finish before enumeration, so sharing a handle is safe.
86unsafe impl Send for Device {}
87unsafe impl Sync for Device {}
88
89impl PartialEq for Device {
90    fn eq(&self, other: &Self) -> bool {
91        self.handle == other.handle
92    }
93}
94
95impl Eq for Device {}
96
97impl Device {
98    /// Build a [`Device`] from the raw FFI struct filled by the library.
99    /// `index` is the registry index when the device came from enumeration,
100    /// or `None` when it came from a loaded model.
101    pub(crate) fn from_raw(
102        raw: &sys::transcribe_device_info,
103        handle: sys::transcribe_device_t,
104        index: Option<usize>,
105    ) -> Device {
106        Device {
107            name: owned_str(raw.name),
108            description: owned_str(raw.description),
109            kind: owned_str(raw.kind),
110            device_type: DeviceType::from_raw(raw.device_type),
111            device_id: owned_opt_str(raw.device_id),
112            memory_total: raw.memory_total,
113            memory_free: raw.memory_free,
114            index,
115            handle,
116        }
117    }
118}
119
120/// Load backend modules from `dir` (dynamic builds) and register their
121/// devices. A no-op in static builds. Call once, before the first model load.
122///
123/// Errors with [`crate::Error::Backend`] if, after loading, the process has no
124/// registered compute device (a dynamic build pointed at a directory with no
125/// usable modules). This call is idempotent per directory and NOT retryable in
126/// the same process.
127///
128/// This mutates the native process-global device registry. Complete every
129/// backend-init call before other threads enumerate devices, query backend
130/// availability, or load models; the native registry does not support racing
131/// registration against those operations.
132pub fn init_backends(dir: impl AsRef<Path>) -> Result<()> {
133    let dir = dir.as_ref();
134    // Pass the path bytes through faithfully (Unix) / reject non-UTF-8 (Windows),
135    // matching model loading — never lossily mangle a path with to_string_lossy.
136    let c_dir = CString::new(crate::model::path_bytes(dir)?)?;
137    let status = unsafe { sys::transcribe_init_backends(c_dir.as_ptr()) };
138    check(status, "init_backends")
139}
140
141/// Load the backend modules the way the build expects, with no caller-supplied
142/// path.
143///
144/// - **Compiled-in builds** (the default static build, or a plain `shared`
145///   build): a no-op returning `Ok(())` — the backends are already registered.
146/// - **`dynamic-backends` builds**: resolves the directory containing the
147///   loaded libtranscribe and scans only that directory. Ship the backend
148///   modules next to libtranscribe for this helper to find them.
149///
150/// If your app uses a different layout, call [`init_backends`] with that
151/// resolved module directory instead. Like [`init_backends`], this is
152/// idempotent and must run once before the first model load. It must also
153/// complete before other threads enumerate devices or query backend
154/// availability; see [`init_backends`] for the registry-ordering contract.
155pub fn init_backends_default() -> Result<()> {
156    let status = unsafe { sys::transcribe_init_backends_default() };
157    check(status, "init_backends_default")
158}
159
160/// The number of compute devices currently registered.
161///
162/// Do not race this query with [`init_backends`] or [`init_backends_default`].
163pub fn device_count() -> usize {
164    let n = unsafe { sys::transcribe_device_count() };
165    n.max(0) as usize
166}
167
168/// Every registered compute device.
169///
170/// Do not race enumeration with [`init_backends`] or
171/// [`init_backends_default`]. Finish backend registration before sharing
172/// devices across threads.
173pub fn devices() -> Vec<Device> {
174    let mut out = Vec::with_capacity(device_count());
175    for i in 0..device_count() as i32 {
176        let handle = unsafe { sys::transcribe_device_get(i) };
177        if handle.is_null() {
178            continue;
179        }
180        let mut raw: sys::transcribe_device_info = unsafe { std::mem::zeroed() };
181        unsafe { sys::transcribe_device_info_init(&mut raw) };
182        let status = unsafe { sys::transcribe_device_get_info(handle, &mut raw) };
183        if status == sys::transcribe_status::TRANSCRIBE_OK {
184            out.push(Device::from_raw(&raw, handle, Some(i as usize)));
185        }
186    }
187    out
188}
189
190/// Whether a backend request can be satisfied by some registered device. This
191/// is the probe to turn `Backend::Vulkan` on a machine without Vulkan into a
192/// clear error instead of a failed model load. Do not race this query with
193/// [`init_backends`] or [`init_backends_default`].
194pub fn backend_available(backend: Backend) -> bool {
195    unsafe { sys::transcribe_backend_available(backend.to_raw()) }
196}