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    pub(crate) resident_handles: Arc<
105        dashmap::DashMap<
106            u64,
107            crate::buffer::GpuBufferHandle,
108            BuildHasherDefault<rustc_hash::FxHasher>,
109        >,
110    >,
111    pub(crate) device_lost: Arc<AtomicBool>,
112    pub(crate) enabled_features: crate::runtime::device::EnabledFeatures,
113    pub(crate) recovery_target: AdapterRecoveryTarget,
114}
115
116#[derive(Clone, Debug)]
117pub(crate) struct PredictedProgram {
118    pub(crate) program: Arc<Program>,
119    pub(crate) config: DispatchConfig,
120}
121
122/// Backend-owned dispatch buffer arena.
123#[derive(Clone)]
124pub struct DispatchArena {
125    pool: crate::buffer::BufferPool,
126    readback_rings: Arc<runtime::readback_ring::ReadbackRingSet>,
127}
128
129impl std::fmt::Debug for DispatchArena {
130    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131        formatter.write_str("DispatchArena { pool: size-classed }")
132    }
133}
134
135impl DispatchArena {
136    /// Create a dispatch arena backed by the canonical persistent buffer pool.
137    #[must_use]
138    #[inline]
139    pub fn new(device: wgpu::Device, queue: wgpu::Queue, config: &DispatchConfig) -> Self {
140        Self {
141            pool: crate::buffer::BufferPool::new(device, queue, config),
142            readback_rings: Arc::new(runtime::readback_ring::ReadbackRingSet::new()),
143        }
144    }
145
146    pub(crate) fn pool(&self) -> &crate::buffer::BufferPool {
147        &self.pool
148    }
149
150    pub(crate) fn readback_rings(&self) -> &Arc<runtime::readback_ring::ReadbackRingSet> {
151        &self.readback_rings
152    }
153}
154
155impl BackendValidationCapabilities for WgpuBackend {
156    fn backend_name(&self) -> &'static str {
157        "wgpu"
158    }
159
160    fn supports_cast_target(&self, target: &DataType) -> bool {
161        matches!(
162            target,
163            DataType::Bool
164                | DataType::U8
165                | DataType::U16
166                | DataType::U32
167                | DataType::U64
168                | DataType::I8
169                | DataType::I16
170                | DataType::I32
171                | DataType::F32
172                | DataType::Vec2U32
173                | DataType::Vec4U32
174        )
175    }
176
177    fn supports_subgroup_ops(&self) -> bool {
178        self.device_profile().supports_subgroup_ops
179    }
180
181    fn supports_indirect_dispatch(&self) -> bool {
182        self.device_profile().supports_indirect_dispatch
183    }
184
185    fn supports_specialization_constants(&self) -> bool {
186        self.device_profile().supports_specialization_constants
187    }
188
189    fn supports_distributed_collectives(&self) -> bool {
190        self.device_profile().supports_distributed_collectives
191    }
192}
193
194inventory::submit! {
195    vyre_driver::BackendRegistration {
196        id: "wgpu",
197        factory: || WgpuBackend::acquire().map(|backend| {
198            Box::new(backend) as Box<dyn vyre_driver::VyreBackend>
199        }),
200        supported_ops: vyre_driver::backend::validation::default_supported_ops_with_trap,
201    }
202}
203
204inventory::submit! {
205    vyre_driver::backend::BackendPrecedence {
206        id: "wgpu",
207        rank: 30,
208    }
209}
210
211inventory::submit! {
212    vyre_driver::backend::BackendCapability {
213        id: "wgpu",
214        dispatches: true,
215    }
216}
217
218impl vyre_driver::backend::private::Sealed for crate::pipeline::WgpuPipeline {}
219impl vyre_driver::backend::private::Sealed for WgpuBackend {}
vyre_driver_wgpu/lib.rs

vyre_driver_wgpu/
lib.rs
