Skip to main content

svod_runtime/
device_registry.rs

1//! Device factory registry for runtime device creation and caching.
2//!
3//! This module provides a registry for full Device objects (renderer + compiler + runtime + allocator).
4//! It's separate from `svod_device::registry::DeviceRegistry` (which only manages allocators)
5//! to avoid circular dependencies between `device` and `runtime` crates.
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use once_cell::sync::Lazy;
11use parking_lot::RwLock;
12use snafu::ResultExt;
13use svod_device::Result as DeviceResult;
14use svod_device::device::Device;
15use svod_device::registry::DeviceRegistry;
16use svod_dtype::DeviceSpec;
17
18use crate::error::{DeviceSnafu, Result, UnsupportedDeviceSnafu};
19
20/// Factory function that creates a Device for a given DeviceSpec.
21///
22/// The factory receives both the device specification and the allocator registry,
23/// allowing it to obtain the correct allocator for the device.
24///
25/// Returns `DeviceResult<Device>` (from svod_device) since device creation
26/// errors come from the device crate.
27pub type DeviceFactory = Arc<dyn Fn(&DeviceSpec, &DeviceRegistry) -> DeviceResult<Device> + Send + Sync>;
28
29/// Registry for full Device objects with caching and factory registration.
30///
31/// # Thread Safety
32///
33/// This registry uses `parking_lot::RwLock` for efficient concurrent access:
34/// - Multiple readers can access cached devices simultaneously
35/// - Writers acquire exclusive lock only when creating new devices
36/// - Double-checked locking pattern prevents redundant device creation
37///
38/// # Example
39///
40/// ```ignore
41/// // Get a device (creates if not cached)
42/// let alloc_registry = svod_device::registry::registry();
43/// let device = DEVICE_FACTORIES.device(&DeviceSpec::Cpu, alloc_registry)?;
44///
45/// // Register a custom factory
46/// DEVICE_FACTORIES.register_factory("CUSTOM", Arc::new(|spec, reg| {
47///     // Create custom device...
48/// }));
49/// ```
50pub struct DeviceFactoryRegistry {
51    /// Cached device instances (DeviceSpec -> Device)
52    devices: RwLock<HashMap<DeviceSpec, Arc<Device>>>,
53    /// Registered factories (device type string -> factory function)
54    factories: RwLock<HashMap<String, DeviceFactory>>,
55}
56
57impl DeviceFactoryRegistry {
58    /// Create a new registry with built-in device factories registered.
59    pub fn new() -> Self {
60        let registry = Self { devices: RwLock::new(HashMap::new()), factories: RwLock::new(HashMap::new()) };
61
62        // Register built-in CPU factory
63        registry
64            .register_factory("CPU", Arc::new(|_spec, alloc_reg| crate::devices::cpu::create_cpu_device(alloc_reg)));
65
66        // Future: Register CUDA, Metal, WebGPU factories when implemented
67        // registry.register_factory("CUDA", Arc::new(|spec, reg| create_cuda_device(spec, reg)));
68
69        registry
70    }
71
72    /// Register a device factory for a device type.
73    ///
74    /// The device type string is case-insensitive (converted to uppercase).
75    /// This allows plugins or extensions to register new device types at runtime.
76    ///
77    /// # Arguments
78    ///
79    /// * `device_type` - Device type identifier (e.g., "CPU", "CUDA", "METAL")
80    /// * `factory` - Factory function that creates Device instances
81    pub fn register_factory(&self, device_type: &str, factory: DeviceFactory) {
82        self.factories.write().insert(device_type.to_uppercase(), factory);
83    }
84
85    /// Get or create a Device for the given specification.
86    ///
87    /// This method uses double-checked locking for efficiency:
88    /// 1. Fast path: Read lock to check cache
89    /// 2. Slow path: Write lock to create and cache new device
90    ///
91    /// # Arguments
92    ///
93    /// * `spec` - Device specification (e.g., `DeviceSpec::Cpu`)
94    /// * `alloc_registry` - Allocator registry for obtaining device allocators
95    ///
96    /// # Returns
97    ///
98    /// Arc-wrapped Device for the specification, or error if device type unsupported.
99    pub fn device(&self, spec: &DeviceSpec, alloc_registry: &DeviceRegistry) -> Result<Arc<Device>> {
100        // Fast path: read lock to check cache
101        if let Some(dev) = self.devices.read().get(spec) {
102            return Ok(Arc::clone(dev));
103        }
104
105        // Slow path: write lock to create
106        let mut devices = self.devices.write();
107
108        // Double-check after acquiring write lock (another thread may have created it)
109        if let Some(dev) = devices.get(spec) {
110            return Ok(Arc::clone(dev));
111        }
112
113        // Look up factory for this device type
114        let device_type = spec.base_type();
115        let factory = self
116            .factories
117            .read()
118            .get(device_type)
119            .cloned()
120            .ok_or_else(|| UnsupportedDeviceSnafu { device: device_type.to_string() }.build())?;
121
122        // Create device via factory
123        let device = factory(spec, alloc_registry).context(DeviceSnafu)?;
124        let arc = Arc::new(device);
125        devices.insert(spec.clone(), Arc::clone(&arc));
126        Ok(arc)
127    }
128}
129
130impl Default for DeviceFactoryRegistry {
131    fn default() -> Self {
132        Self::new()
133    }
134}
135
136/// Global device factory registry.
137///
138/// This static instance is lazily initialized on first access,
139/// with built-in device factories automatically registered.
140///
141/// # Example
142///
143/// ```ignore
144/// let device = svod_runtime::DEVICE_FACTORIES
145///     .device(&DeviceSpec::Cpu, svod_device::registry::registry())?;
146/// ```
147pub static DEVICE_FACTORIES: Lazy<DeviceFactoryRegistry> = Lazy::new(DeviceFactoryRegistry::new);