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}