GPU-accelerated UI toolkit (Vulkan)
git clone https://git.lucas.co/cce-ui.git
src/vk/compute.rs (33.1K)
1 //! The compute-job API: upload buffers, dispatch a WGSL kernel, read back.
2 //!
3 //! The renderer already had everything a compute job needs — naga compiles
4 //! WGSL to SPIR-V at runtime, the path tracer builds compute pipelines over
5 //! storage buffers, and `RtOffscreen` runs a headless device with no window
6 //! — but all of it was internal to the RT pass and read back only an image.
7 //! This module is that machinery with a general face: a [`ComputeDevice`]
8 //! owns a headless [`VkCore`], and [`ComputeDevice::run`] takes a
9 //! [`Kernel`] and a list of [`Binding`]s, uploads them, dispatches, waits,
10 //! and copies every read-write binding back into the caller's slice.
11 //!
12 //! Designed for cce-designer's Phase 7 step 4 (see its `shapeshifter.md`):
13 //! the per-point solver operators — relax, diffuse, collide — written once
14 //! in WGSL over the columnar attribute arrays a `Detail` already keeps, with
15 //! the CPU evaluator as the reference each is held to. The shape of the API
16 //! follows from that use: the caller has arrays in memory and wants them
17 //! transformed, so buffers are HOST-VISIBLE and mapped, upload and readback
18 //! are memcpys through the mapping, and there is no staging copy. On an
19 //! integrated GPU that is the fastest path there is; on a discrete one it is
20 //! correct and simple, and a device-local tier can be added behind the same
21 //! API if a workload ever asks for it.
22 //!
23 //! Every failure is an `Err(String)`, never a panic, because a kernel may be
24 //! user-authored: a WGSL error comes back with naga's own diagnostic, a
25 //! missing entry point names what the module does offer, and a device
26 //! without Vulkan reports as such from [`ComputeDevice::new`].
27 //!
28 //! Not `Send`: it owns a device and a command buffer. Make one per thread
29 //! that computes, and keep it — pipelines cache by source and entry point,
30 //! and buffers are reused across runs when they fit.
31
32 use super::core::VkCore;
33 use super::renderer::{create_cpu_buffer, destroy_cpu_buffer, AllocatedBuffer};
34 use ash::vk;
35 use std::collections::HashMap;
36 use std::ffi::CString;
37
38 /// A compute shader: WGSL source and the `@compute` entry point to run.
39 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
40 pub struct Kernel {
41 pub source: String,
42 pub entry: String,
43 }
44
45 impl Kernel {
46 pub fn new(source: impl Into<String>, entry: impl Into<String>) -> Self {
47 Kernel { source: source.into(), entry: entry.into() }
48 }
49 }
50
51 /// How a binding is declared to the shader, in `@binding(i)` order.
52 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
53 pub enum BindKind {
54 /// `var<storage, read>` or `var<storage, read_write>`.
55 Storage,
56 /// `var<uniform>`: a small parameter block, 16-byte layout rules apply.
57 Uniform,
58 }
59
60 /// One buffer of a job, bound at `@group(0) @binding(i)` for its index in
61 /// the list handed to [`ComputeDevice::run`].
62 pub enum Binding<'a> {
63 /// Read-write storage: uploaded before the dispatch and READ BACK into
64 /// the same slice after it.
65 Storage(&'a mut [u8]),
66 /// Read-only storage: uploaded, never read back.
67 Input(&'a [u8]),
68 /// A uniform block: uploaded, never read back.
69 Uniform(&'a [u8]),
70 }
71
72 impl<'a> Binding<'a> {
73 /// A read-write binding over a typed slice (`&mut [f32]`, `&mut [[f32; 3]]`, …).
74 pub fn rw<T: bytemuck::Pod>(data: &'a mut [T]) -> Self {
75 Binding::Storage(bytemuck::cast_slice_mut(data))
76 }
77
78 /// A read-only storage binding over a typed slice.
79 pub fn input<T: bytemuck::Pod>(data: &'a [T]) -> Self {
80 Binding::Input(bytemuck::cast_slice(data))
81 }
82
83 /// A uniform binding over one `Pod` struct.
84 pub fn uniform<T: bytemuck::Pod>(value: &'a T) -> Self {
85 Binding::Uniform(bytemuck::bytes_of(value))
86 }
87
88 fn kind(&self) -> BindKind {
89 match self {
90 Binding::Storage(_) | Binding::Input(_) => BindKind::Storage,
91 Binding::Uniform(_) => BindKind::Uniform,
92 }
93 }
94
95 fn bytes(&self) -> &[u8] {
96 match self {
97 Binding::Storage(b) => b,
98 Binding::Input(b) => b,
99 Binding::Uniform(b) => b,
100 }
101 }
102 }
103
104 /// Workgroups needed to cover `items` at `per_group` invocations each — the
105 /// `@workgroup_size` of the entry point, which [`ComputeDevice::run_over`]
106 /// reads for you.
107 pub fn workgroups(items: u32, per_group: u32) -> u32 {
108 items.div_ceil(per_group.max(1)).max(1)
109 }
110
111 #[derive(Clone, PartialEq, Eq, Hash)]
112 struct PipelineKey {
113 kernel: Kernel,
114 kinds: Vec<BindKind>,
115 }
116
117 struct Pipeline {
118 set_layout: vk::DescriptorSetLayout,
119 layout: vk::PipelineLayout,
120 module: vk::ShaderModule,
121 pipeline: vk::Pipeline,
122 workgroup_size: [u32; 3],
123 }
124
125 /// Storage bindings are bound whole, so a buffer's size has to be a multiple
126 /// of the widest element stride a shader may declare; 16 covers `vec4<f32>`.
127 const BUFFER_ALIGN: usize = 16;
128
129 /// A headless device that runs compute jobs. See the module docs.
130 pub struct ComputeDevice {
131 pipelines: HashMap<PipelineKey, Pipeline>,
132 /// One buffer per binding index, grown when a job needs more room.
133 slots: Vec<AllocatedBuffer>,
134 descriptor_pool: vk::DescriptorPool,
135 cmd: vk::CommandBuffer,
136 fence: vk::Fence,
137 /// Declared last: everything above is destroyed before the device.
138 core: VkCore,
139 }
140
141 /// The most bindings one job may carry (the descriptor pool is sized to it).
142 pub const MAX_BINDINGS: usize = 16;
143
144 impl ComputeDevice {
145 /// A device on the machine's preferred GPU (`CCE_VK_DEVICE` steers it,
146 /// as for every renderer). `Err` when there is no usable Vulkan at all.
147 pub fn new() -> Result<Self, String> {
148 // VkCore reports an absent driver by panicking, as a renderer with no
149 // window to draw into has nothing better to do; a compute consumer
150 // has a CPU path to fall back to, so the panic is caught here.
151 let core = std::panic::catch_unwind(VkCore::new_headless).map_err(|e| {
152 let msg = e
153 .downcast_ref::<String>()
154 .cloned()
155 .or_else(|| e.downcast_ref::<&str>().map(|s| s.to_string()))
156 .unwrap_or_else(|| "unknown".to_string());
157 format!("no Vulkan compute device: {msg}")
158 })?;
159 let device = core.device.clone();
160 unsafe {
161 let cmd = device
162 .allocate_command_buffers(
163 &vk::CommandBufferAllocateInfo::default()
164 .command_pool(core.command_pool)
165 .level(vk::CommandBufferLevel::PRIMARY)
166 .command_buffer_count(1),
167 )
168 .map_err(|e| format!("command buffer: {e}"))?[0];
169 let fence = device
170 .create_fence(&vk::FenceCreateInfo::default(), None)
171 .map_err(|e| format!("fence: {e}"))?;
172 let pool_sizes = [
173 vk::DescriptorPoolSize::default()
174 .ty(vk::DescriptorType::STORAGE_BUFFER)
175 .descriptor_count(2 * MAX_BINDINGS as u32),
176 vk::DescriptorPoolSize::default()
177 .ty(vk::DescriptorType::UNIFORM_BUFFER)
178 .descriptor_count(2 * MAX_BINDINGS as u32),
179 ];
180 let descriptor_pool = device
181 .create_descriptor_pool(
182 &vk::DescriptorPoolCreateInfo::default().max_sets(2).pool_sizes(&pool_sizes),
183 None,
184 )
185 .map_err(|e| format!("descriptor pool: {e}"))?;
186 Ok(ComputeDevice {
187 pipelines: HashMap::new(),
188 slots: Vec::new(),
189 descriptor_pool,
190 cmd,
191 fence,
192 core,
193 })
194 }
195 }
196
197 /// The physical device's name, for a log line or a status readout.
198 pub fn device_name(&self) -> String {
199 unsafe {
200 let props = self.core.instance.get_physical_device_properties(self.core.physical_device);
201 std::ffi::CStr::from_ptr(props.device_name.as_ptr()).to_string_lossy().into_owned()
202 }
203 }
204
205 /// The entry point's `@workgroup_size`, compiling the kernel if needed.
206 pub fn workgroup_size(&mut self, kernel: &Kernel, kinds: &[BindKind]) -> Result<[u32; 3], String> {
207 let key = PipelineKey { kernel: kernel.clone(), kinds: kinds.to_vec() };
208 Ok(self.pipeline(&key)?.workgroup_size)
209 }
210
211 /// Run the kernel over `items` invocations along x — the common case, a
212 /// job that is one invocation per element — computing the workgroup
213 /// count from the entry point's own `@workgroup_size`. A kernel should
214 /// still guard `id.x < arrayLength(...)`: the last group is padded.
215 pub fn run_over(&mut self, kernel: &Kernel, bindings: &mut [Binding<'_>], items: u32) -> Result<(), String> {
216 let kinds: Vec<BindKind> = bindings.iter().map(Binding::kind).collect();
217 let wg = self.workgroup_size(kernel, &kinds)?;
218 self.run(kernel, bindings, [workgroups(items, wg[0]), 1, 1])
219 }
220
221 /// Upload every binding, dispatch `groups` workgroups of the kernel, wait
222 /// for the GPU, and read every [`Binding::Storage`] back into its slice.
223 pub fn run(&mut self, kernel: &Kernel, bindings: &mut [Binding<'_>], groups: [u32; 3]) -> Result<(), String> {
224 self.execute(kernel, bindings, groups, 1, None)
225 }
226
227 /// [`run_passes`](Self::run_passes) with the dispatch sized from the entry
228 /// point's `@workgroup_size` over `items`, like [`run_over`](Self::run_over).
229 pub fn run_passes_over(
230 &mut self,
231 kernel: &Kernel,
232 bindings: &mut [Binding<'_>],
233 items: u32,
234 passes: u32,
235 ping_pong: Option<(usize, usize)>,
236 ) -> Result<(), String> {
237 let kinds: Vec<BindKind> = bindings.iter().map(Binding::kind).collect();
238 let wg = self.workgroup_size(kernel, &kinds)?;
239 self.execute(kernel, bindings, [workgroups(items, wg[0]), 1, 1], passes, ping_pong)
240 }
241
242 /// `passes` dispatches of the kernel in ONE submission — uploaded once,
243 /// a memory barrier between passes, waited on once, read back once —
244 /// which is what an iterative solve needs: measured on an integrated
245 /// GPU, a pass submitted on its own costs about half a millisecond of
246 /// round trip whatever its size, and sixteen of those lose to the CPU
247 /// at every mesh size a designer works at.
248 ///
249 /// `ping_pong = Some((a, b))` makes passes alternate the roles of two
250 /// bindings: `a` must be a [`Binding::Input`] (the first pass reads it)
251 /// and `b` a [`Binding::Storage`] of the same length (the first pass
252 /// writes it); the second pass reads `b` and writes `a`'s buffer, and so
253 /// on. Whichever buffer the LAST pass wrote is read back into `b`'s
254 /// slice, so the caller always finds the result where it bound the
255 /// output. A Jacobi solve is exactly this shape.
256 pub fn run_passes(
257 &mut self,
258 kernel: &Kernel,
259 bindings: &mut [Binding<'_>],
260 groups: [u32; 3],
261 passes: u32,
262 ping_pong: Option<(usize, usize)>,
263 ) -> Result<(), String> {
264 self.execute(kernel, bindings, groups, passes, ping_pong)
265 }
266
267 fn execute(
268 &mut self,
269 kernel: &Kernel,
270 bindings: &mut [Binding<'_>],
271 groups: [u32; 3],
272 passes: u32,
273 ping_pong: Option<(usize, usize)>,
274 ) -> Result<(), String> {
275 if bindings.len() > MAX_BINDINGS {
276 return Err(format!("{} bindings; a job may carry at most {MAX_BINDINGS}", bindings.len()));
277 }
278 if groups.iter().any(|&g| g == 0) {
279 return Err(format!("workgroup count {groups:?} has a zero"));
280 }
281 if passes == 0 {
282 return Err("a job needs at least one pass".to_string());
283 }
284 if let Some((a, b)) = ping_pong {
285 if a == b || a >= bindings.len() || b >= bindings.len() {
286 return Err(format!("ping-pong pair ({a}, {b}) does not name two distinct bindings of {}", bindings.len()));
287 }
288 if !matches!(bindings[a], Binding::Input(_)) {
289 return Err(format!("ping-pong binding {a} must be a read-only Input: it is where the first pass reads"));
290 }
291 if !matches!(bindings[b], Binding::Storage(_)) {
292 return Err(format!("ping-pong binding {b} must be a read-write Storage: it is where the result lands"));
293 }
294 if bindings[a].bytes().len() != bindings[b].bytes().len() {
295 return Err(format!(
296 "ping-pong bindings {a} and {b} differ in length ({} vs {} bytes)",
297 bindings[a].bytes().len(),
298 bindings[b].bytes().len()
299 ));
300 }
301 }
302 let kinds: Vec<BindKind> = bindings.iter().map(Binding::kind).collect();
303 let key = PipelineKey { kernel: kernel.clone(), kinds };
304 let (pipeline, layout, set_layout) = {
305 let p = self.pipeline(&key)?;
306 (p.pipeline, p.layout, p.set_layout)
307 };
308
309 // Buffers: one per binding index, reused when big enough. Uploads are
310 // memcpys through the persistent mapping.
311 let device = self.core.device.clone();
312 let mut sizes = Vec::with_capacity(bindings.len());
313 for (i, b) in bindings.iter().enumerate() {
314 let bytes = b.bytes();
315 let padded = bytes.len().max(BUFFER_ALIGN).div_ceil(BUFFER_ALIGN) * BUFFER_ALIGN;
316 if b.kind() == BindKind::Uniform {
317 let cap = unsafe {
318 self.core.instance.get_physical_device_properties(self.core.physical_device).limits.max_uniform_buffer_range
319 } as usize;
320 if padded > cap {
321 return Err(format!("binding {i}: a uniform block of {} bytes exceeds the device's {cap}", bytes.len()));
322 }
323 }
324 if i >= self.slots.len() {
325 self.slots.push(AllocatedBuffer::null());
326 }
327 if (self.slots[i].size as usize) < padded {
328 let allocator = self.core.allocator.as_mut().unwrap();
329 destroy_cpu_buffer(&device, allocator, &mut self.slots[i]);
330 self.slots[i] = create_cpu_buffer(
331 &device,
332 allocator,
333 padded as vk::DeviceSize,
334 vk::BufferUsageFlags::STORAGE_BUFFER | vk::BufferUsageFlags::UNIFORM_BUFFER,
335 "compute-binding",
336 );
337 }
338 let mapped = self.slots[i]
339 .allocation
340 .as_mut()
341 .and_then(|a| a.mapped_slice_mut())
342 .ok_or_else(|| format!("binding {i}: buffer memory is not host-visible"))?;
343 mapped[..bytes.len()].copy_from_slice(bytes);
344 // The padding is defined too, so an `arrayLength` that counts it
345 // reads zeros rather than whatever the last job left there.
346 for b in &mut mapped[bytes.len()..padded] {
347 *b = 0;
348 }
349 sizes.push(padded);
350 }
351
352 unsafe {
353 // One descriptor set per run, from a pool reset each time.
354 device
355 .reset_descriptor_pool(self.descriptor_pool, vk::DescriptorPoolResetFlags::empty())
356 .map_err(|e| format!("descriptor pool reset: {e}"))?;
357 // One descriptor set, or two with the ping-pong pair swapped in
358 // the second, so alternate passes bind the buffers the other
359 // way round without a write between dispatches.
360 let set_count = if ping_pong.is_some() { 2 } else { 1 };
361 let set_layouts = vec![set_layout; set_count];
362 let sets = device
363 .allocate_descriptor_sets(
364 &vk::DescriptorSetAllocateInfo::default()
365 .descriptor_pool(self.descriptor_pool)
366 .set_layouts(&set_layouts),
367 )
368 .map_err(|e| format!("descriptor set: {e}"))?;
369 let slot_for = |binding: usize, swapped: bool| -> usize {
370 match ping_pong {
371 Some((a, b)) if swapped && binding == a => b,
372 Some((a, b)) if swapped && binding == b => a,
373 _ => binding,
374 }
375 };
376 let mut infos: Vec<[vk::DescriptorBufferInfo; 1]> = Vec::with_capacity(set_count * bindings.len());
377 for (si, _) in sets.iter().enumerate() {
378 for i in 0..bindings.len() {
379 let slot = slot_for(i, si == 1);
380 infos.push([vk::DescriptorBufferInfo::default()
381 .buffer(self.slots[slot].buffer)
382 .offset(0)
383 .range(sizes[slot] as vk::DeviceSize)]);
384 }
385 }
386 let mut writes: Vec<vk::WriteDescriptorSet> = Vec::with_capacity(infos.len());
387 for (si, set) in sets.iter().enumerate() {
388 for (i, b) in bindings.iter().enumerate() {
389 writes.push(
390 vk::WriteDescriptorSet::default()
391 .dst_set(*set)
392 .dst_binding(i as u32)
393 .descriptor_type(match b.kind() {
394 BindKind::Storage => vk::DescriptorType::STORAGE_BUFFER,
395 BindKind::Uniform => vk::DescriptorType::UNIFORM_BUFFER,
396 })
397 .buffer_info(&infos[si * bindings.len() + i]),
398 );
399 }
400 }
401 device.update_descriptor_sets(&writes, &[]);
402
403 device
404 .begin_command_buffer(
405 self.cmd,
406 &vk::CommandBufferBeginInfo::default().flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT),
407 )
408 .map_err(|e| format!("begin: {e}"))?;
409 device.cmd_bind_pipeline(self.cmd, vk::PipelineBindPoint::COMPUTE, pipeline);
410 for pass in 0..passes {
411 let set = sets[if ping_pong.is_some() && pass % 2 == 1 { 1 } else { 0 }];
412 device.cmd_bind_descriptor_sets(self.cmd, vk::PipelineBindPoint::COMPUTE, layout, 0, &[set], &[]);
413 device.cmd_dispatch(self.cmd, groups[0], groups[1], groups[2]);
414 if pass + 1 < passes {
415 // The next pass reads what this one wrote.
416 device.cmd_pipeline_barrier(
417 self.cmd,
418 vk::PipelineStageFlags::COMPUTE_SHADER,
419 vk::PipelineStageFlags::COMPUTE_SHADER,
420 vk::DependencyFlags::empty(),
421 &[vk::MemoryBarrier::default()
422 .src_access_mask(vk::AccessFlags::SHADER_WRITE)
423 .dst_access_mask(vk::AccessFlags::SHADER_READ | vk::AccessFlags::SHADER_WRITE)],
424 &[],
425 &[],
426 );
427 }
428 }
429 // Shader writes become host-visible before the fence is signalled.
430 device.cmd_pipeline_barrier(
431 self.cmd,
432 vk::PipelineStageFlags::COMPUTE_SHADER,
433 vk::PipelineStageFlags::HOST,
434 vk::DependencyFlags::empty(),
435 &[vk::MemoryBarrier::default()
436 .src_access_mask(vk::AccessFlags::SHADER_WRITE)
437 .dst_access_mask(vk::AccessFlags::HOST_READ)],
438 &[],
439 &[],
440 );
441 device.end_command_buffer(self.cmd).map_err(|e| format!("end: {e}"))?;
442
443 let cmds = [self.cmd];
444 device
445 .queue_submit(self.core.queue, &[vk::SubmitInfo::default().command_buffers(&cmds)], self.fence)
446 .map_err(|e| format!("submit: {e}"))?;
447 device
448 .wait_for_fences(&[self.fence], true, u64::MAX)
449 .map_err(|e| format!("fence wait: {e}"))?;
450 device.reset_fences(&[self.fence]).map_err(|e| format!("fence reset: {e}"))?;
451 }
452
453 // Read back the read-write bindings — the ping-pong output from
454 // whichever buffer the last pass wrote.
455 let last_written = |i: usize| -> usize {
456 match ping_pong {
457 Some((a, b)) if i == b && passes % 2 == 0 => a,
458 _ => i,
459 }
460 };
461 for (i, b) in bindings.iter_mut().enumerate() {
462 if let Binding::Storage(out) = b {
463 let mapped = self.slots[last_written(i)]
464 .allocation
465 .as_ref()
466 .and_then(|a| a.mapped_slice())
467 .ok_or_else(|| format!("binding {i}: buffer memory is not host-visible"))?;
468 out.copy_from_slice(&mapped[..out.len()]);
469 }
470 }
471 Ok(())
472 }
473
474 fn pipeline(&mut self, key: &PipelineKey) -> Result<&Pipeline, String> {
475 if !self.pipelines.contains_key(key) {
476 let built = self.build_pipeline(key)?;
477 self.pipelines.insert(key.clone(), built);
478 }
479 Ok(&self.pipelines[key])
480 }
481
482 fn build_pipeline(&self, key: &PipelineKey) -> Result<Pipeline, String> {
483 let (spirv, workgroup_size) = compile_kernel(&key.kernel)?;
484 let device = &self.core.device;
485 unsafe {
486 let bindings: Vec<vk::DescriptorSetLayoutBinding> = key
487 .kinds
488 .iter()
489 .enumerate()
490 .map(|(i, k)| {
491 vk::DescriptorSetLayoutBinding::default()
492 .binding(i as u32)
493 .descriptor_type(match k {
494 BindKind::Storage => vk::DescriptorType::STORAGE_BUFFER,
495 BindKind::Uniform => vk::DescriptorType::UNIFORM_BUFFER,
496 })
497 .descriptor_count(1)
498 .stage_flags(vk::ShaderStageFlags::COMPUTE)
499 })
500 .collect();
501 let set_layout = device
502 .create_descriptor_set_layout(&vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings), None)
503 .map_err(|e| format!("descriptor set layout: {e}"))?;
504 let set_layouts = [set_layout];
505 let layout = match device.create_pipeline_layout(&vk::PipelineLayoutCreateInfo::default().set_layouts(&set_layouts), None) {
506 Ok(l) => l,
507 Err(e) => {
508 device.destroy_descriptor_set_layout(set_layout, None);
509 return Err(format!("pipeline layout: {e}"));
510 }
511 };
512 let module = match device.create_shader_module(&vk::ShaderModuleCreateInfo::default().code(&spirv), None) {
513 Ok(m) => m,
514 Err(e) => {
515 device.destroy_pipeline_layout(layout, None);
516 device.destroy_descriptor_set_layout(set_layout, None);
517 return Err(format!("shader module: {e}"));
518 }
519 };
520 let entry = CString::new(key.kernel.entry.as_str()).map_err(|e| format!("entry point name: {e}"))?;
521 let pipeline = device.create_compute_pipelines(
522 vk::PipelineCache::null(),
523 &[vk::ComputePipelineCreateInfo::default()
524 .stage(
525 vk::PipelineShaderStageCreateInfo::default()
526 .stage(vk::ShaderStageFlags::COMPUTE)
527 .module(module)
528 .name(&entry),
529 )
530 .layout(layout)],
531 None,
532 );
533 let pipeline = match pipeline {
534 Ok(p) => p[0],
535 Err((_, e)) => {
536 device.destroy_shader_module(module, None);
537 device.destroy_pipeline_layout(layout, None);
538 device.destroy_descriptor_set_layout(set_layout, None);
539 return Err(format!("compute pipeline: {e}"));
540 }
541 };
542 Ok(Pipeline { set_layout, layout, module, pipeline, workgroup_size })
543 }
544 }
545 }
546
547 impl Drop for ComputeDevice {
548 fn drop(&mut self) {
549 let device = self.core.device.clone();
550 unsafe {
551 let _ = device.device_wait_idle();
552 for (_, p) in self.pipelines.drain() {
553 device.destroy_pipeline(p.pipeline, None);
554 device.destroy_shader_module(p.module, None);
555 device.destroy_pipeline_layout(p.layout, None);
556 device.destroy_descriptor_set_layout(p.set_layout, None);
557 }
558 device.destroy_descriptor_pool(self.descriptor_pool, None);
559 device.destroy_fence(self.fence, None);
560 // The command buffer dies with the pool in VkCore's Drop.
561 }
562 let allocator = self.core.allocator.as_mut().unwrap();
563 for slot in &mut self.slots {
564 destroy_cpu_buffer(&device, allocator, slot);
565 }
566 }
567 }
568
569 /// WGSL to SPIR-V with every failure reported, plus the entry point's
570 /// workgroup size. `compile_wgsl` in the renderer panics on a bad shader,
571 /// which is right for the toolkit's own; a kernel here may be a user's.
572 fn compile_kernel(kernel: &Kernel) -> Result<(Vec<u32>, [u32; 3]), String> {
573 let module = naga::front::wgsl::parse_str(&kernel.source)
574 .map_err(|e| format!("WGSL parse error: {}", e.emit_to_string(&kernel.source).trim_end()))?;
575 let entry = module
576 .entry_points
577 .iter()
578 .find(|ep| ep.name == kernel.entry && ep.stage == naga::ShaderStage::Compute)
579 .ok_or_else(|| {
580 let offered: Vec<&str> = module
581 .entry_points
582 .iter()
583 .filter(|ep| ep.stage == naga::ShaderStage::Compute)
584 .map(|ep| ep.name.as_str())
585 .collect();
586 format!(
587 "no @compute entry point named `{}`; the module offers {}",
588 kernel.entry,
589 if offered.is_empty() { "none".to_string() } else { offered.join(", ") }
590 )
591 })?;
592 let workgroup_size = entry.workgroup_size;
593 let info = naga::valid::Validator::new(naga::valid::ValidationFlags::all(), naga::valid::Capabilities::empty())
594 .validate(&module)
595 .map_err(|e| format!("WGSL validation error: {}", e.emit_to_string(&kernel.source).trim_end()))?;
596 let options = naga::back::spv::Options {
597 lang_version: (1, 0),
598 flags: naga::back::spv::WriterFlags::LABEL_VARYINGS,
599 ..Default::default()
600 };
601 let spirv = naga::back::spv::write_vec(&module, &info, &options, None).map_err(|e| format!("SPIR-V: {e}"))?;
602 Ok((spirv, workgroup_size))
603 }
604
605 #[cfg(test)]
606 mod tests {
607 use super::*;
608
609 /// A device, or None with a note: these tests run on whatever Vulkan
610 /// the machine has (lavapipe counts) and skip where there is none.
611 fn device() -> Option<ComputeDevice> {
612 match ComputeDevice::new() {
613 Ok(d) => Some(d),
614 Err(e) => {
615 println!("skipping compute test: {e}");
616 None
617 }
618 }
619 }
620
621 const DOUBLE: &str = r#"
622 @group(0) @binding(0) var<storage, read_write> data: array<f32>;
623 @compute @workgroup_size(64)
624 fn main(@builtin(global_invocation_id) id: vec3<u32>) {
625 let i = id.x;
626 if (i < arrayLength(&data)) {
627 data[i] = data[i] * 2.0;
628 }
629 }"#;
630
631 #[test]
632 fn a_storage_binding_round_trips_through_the_kernel() {
633 let Some(mut dev) = device() else { return };
634 println!("compute on {}", dev.device_name());
635 // 1001 floats: not a multiple of the 16-byte padding, so the last
636 // element sits beside padding and must still come back doubled.
637 let mut data: Vec<f32> = (0..1001).map(|i| i as f32 * 0.5).collect();
638 let kernel = Kernel::new(DOUBLE, "main");
639 dev.run_over(&kernel, &mut [Binding::rw(&mut data)], 1001).unwrap();
640 for (i, v) in data.iter().enumerate() {
641 assert_eq!(*v, i as f32, "element {i}");
642 }
643 // Again, on the same device: the pipeline and the buffer are reused.
644 dev.run_over(&kernel, &mut [Binding::rw(&mut data)], 1001).unwrap();
645 assert_eq!(data[1000], 2000.0);
646 assert_eq!(dev.pipelines.len(), 1, "one pipeline for one kernel");
647 assert_eq!(dev.slots.len(), 1);
648 assert_eq!(dev.workgroup_size(&kernel, &[BindKind::Storage]).unwrap(), [64, 1, 1]);
649 }
650
651 #[test]
652 fn inputs_and_uniforms_bind_beside_the_output() {
653 let Some(mut dev) = device() else { return };
654 #[repr(C)]
655 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
656 struct Params {
657 scale: f32,
658 offset: f32,
659 _pad: [f32; 2],
660 }
661 const AXPY: &str = r#"
662 struct Params { scale: f32, offset: f32, pad: vec2<f32> }
663 @group(0) @binding(0) var<storage, read> a: array<vec4<f32>>;
664 @group(0) @binding(1) var<uniform> params: Params;
665 @group(0) @binding(2) var<storage, read_write> out: array<vec4<f32>>;
666 @compute @workgroup_size(32)
667 fn axpy(@builtin(global_invocation_id) id: vec3<u32>) {
668 let i = id.x;
669 if (i < arrayLength(&out)) {
670 out[i] = a[i] * params.scale + vec4<f32>(params.offset);
671 }
672 }"#;
673 let a: Vec<[f32; 4]> = (0..300).map(|i| [i as f32; 4]).collect();
674 let mut out = vec![[0.0f32; 4]; 300];
675 let params = Params { scale: 3.0, offset: 1.0, _pad: [0.0; 2] };
676 dev.run_over(
677 &Kernel::new(AXPY, "axpy"),
678 &mut [Binding::input(&a), Binding::uniform(¶ms), Binding::rw(&mut out)],
679 300,
680 )
681 .unwrap();
682 for (i, v) in out.iter().enumerate() {
683 assert_eq!(*v, [i as f32 * 3.0 + 1.0; 4], "element {i}");
684 }
685 assert!(a.iter().enumerate().all(|(i, v)| *v == [i as f32; 4]), "an input is never written back");
686 }
687
688 #[test]
689 fn a_bad_kernel_is_an_error_not_a_panic() {
690 let Some(mut dev) = device() else { return };
691 let mut data = vec![1.0f32; 4];
692 let err = dev.run_over(&Kernel::new("fn main( {", "main"), &mut [Binding::rw(&mut data)], 4).unwrap_err();
693 assert!(err.starts_with("WGSL parse error"), "{err}");
694 let err = dev.run_over(&Kernel::new(DOUBLE, "nope"), &mut [Binding::rw(&mut data)], 4).unwrap_err();
695 assert!(err.contains("`nope`") && err.contains("main"), "names the missing entry and the offer: {err}");
696 let typed = "@group(0) @binding(0) var<storage, read_write> d: array<f32>;\n@compute @workgroup_size(1) fn main() { d[0] = 1u; }";
697 let err = dev.run_over(&Kernel::new(typed, "main"), &mut [Binding::rw(&mut data)], 1).unwrap_err();
698 // naga's front end types as it parses, so a type error is a parse
699 // error; what matters is that it is a WGSL diagnostic, not a panic.
700 assert!(err.starts_with("WGSL") && err.contains("u32"), "{err}");
701 assert_eq!(data, vec![1.0; 4], "nothing ran");
702 // The device is still good after every failure.
703 dev.run_over(&Kernel::new(DOUBLE, "main"), &mut [Binding::rw(&mut data)], 4).unwrap();
704 assert_eq!(data, vec![2.0; 4]);
705 }
706
707 /// Passes chain inside one submission, and a ping-pong pair alternates
708 /// so the result lands in the output slice whether the count is odd or
709 /// even.
710 #[test]
711 fn passes_chain_and_ping_pong_lands_in_the_output() {
712 let Some(mut dev) = device() else { return };
713 // In place: three doublings are one octupling.
714 let mut data: Vec<f32> = (0..500).map(|i| i as f32).collect();
715 dev.run_passes_over(&Kernel::new(DOUBLE, "main"), &mut [Binding::rw(&mut data)], 500, 3, None).unwrap();
716 assert!(data.iter().enumerate().all(|(i, v)| *v == i as f32 * 8.0));
717
718 const COPY_DOUBLE: &str = r#"
719 @group(0) @binding(0) var<storage, read> src: array<f32>;
720 @group(0) @binding(1) var<storage, read_write> dst: array<f32>;
721 @compute @workgroup_size(64)
722 fn main(@builtin(global_invocation_id) id: vec3<u32>) {
723 let i = id.x;
724 if (i < arrayLength(&dst)) { dst[i] = src[i] * 2.0; }
725 }"#;
726 let src: Vec<f32> = (0..500).map(|i| i as f32).collect();
727 let kernel = Kernel::new(COPY_DOUBLE, "main");
728 for passes in [1u32, 2, 3, 4] {
729 let mut dst = vec![0.0f32; 500];
730 dev.run_passes_over(&kernel, &mut [Binding::input(&src), Binding::rw(&mut dst)], 500, passes, Some((0, 1))).unwrap();
731 let factor = 2f32.powi(passes as i32);
732 assert!(dst.iter().enumerate().all(|(i, v)| *v == i as f32 * factor), "{passes} passes give x{factor}");
733 }
734 assert!(src.iter().enumerate().all(|(i, v)| *v == i as f32), "the input slice is never written");
735
736 let mut dst = vec![0.0f32; 500];
737 let err = dev.run_passes_over(&kernel, &mut [Binding::input(&src), Binding::rw(&mut dst)], 500, 2, Some((1, 0))).unwrap_err();
738 assert!(err.contains("must be a read-only Input"), "{err}");
739 let err = dev.run_passes_over(&kernel, &mut [Binding::input(&src), Binding::rw(&mut dst)], 500, 0, None).unwrap_err();
740 assert!(err.contains("at least one pass"), "{err}");
741 }
742
743 #[test]
744 fn workgroup_arithmetic() {
745 assert_eq!(workgroups(0, 64), 1, "a dispatch of zero groups is invalid");
746 assert_eq!(workgroups(1, 64), 1);
747 assert_eq!(workgroups(64, 64), 1);
748 assert_eq!(workgroups(65, 64), 2);
749 assert_eq!(workgroups(1000, 0), 1000, "a zero group size is treated as one");
750 }
751 }