Skip to main content

dynamis_gpu/
dispatch.rs

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}