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

vyre_driver_wgpu/
lib.rs
