Skip to main content

ironaccelerator_core/
backend.rs

1//! [`Backend`] is the single trait every accelerator backend implements.
2//!
3//! The trait is intentionally small and object-safe so the runtime can hold a
4//! `&'static dyn Backend` table and dispatch without monomorphisation overhead
5//! for the rare cross-backend code paths. Hot paths (kernel launches, memcpys)
6//! live on the concrete backend type and are inlined.
7
8use crate::{capability::CapabilityFlags, device::DeviceDescriptor, error::Result};
9
10#[cfg(not(feature = "std"))]
11use alloc::vec::Vec;
12#[cfg(feature = "std")]
13use std::vec::Vec;
14
15/// The set of accelerator families IronAccelerator can target.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
18#[non_exhaustive]
19pub enum BackendKind {
20    /// NVIDIA CUDA (Toolkit 13.2+ targeted).
21    Cuda,
22    /// AMD ROCm / HIP.
23    Rocm,
24    /// Apple Metal (Performance Shaders + MLX kernels).
25    Metal,
26    /// Qualcomm AI Engine — Hexagon NPU via QNN SDK.
27    QualcommNpu,
28    /// CPU SIMD reference path (used as a fallback / oracle).
29    Cpu,
30    /// Vulkan Compute — cross-vendor GPU compute.
31    Vulkan,
32    /// OpenGL 4.3+ compute shaders — legacy / embedded GPU fallback.
33    OpenGl,
34    /// Direct3D 12 — the Windows-native GPU API, across all vendors.
35    Dx12,
36    /// WebGPU in the browser — the WASM compute path.
37    WebGpu,
38    /// Google TPU (v4 / v5 / v6e) via the PJRT plugin interface.
39    Tpu,
40    /// Intel GPU + NPU via oneAPI Level Zero (`ze_loader`).
41    LevelZero,
42    /// AWS Trainium / Inferentia via the Neuron Runtime (`libnrt`).
43    Neuron,
44}
45
46impl BackendKind {
47    pub const ALL: &'static [BackendKind] = &[
48        BackendKind::Cuda,
49        BackendKind::Rocm,
50        BackendKind::Metal,
51        BackendKind::QualcommNpu,
52        BackendKind::Cpu,
53        BackendKind::Vulkan,
54        BackendKind::OpenGl,
55        BackendKind::Dx12,
56        BackendKind::WebGpu,
57        BackendKind::Tpu,
58        BackendKind::LevelZero,
59        BackendKind::Neuron,
60    ];
61
62    pub const fn name(self) -> &'static str {
63        match self {
64            BackendKind::Cuda => "cuda",
65            BackendKind::Rocm => "rocm",
66            BackendKind::Metal => "metal",
67            BackendKind::QualcommNpu => "qnn",
68            BackendKind::Cpu => "cpu",
69            BackendKind::Vulkan => "vulkan",
70            BackendKind::OpenGl => "opengl",
71            BackendKind::Dx12 => "dx12",
72            BackendKind::WebGpu => "webgpu",
73            BackendKind::Tpu => "tpu",
74            BackendKind::LevelZero => "level-zero",
75            BackendKind::Neuron => "neuron",
76        }
77    }
78}
79
80/// The minimum capability surface a backend must expose.
81///
82/// Object-safe; concrete backends typically also expose a non-trait API with
83/// inlined fast paths.
84///
85/// The trait is deliberately limited to *discovery* — identity, availability,
86/// devices, capability bits. Choosing what to run on a device is a planner
87/// concern and lives in the consumer, not here.
88pub trait Backend: Send + Sync + 'static {
89    /// Static identity.
90    fn kind(&self) -> BackendKind;
91
92    /// Whether the backend's runtime libraries were located on this host.
93    fn is_available(&self) -> bool;
94
95    /// Enumerate every visible device.
96    fn enumerate(&self) -> Result<Vec<DeviceDescriptor>>;
97
98    /// Coarse capability bits for a device, translated from the vendor's
99    /// capability table into the common [`CapabilityFlags`] space.
100    fn capabilities(&self, device: u32) -> Result<CapabilityFlags>;
101}
102
103/// Process-wide registry of compiled-in backends. Backends register themselves
104/// from their crate's `init()` (called from `ironaccelerator::init`).
105pub struct BackendRegistry {
106    entries: Vec<&'static dyn Backend>,
107}
108
109impl BackendRegistry {
110    pub fn new() -> Self {
111        Self {
112            entries: Vec::new(),
113        }
114    }
115
116    pub fn register(&mut self, backend: &'static dyn Backend) {
117        if !self.entries.iter().any(|b| b.kind() == backend.kind()) {
118            self.entries.push(backend);
119        }
120    }
121
122    pub fn get(&self, kind: BackendKind) -> Option<&'static dyn Backend> {
123        self.entries.iter().copied().find(|b| b.kind() == kind)
124    }
125
126    pub fn iter(&self) -> impl Iterator<Item = &'static dyn Backend> + '_ {
127        self.entries.iter().copied()
128    }
129
130    pub fn available(&self) -> impl Iterator<Item = &'static dyn Backend> + '_ {
131        self.iter().filter(|b| b.is_available())
132    }
133
134    /// Enumerate every device across every *available* backend. Errors
135    /// from any one backend are swallowed (we return an empty list for
136    /// it) so a single broken backend can't mask the rest — callers
137    /// that want error visibility should iterate manually.
138    pub fn describe_all(&self) -> Vec<DeviceDescriptor> {
139        self.available()
140            .flat_map(|b| b.enumerate().unwrap_or_default())
141            .collect()
142    }
143}
144
145impl Default for BackendRegistry {
146    fn default() -> Self {
147        Self::new()
148    }
149}