vyre_driver_wgpu/
lib.rs

1#![allow(
2    unstable_name_collisions,
3    clippy::field_reassign_with_default,
4    clippy::double_must_use,
5    clippy::type_complexity,
6    clippy::missing_errors_doc,
7    clippy::too_many_arguments,
8    clippy::manual_clamp,
9    clippy::module_inception,
10    clippy::empty_line_after_doc_comments,
11    clippy::let_and_return,
12    clippy::missing_safety_doc
13)]
14#![deny(unsafe_code)]
15#![deny(missing_docs)]
16
17//! # vyre-wgpu  -  wgpu backend for the vyre GPU compute specification
18
19mod allocation;
20mod async_dispatch;
21mod backend_impl;
22pub mod buffer;
23mod capabilities;
24mod descriptor_mapping;
25mod device_buffer;
26pub mod emit;
27pub mod engine;
28mod executable_api;
29pub mod ext;
30pub mod megakernel;
31mod numeric;
32mod padded_upload;
33#[cfg(feature = "parity-testing")]
34mod parity_probe;
35pub mod pipeline;
36mod resident_dispatch;
37mod resident_download;
38mod resident_resource;
39mod resident_upload;
40pub mod runtime;
41pub mod spirv_backend;
42mod staging_reserve;
43mod stats;
44mod thread_pool;
45mod wait_backoff;
46
47pub use device_buffer::{WgpuDeviceBuffer, WGPU_BACKEND_ID};
48pub use executable_api::WgpuIR;
49pub use stats::WgpuBackendStats;
50use std::hash::BuildHasherDefault;
51use std::sync::{atomic::AtomicBool, Arc};
52use vyre_driver::shape_prediction::{ShapeFingerprint, ShapeHistory};
53use vyre_driver::DispatchConfig;
54use vyre_foundation::ir::DataType;
55use vyre_foundation::ir::Program;
56use vyre_foundation::validate::BackendValidationCapabilities;
57
58#[derive(Clone, Debug)]
59pub(crate) enum AdapterRecoveryTarget {
60    Index(usize),
61    Identity(crate::runtime::device::AdapterIdentity),
62}
63
64/// A real wgpu backend for vyre.
65#[derive(Clone, Debug)]
66pub struct WgpuBackend {
67    pub(crate) adapter_info: wgpu::AdapterInfo,
68    pub(crate) adapter_name: Arc<str>,
69    pub(crate) device_limits: wgpu::Limits,
70    pub(crate) device_queue: Arc<arc_swap::ArcSwap<(wgpu::Device, wgpu::Queue)>>,
71    pub(crate) dispatch_arena: Arc<arc_swap::ArcSwap<DispatchArena>>,
72    pub(crate) persistent_pool: Arc<arc_swap::ArcSwap<crate::buffer::BufferPool>>,
73    pub(crate) pipeline_cache: Arc<runtime::cache::pipeline::LruPipelineCache>,
74    pub(crate) wgsl_dispatch_pipeline_cache: Arc<
75        dashmap::DashMap<
76            [u8; 32],
77            Arc<wgpu::ComputePipeline>,
78            BuildHasherDefault<rustc_hash::FxHasher>,
79        >,
80    >,
81    pub(crate) resident_pipeline_cache: Arc<
82        dashmap::DashMap<
83            (u64, u64, usize),
84            Arc<crate::pipeline::WgpuPipeline>,
85            BuildHasherDefault<rustc_hash::FxHasher>,
86        >,
87    >,
88    pub(crate) validation_cache: Arc<vyre_driver::validation::ValidationCache>,
89    pub(crate) shape_history: Arc<std::sync::Mutex<ShapeHistory>>,
90    pub(crate) predicted_programs: Arc<
91        dashmap::DashMap<
92            ShapeFingerprint,
93            PredictedProgram,
94            BuildHasherDefault<rustc_hash::FxHasher>,
95        >,
96    >,
97    pub(crate) bind_group_layout_cache: Arc<
98        dashmap::DashMap<
99            vyre_driver::BackendLayoutFingerprint,
100            Arc<[Arc<wgpu::BindGroupLayout>]>,
101            BuildHasherDefault<rustc_hash::FxHasher>,
102        >,
103    >,
104    /// Live resident buffers this backend instance is keeping alive.
105    ///
106    /// The WGPU resident namespace itself is process-wide (see
107    /// `buffer::handle`), so a handle stays resolvable for as long as its
108    /// buffer lives; this table is the strong reference that decides that.
109    pub(crate) resident_handles: Arc<
110        dashmap::DashMap<
111            vyre_driver::ResidentHandle,
112            crate::buffer::GpuBufferHandle,
113            BuildHasherDefault<rustc_hash::FxHasher>,
114        >,
115    >,
116    pub(crate) device_lost: Arc<AtomicBool>,
117    pub(crate) enabled_features: crate::runtime::device::EnabledFeatures,
118    pub(crate) recovery_target: AdapterRecoveryTarget,
119}
120
121#[derive(Clone, Debug)]
122pub(crate) struct PredictedProgram {
123    pub(crate) program: Arc<Program>,
124    pub(crate) config: DispatchConfig,
125}
126
127/// Backend-owned dispatch buffer arena.
128#[derive(Clone)]
129pub struct DispatchArena {
130    pool: crate::buffer::BufferPool,
131    readback_rings: Arc<runtime::readback_ring::ReadbackRingSet>,
132}
133
134impl std::fmt::Debug for DispatchArena {
135    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136        formatter.write_str("DispatchArena { pool: size-classed }")
137    }
138}
139
140impl DispatchArena {
141    /// Create a dispatch arena backed by the canonical persistent buffer pool.
142    #[must_use]
143    #[inline]
144    pub fn new(device: wgpu::Device, queue: wgpu::Queue, config: &DispatchConfig) -> Self {
145        Self {
146            pool: crate::buffer::BufferPool::new(device, queue, config),
147            readback_rings: Arc::new(runtime::readback_ring::ReadbackRingSet::new()),
148        }
149    }
150
151    pub(crate) fn pool(&self) -> &crate::buffer::BufferPool {
152        &self.pool
153    }
154
155    pub(crate) fn readback_rings(&self) -> &Arc<runtime::readback_ring::ReadbackRingSet> {
156        &self.readback_rings
157    }
158}
159
160impl BackendValidationCapabilities for WgpuBackend {
161    fn backend_name(&self) -> &'static str {
162        "wgpu"
163    }
164
165    fn supports_cast_target(&self, target: &DataType) -> bool {
166        matches!(
167            target,
168            DataType::Bool
169                | DataType::U8
170                | DataType::U16
171                | DataType::U32
172                | DataType::U64
173                | DataType::I8
174                | DataType::I16
175                | DataType::I32
176                // I64 is backed by the SAME `vec2<u32>` emulation as U64 (see
177                // `binding_helpers`/`setup` buffer lowering and the op_dispatch
178                // cast path, which treat `U64 | I64` identically and sign-extend
179                // the high word for a signed source). Omitting it here falsely
180                // rejected `i32 -> i64` casts the emitter and hardware handle
181                // correctly, a capability/coherence gap, NOT a real limit
182                // (verified on the live 5090 by `widening_cast_64_parity`).
183                | DataType::I64
184                | DataType::F32
185                | DataType::Vec2U32
186                | DataType::Vec4U32
187        )
188    }
189
190    fn supports_subgroup_ops(&self) -> bool {
191        self.device_profile().supports_subgroup_ops
192    }
193
194    fn supports_indirect_dispatch(&self) -> bool {
195        self.device_profile().supports_indirect_dispatch
196    }
197
198    fn supports_specialization_constants(&self) -> bool {
199        self.device_profile().supports_specialization_constants
200    }
201
202    fn supports_distributed_collectives(&self) -> bool {
203        self.device_profile().supports_distributed_collectives
204    }
205}
206
207inventory::submit! {
208    vyre_driver::BackendRegistration {
209        id: "wgpu",
210        factory: || WgpuBackend::acquire().map(|backend| {
211            Box::new(backend) as Box<dyn vyre_driver::VyreBackend>
212        }),
213        supported_ops: vyre_driver::backend::validation::default_supported_ops_with_trap,
214    }
215}
216
217inventory::submit! {
218    vyre_driver::backend::BackendPrecedence {
219        id: "wgpu",
220        rank: 30,
221    }
222}
223
224inventory::submit! {
225    vyre_driver::backend::BackendCapability {
226        id: "wgpu",
227        dispatches: true,
228    }
229}
230
231impl vyre_driver::backend::private::Sealed for crate::pipeline::WgpuPipeline {}
232impl vyre_driver::backend::private::Sealed for WgpuBackend {}
vyre_driver_wgpu/lib.rs

vyre_driver_wgpu/
lib.rs
