1use wgpu::{BufferAddress, BufferUsages, Device};
2
3#[repr(C)]
4#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
5struct DispatchArgs {
6 per_row: u32,
7 rows: u32,
8 layers: u32,
9 _pad: u32,
10}
11
12pub const DISPATCH_ARGS_BYTES: BufferAddress = size_of::<DispatchArgs>() as BufferAddress;
13
14const _: () = assert!(DISPATCH_ARGS_BYTES == 16);
15
16pub struct DispatchTable {
17 buffer: crate::GpuBuffer,
18 slots: u32,
19}
20
21impl DispatchTable {
22 pub fn new(device: &Device, label: &str, slots: u32) -> Self {
23 assert!(slots > 0, "a dispatch table needs at least one slot");
24 Self {
25 buffer: crate::GpuBuffer::new(
26 device,
27 label,
28 slots as BufferAddress * DISPATCH_ARGS_BYTES,
29 BufferUsages::STORAGE
30 .union(BufferUsages::INDIRECT)
31 .union(BufferUsages::COPY_DST),
32 ),
33 slots,
34 }
35 }
36
37 pub fn buffer(&self) -> &crate::GpuBuffer {
38 &self.buffer
39 }
40
41 pub fn slots(&self) -> u32 {
42 self.slots
43 }
44
45 pub fn offset(&self, slot: u32) -> BufferAddress {
46 assert!(
47 slot < self.slots,
48 "dispatch slot {slot} is outside the {} table slots",
49 self.slots
50 );
51 slot as BufferAddress * DISPATCH_ARGS_BYTES
52 }
53}