ferrotensor_backend/
context.rs

1//! GPU device setup.
2
3use crate::pool::BufferPool;
4use std::sync::Arc;
5use wgpu::{Adapter, Device, Instance, Queue};
6
7/// A GPU device, its queue and a buffer pool shared by tensors.
8#[derive(Debug)]
9pub struct WgpuContext {
10    /// The wgpu instance the adapter came from.
11    pub instance: Instance,
12    /// The physical GPU.
13    pub adapter: Adapter,
14    /// Logical device; requested with the adapter's full limits.
15    pub device: Arc<Device>,
16    /// Queue that every kernel dispatch is submitted to.
17    pub queue: Arc<Queue>,
18    /// Pool that [`UnifiedTensor::empty`](crate::UnifiedTensor::empty) draws from.
19    pub pool: Arc<BufferPool>,
20}
21
22impl WgpuContext {
23    /// Opens the high-performance adapter.
24    ///
25    /// # Panics
26    /// Panics if no adapter or device is available.
27    pub async fn new() -> Self {
28        Self::new_with_options(wgpu::PowerPreference::HighPerformance, false).await
29    }
30
31    /// Opens an adapter with the given power preference, optionally forcing
32    /// the software fallback adapter.
33    ///
34    /// # Panics
35    /// Panics if no adapter or device is available.
36    pub async fn new_with_options(
37        power_preference: wgpu::PowerPreference,
38        force_fallback_adapter: bool,
39    ) -> Self {
40        let instance = Instance::default();
41        let adapter = instance
42            .request_adapter(&wgpu::RequestAdapterOptions {
43                power_preference,
44                force_fallback_adapter,
45                compatible_surface: None,
46            })
47            .await
48            .expect("Failed to find adapter");
49
50        let (device, queue) = adapter
51            .request_device(
52                &wgpu::DeviceDescriptor {
53                    label: Some("Ferrotensor Device"),
54                    required_features: wgpu::Features::empty(),
55                    // Use what the hardware supports rather than the WebGPU
56                    // defaults (e.g. 8 storage buffers per stage, 128 MiB
57                    // bindings), which are too small for model weights.
58                    required_limits: adapter.limits(),
59                },
60                None,
61            )
62            .await
63            .expect("Failed to create device");
64        let pool = Arc::new(BufferPool::new(
65            wgpu::BufferUsages::STORAGE
66                | wgpu::BufferUsages::COPY_DST
67                | wgpu::BufferUsages::COPY_SRC,
68        ));
69        Self {
70            instance,
71            adapter,
72            device: Arc::new(device),
73            queue: Arc::new(queue),
74            pool,
75        }
76    }
77
78    /// Adapter name and backend, e.g. `"Apple M2 (Metal)"`.
79    pub fn adapter_info(&self) -> String {
80        let info = self.adapter.get_info();
81        format!("{} ({:?})", info.name, info.backend)
82    }
83}
ferrotensor_backend/context.rs

ferrotensor_backend/
context.rs
