Skip to main content

vyre_driver_spirv/
lib.rs

1//! SPIR-V backend for vyre.
2//!
3//! Reuses the shared naga::Module builder family and emits
4//! SPIR-V rather than WGSL. Intended for consumers targeting Vulkan-compatible
5//! compute pipelines (Vulkan 1.0+, Android NDK compute, desktop Vulkan).
6//!
7//! ```no_run
8//! use vyre_driver_spirv::SpirvBackend;
9//! // let module: naga::Module = ...;   // built via shared vyre naga emit
10//! // let words: Vec<u32> = SpirvBackend::emit_spv(&module).unwrap();
11//! ```
12//!
13//! The crate registers a `BackendRegistration` named `"spirv"` via inventory
14//! so `vyre::registered_backends()` enumerates it alongside other
15//! `photonic`.
16
17// Vulkan driver bindings (`ash::vk::*`) are inherently unsafe FFI;
18// every call site is the boundary between safe vyre code and the Vulkan
19// driver API. Allow `unsafe` here so the rest of the workspace can keep
20// `unsafe_code = "deny"` while this backend wraps ash properly with
21// per-call Safety: comments.
22#![allow(unsafe_code)]
23#![deny(rust_2018_idioms)]
24#![deny(missing_docs)]
25
26/// SPIR-V backend implementation. Contains `SpirvBackend` and the
27/// naga::back::spv glue that turns a `vyre::Program` into SPIR-V
28/// bytes.
29/// SpirV element.
30/// SpirV element.
31pub mod backend;
32/// Vulkan compute dispatch implementation.
33mod vulkan;
34/// The SPIR-V `VyreBackend` implementation.
35/// SpirV element.
36/// SpirV element.
37pub use backend::SpirvBackend;
38
39use std::sync::Arc;
40
41use vyre_driver::{BackendError, BackendRegistration, DispatchConfig, VyreBackend};
42use vyre_foundation::ir::Program;
43
44/// Stable backend identifier for conform certificates.
45pub const SPIRV_BACKEND_ID: &str = "spirv";
46
47/// Live Vulkan-backed SPIR-V backend.
48///
49/// Acquires a Vulkan device on construction and uses it to dispatch
50/// SPIR-V compute pipelines at runtime. If no Vulkan device is available,
51/// acquisition returns a structured error.
52#[derive(Debug, Clone)]
53pub struct SpirvBackendRegistration {
54    device: Arc<vulkan::VulkanDevice>,
55}
56
57impl SpirvBackendRegistration {
58    /// Acquire a new SPIR-V backend by probing for a Vulkan compute device.
59    ///
60    /// # Errors
61    /// Returns [`BackendError`] when no Vulkan loader or compatible GPU is found.
62    pub fn acquire() -> Result<Self, BackendError> {
63        let device = vulkan::VulkanDevice::acquire()?;
64        Ok(Self {
65            device: Arc::new(device),
66        })
67    }
68}
69
70impl vyre_driver::backend::private::Sealed for SpirvBackendRegistration {}
71
72fn spirv_device_buffer_unsupported() -> BackendError {
73    BackendError::UnsupportedFeature {
74        name: format!(
75            "{} requires native Vulkan-resident buffers; HostShimBuffer dispatch is forbidden",
76            vyre_driver::DEVICE_BUFFER_FEATURE
77        ),
78        backend: SPIRV_BACKEND_ID.to_string(),
79    }
80}
81
82impl VyreBackend for SpirvBackendRegistration {
83    fn id(&self) -> &'static str {
84        SPIRV_BACKEND_ID
85    }
86
87    fn version(&self) -> &'static str {
88        env!("CARGO_PKG_VERSION")
89    }
90
91    fn dispatch(
92        &self,
93        program: &Program,
94        inputs: &[Vec<u8>],
95        config: &DispatchConfig,
96    ) -> Result<Vec<Vec<u8>>, BackendError> {
97        let borrowed: Vec<&[u8]> = inputs.iter().map(Vec::as_slice).collect();
98        self.dispatch_borrowed(program, &borrowed, config)
99    }
100
101    fn dispatch_borrowed(
102        &self,
103        program: &Program,
104        inputs: &[&[u8]],
105        config: &DispatchConfig,
106    ) -> Result<Vec<Vec<u8>>, BackendError> {
107        let spv_words = SpirvBackend::program_to_spv(program).map_err(|e| {
108            BackendError::KernelCompileFailed {
109                backend: SPIRV_BACKEND_ID.to_string(),
110                compiler_message: format!("{e}. Fix: validate the Program IR before dispatch."),
111            }
112        })?;
113
114        // SAFETY: FFI to ash::vk. Handle lifetimes are documented at the
115        // surrounding VulkanDevice construction site; the Drop impl owns
116        // destruction.
117        unsafe { vulkan::dispatch_program(&self.device, program, &spv_words, inputs, config) }
118    }
119
120    fn allocate_device_buffer(
121        &self,
122        byte_len: usize,
123    ) -> Result<Box<dyn vyre_driver::DeviceBuffer>, BackendError> {
124        let _ = byte_len;
125        Err(spirv_device_buffer_unsupported())
126    }
127
128    fn upload_device_buffer(
129        &self,
130        buffer: &mut dyn vyre_driver::DeviceBuffer,
131        bytes: &[u8],
132    ) -> Result<(), BackendError> {
133        let _ = (buffer, bytes);
134        Err(spirv_device_buffer_unsupported())
135    }
136
137    fn download_device_buffer(
138        &self,
139        buffer: &dyn vyre_driver::DeviceBuffer,
140    ) -> Result<Vec<u8>, BackendError> {
141        let _ = buffer;
142        Err(spirv_device_buffer_unsupported())
143    }
144
145    fn free_device_buffer(
146        &self,
147        _buffer: Box<dyn vyre_driver::DeviceBuffer>,
148    ) -> Result<(), BackendError> {
149        // The SPIRV backend never allocates DeviceBuffers (allocate_device_buffer
150        // always returns Err). Reaching this function means the caller obtained a
151        // DeviceBuffer through some path the SPIRV backend does not own, which is a
152        // caller contract violation. Avoid silently dropping the buffer and then
153        // returning an error (which would imply the caller did something wrong after
154        // the buffer was already destroyed).
155        unreachable!(
156            "SPIRV backend never allocates DeviceBuffer; \
157             free_device_buffer cannot be called on a SPIRV-backend buffer. \
158             Fix: do not call free_device_buffer on a DeviceBuffer that was \
159             not produced by this backend's allocate_device_buffer."
160        )
161    }
162
163    fn dispatch_with_device_buffers(
164        &self,
165        program: &Program,
166        inputs: &[&dyn vyre_driver::DeviceBuffer],
167        outputs: &mut [&mut dyn vyre_driver::DeviceBuffer],
168        config: &DispatchConfig,
169    ) -> Result<(), BackendError> {
170        vyre_driver::default_dispatch_with_device_buffers(self, program, inputs, outputs, config)
171    }
172
173    fn max_workgroup_size(&self) -> [u32; 3] {
174        let max = self.device.properties.limits.max_compute_work_group_size;
175        [max[0], max[1], max[2]]
176    }
177
178    fn max_compute_workgroups_per_dimension(&self) -> u32 {
179        let max = self.device.properties.limits.max_compute_work_group_count;
180        max[0]
181    }
182
183    fn max_compute_invocations_per_workgroup(&self) -> u32 {
184        self.device
185            .properties
186            .limits
187            .max_compute_work_group_invocations
188    }
189
190    fn max_storage_buffer_bytes(&self) -> u64 {
191        self.device.properties.limits.max_storage_buffer_range as u64
192    }
193
194    fn supports_grid_sync(&self) -> bool {
195        false
196    }
197
198    fn supports_subgroup_ops(&self) -> bool {
199        false
200    }
201
202    fn supports_f16(&self) -> bool {
203        false
204    }
205
206    fn supports_bf16(&self) -> bool {
207        false
208    }
209
210    fn supports_tensor_cores(&self) -> bool {
211        false
212    }
213
214    fn supports_async_compute(&self) -> bool {
215        false
216    }
217
218    fn supports_indirect_dispatch(&self) -> bool {
219        false
220    }
221
222    fn device_profile(&self) -> vyre_driver::DeviceProfile {
223        let max_workgroup_size = self.max_workgroup_size();
224        vyre_driver::DeviceProfile {
225            backend: self.id(),
226            supports_subgroup_ops: false,
227            supports_indirect_dispatch: false,
228            supports_distributed_collectives: false,
229            supports_specialization_constants: false,
230            supports_f16: false,
231            supports_bf16: false,
232            supports_trap_propagation: false,
233            supports_tensor_cores: false,
234            has_mul_high: false,
235            has_dual_issue_fp32_int32: false,
236            has_subgroup_shuffle: false,
237            has_shared_memory: false,
238            max_native_int_width: 32,
239            max_workgroup_size,
240            max_invocations_per_workgroup: self.max_compute_invocations_per_workgroup(),
241            max_shared_memory_bytes: self.device.properties.limits.max_compute_shared_memory_size,
242            max_storage_buffer_binding_size: self.max_storage_buffer_bytes(),
243            subgroup_size: 0,
244            compute_units: 0,
245            regs_per_thread_max: 0,
246            l1_cache_bytes: 0,
247            l2_cache_bytes: 0,
248            mem_bw_gbps: 0,
249            timing_quality: vyre_driver::DeviceTimingQuality::HostOnly,
250            supports_device_timestamps: false,
251            supports_hardware_counters: false,
252            ideal_unroll_depth: 0,
253            ideal_vector_pack_bits: 0,
254            ideal_workgroup_tile: [0, 0, 0],
255            shared_memory_bank_count: 0,
256            shared_memory_bank_width_bytes: 0,
257        }
258    }
259}
260
261/// Factory for the inventory registration path.
262pub fn spirv_factory() -> Result<Box<dyn VyreBackend>, BackendError> {
263    SpirvBackendRegistration::acquire().map(|backend| Box::new(backend) as Box<dyn VyreBackend>)
264}
265
266/// Op-support set for the SPIR-V backend.
267///
268/// The SPIRV/naga path supports the same core IR ops as every other vyre backend
269/// (arithmetic, bitwise, control-flow, memory, collectives). Using `core_supported_ops`
270/// here keeps the inventory-registered op set consistent with what the router
271/// sees at runtime. A permanently-empty set here would cause the router to skip
272/// all SPIRV dispatch even when the op is supported, silently degrading to a
273/// lower-precedence backend.
274pub fn spirv_supported_ops() -> &'static std::collections::HashSet<vyre_foundation::ir::OpId> {
275    vyre_driver::backend::core_supported_ops()
276}
277
278inventory::submit! {
279    BackendRegistration {
280        id: SPIRV_BACKEND_ID,
281        factory: spirv_factory,
282        supported_ops: spirv_supported_ops,
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    /// Before the fix, spirv_supported_ops() returned a permanently-empty OnceLock.
291    /// The router would see zero supported ops and skip ALL SPIRV dispatch silently.
292    /// After the fix, it delegates to core_supported_ops() which is non-empty.
293    #[test]
294    fn test_spirv_supported_ops_is_not_empty() {
295        let ops = spirv_supported_ops();
296        assert!(
297            !ops.is_empty(),
298            "Fix: spirv_supported_ops must not be empty; a permanently-empty set causes \
299             the router to skip all SPIRV dispatch. Got {} ops.",
300            ops.len()
301        );
302    }
303}
304
305// V7-EXT-021: declare router precedence inline. SPIR-V is rank 30.
306inventory::submit! {
307    vyre_driver::backend::BackendPrecedence {
308        id: SPIRV_BACKEND_ID,
309        rank: 30,
310    }
311}
312
313inventory::submit! {
314    vyre_driver::backend::BackendCapability {
315        id: SPIRV_BACKEND_ID,
316        dispatches: true,
317    }
318}