Skip to main content

forge/
device.rs

1use std::sync::Arc;
2
3use crate::backend::wgpu::WgpuContext;
4use crate::error::Result;
5
6/// Where a tensor lives and where its ops execute.
7#[derive(Clone, Debug)]
8pub enum Device {
9    /// Reference implementation used for testing and verification.
10    Cpu,
11    /// Production backend: WebGPU via wgpu (Vulkan/Metal/D3D12/Browser).
12    Wgpu(Arc<WgpuContext>),
13}
14
15impl Device {
16    /// Create the default WebGPU device (highest-performance adapter).
17    /// Sync facade — native only; on wasm use [`Device::wgpu_async`].
18    #[cfg(not(target_arch = "wasm32"))]
19    pub fn wgpu() -> Result<Device> {
20        Ok(Device::Wgpu(WgpuContext::new()?))
21    }
22
23    /// Async device creation — the primary form (roadmap v4, pitfall 14);
24    /// works on native and wasm32/browser.
25    pub async fn wgpu_async() -> Result<Device> {
26        Ok(Device::Wgpu(WgpuContext::new_async().await?))
27    }
28
29    /// Human-readable adapter description.
30    pub fn describe(&self) -> String {
31        match self {
32            Device::Cpu => "CPU (reference)".to_string(),
33            Device::Wgpu(ctx) => format!(
34                "WebGPU: {} ({:?})",
35                ctx.adapter_info.name, ctx.adapter_info.backend
36            ),
37        }
38    }
39
40    /// Largest single buffer binding this device supports, in bytes. Large
41    /// weights (GPT-2's wte) are row-chunked to stay under this.
42    pub fn max_binding_bytes(&self) -> usize {
43        match self {
44            Device::Cpu => usize::MAX,
45            Device::Wgpu(ctx) => ctx.device.limits().max_storage_buffer_binding_size as usize,
46        }
47    }
48
49    /// Two devices are compatible when tensors on them can be combined.
50    pub fn same_as(&self, other: &Device) -> bool {
51        match (self, other) {
52            (Device::Cpu, Device::Cpu) => true,
53            (Device::Wgpu(a), Device::Wgpu(b)) => Arc::ptr_eq(a, b),
54            _ => false,
55        }
56    }
57}