1#![allow(unsafe_code)]
23#![deny(rust_2018_idioms)]
24#![deny(missing_docs)]
25
26pub mod backend;
32mod vulkan;
34pub use backend::SpirvBackend;
38
39use std::sync::Arc;
40
41use vyre_driver::{BackendError, BackendRegistration, DispatchConfig, VyreBackend};
42use vyre_foundation::ir::Program;
43
44pub const SPIRV_BACKEND_ID: &str = "spirv";
46
47#[derive(Debug, Clone)]
53pub struct SpirvBackendRegistration {
54 device: Arc<vulkan::VulkanDevice>,
55}
56
57impl SpirvBackendRegistration {
58 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 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 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
261pub fn spirv_factory() -> Result<Box<dyn VyreBackend>, BackendError> {
263 SpirvBackendRegistration::acquire().map(|backend| Box::new(backend) as Box<dyn VyreBackend>)
264}
265
266pub 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 #[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
305inventory::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}