forge/
device.rs

1use std::sync::Arc;
2
3use crate::backend::wgpu::{DispatchScope, 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    /// Batch every dispatch made while the returned guard is alive into one
30    /// command buffer and one submit. `None` on the CPU backend, which has no
31    /// queue to batch.
32    ///
33    /// Hold it for the span of a logical step — a decode step issues ~100
34    /// kernels, and a submit each leaves the GPU waiting on the CPU between
35    /// every one of them. Readbacks flush automatically, so a scope can wrap
36    /// code that reads results without special care.
37    #[must_use = "the batching ends when the guard drops"]
38    pub fn dispatch_scope(&self) -> Option<DispatchScope> {
39        match self {
40            Device::Cpu => None,
41            Device::Wgpu(ctx) => Some(ctx.scope()),
42        }
43    }
44
45    /// Human-readable adapter description.
46    pub fn describe(&self) -> String {
47        match self {
48            Device::Cpu => "CPU (reference)".to_string(),
49            Device::Wgpu(ctx) => format!(
50                "WebGPU: {} ({:?})",
51                ctx.adapter_info.name, ctx.adapter_info.backend
52            ),
53        }
54    }
55
56    /// Largest single buffer binding this device supports, in bytes. Large
57    /// weights (GPT-2's wte) are row-chunked to stay under this.
58    pub fn max_binding_bytes(&self) -> usize {
59        match self {
60            Device::Cpu => usize::MAX,
61            Device::Wgpu(ctx) => ctx.device.limits().max_storage_buffer_binding_size as usize,
62        }
63    }
64
65    /// Two devices are compatible when tensors on them can be combined.
66    pub fn same_as(&self, other: &Device) -> bool {
67        match (self, other) {
68            (Device::Cpu, Device::Cpu) => true,
69            (Device::Wgpu(a), Device::Wgpu(b)) => Arc::ptr_eq(a, b),
70            _ => false,
71        }
72    }
73}
forge/device.rs

forge/
device.rs
