remove mid pipeline copy to CPU with indirect dispatch

Author: friendlymatthew

src/gpu.rs

@@ -101,13 +101,48 @@ impl Gpu {
         bytemuck::cast_slice(&bytes).to_vec()
     }
 
-    pub fn dispatch(
+    pub fn indirect_buffer(&self, label: &str) -> wgpu::Buffer {
+        self.device.create_buffer(&wgpu::BufferDescriptor {
+            label: Some(label),
+            size: 12,
+            usage: wgpu::BufferUsages::STORAGE
+                | wgpu::BufferUsages::INDIRECT
+                | wgpu::BufferUsages::COPY_SRC
+                | wgpu::BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        })
+    }
+
+    pub fn read_buffer_at<T: bytemuck::Pod>(&self, buffer: &wgpu::Buffer, byte_offset: u64) -> T {
+        let size = std::mem::size_of::<T>() as u64;
+        let staging = self.staging_buffer(size);
+
+        let mut encoder = self
+            .device
+            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
+
+        encoder.copy_buffer_to_buffer(buffer, byte_offset, &staging, 0, size);
+        self.queue.submit(Some(encoder.finish()));
+
+        let slice = staging.slice(..);
+        let (tx, rx) = std::sync::mpsc::channel();
+        slice.map_async(wgpu::MapMode::Read, move |result| {
+            tx.send(result).unwrap();
+        });
+
+        self.device.poll(wgpu::Maintain::Wait);
+        rx.recv().unwrap().unwrap();
+
+        let data = slice.get_mapped_range();
+        *bytemuck::from_bytes(&data)
+    }
+
+    fn build_pipeline(
         &self,
         shader_src: &str,
         entry_point: &str,
         buffers: &[(&wgpu::Buffer, bool)],
-        workgroups: u32,
-    ) {
+    ) -> (wgpu::ComputePipeline, wgpu::BindGroup) {
         let module = self
             .device
             .create_shader_module(wgpu::ShaderModuleDescriptor {
@@ -171,6 +206,18 @@ impl Gpu {
             entries: &bind_group_entries,
         });
 
+        (pipeline, bind_group)
+    }
+
+    pub fn dispatch(
+        &self,
+        shader_src: &str,
+        entry_point: &str,
+        buffers: &[(&wgpu::Buffer, bool)],
+        workgroups: u32,
+    ) {
+        let (pipeline, bind_group) = self.build_pipeline(shader_src, entry_point, buffers);
+
         let mut encoder = self
             .device
             .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
@@ -185,6 +232,29 @@ impl Gpu {
         self.queue.submit(Some(encoder.finish()));
     }
 
+    pub fn dispatch_indirect(
+        &self,
+        shader_src: &str,
+        entry_point: &str,
+        buffers: &[(&wgpu::Buffer, bool)],
+        indirect: &wgpu::Buffer,
+    ) {
+        let (pipeline, bind_group) = self.build_pipeline(shader_src, entry_point, buffers);
+
+        let mut encoder = self
+            .device
+            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
+
+        {
+            let mut cpass = encoder.begin_compute_pass(&Default::default());
+            cpass.set_pipeline(&pipeline);
+            cpass.set_bind_group(0, &bind_group, &[]);
+            cpass.dispatch_workgroups_indirect(indirect, 0);
+        }
+
+        self.queue.submit(Some(encoder.finish()));
+    }
+
     pub fn multi_scan_u32(&self, buf: &wgpu::Buffer, n: usize) {
         self.multi_scan_generic(
             buf,

src/main.rs

@@ -55,27 +55,30 @@ fn parse(gpu: &Gpu, json: &str) -> Vec<TapeEntry> {
         n_wg.try_into().unwrap(),
     );
 
-    let scanned_mask = gpu.read_buffer_as::<u32>(&mask_buf);
-    let struct_count = scanned_mask[padded_len - 1] as usize;
-    let struct_wg = struct_count.div_ceil(256).max(1);
-    let struct_padded = struct_wg * 256;
-
-    let depth_buf = gpu.storage_buffer_empty("depth", (struct_padded * 4) as u64);
+    let indirect_buf = gpu.indirect_buffer("indirect");
     gpu.dispatch(
+        include_str!("shaders/multi/prepare_indirect.wgsl"),
+        "main",
+        &[(&mask_buf, true), (&indirect_buf, false)],
+        1,
+    );
+
+    let depth_buf = gpu.storage_buffer_empty("depth", buf_size(4));
+    gpu.dispatch_indirect(
         include_str!("shaders/multi/map_depth.wgsl"),
         "main",
         &[
             (&input_buf, true),
             (&compact_buf, true),
             (&depth_buf, false),
         ],
-        struct_wg.try_into().unwrap(),
+        &indirect_buf,
     );
 
-    gpu.multi_scan_i32(&depth_buf, struct_padded);
+    gpu.multi_scan_i32(&depth_buf, padded_len);
 
-    let parent_buf = gpu.storage_buffer_empty("parents", (struct_padded * 4) as u64);
-    gpu.dispatch(
+    let parent_buf = gpu.storage_buffer_empty("parents", buf_size(4));
+    gpu.dispatch_indirect(
         include_str!("shaders/multi/parent_link.wgsl"),
         "main",
         &[
@@ -84,14 +87,12 @@ fn parse(gpu: &Gpu, json: &str) -> Vec<TapeEntry> {
             (&depth_buf, true),
             (&parent_buf, false),
         ],
-        struct_wg.try_into().unwrap(),
+        &indirect_buf,
     );
 
-    let tape_buf = gpu.storage_buffer_empty(
-        "tape",
-        (struct_padded * std::mem::size_of::<TapeEntry>()) as u64,
-    );
-    gpu.dispatch(
+    let tape_buf =
+        gpu.storage_buffer_empty("tape", buf_size(std::mem::size_of::<TapeEntry>()));
+    gpu.dispatch_indirect(
         include_str!("shaders/multi/assemble_tape.wgsl"),
         "main",
         &[
@@ -101,9 +102,12 @@ fn parse(gpu: &Gpu, json: &str) -> Vec<TapeEntry> {
             (&parent_buf, true),
             (&tape_buf, false),
         ],
-        struct_wg.try_into().unwrap(),
+        &indirect_buf,
     );
 
+    let struct_count =
+        gpu.read_buffer_at::<u32>(&mask_buf, ((padded_len - 1) * 4) as u64) as usize;
+
     let tape = gpu.read_buffer_as::<TapeEntry>(&tape_buf);
     tape[..struct_count].to_vec()
 }

src/shaders/multi/prepare_indirect.wgsl

@@ -0,0 +1,21 @@
+@group(0) @binding(0)
+var<storage, read> mask: array<u32>;
+
+struct Indirect {
+    x: u32,
+    y: u32,
+    z: u32,
+}
+
+@group(0) @binding(1)
+var<storage, read_write> indirect: Indirect;
+
+@compute
+@workgroup_size(1)
+fn main() {
+    let n = arrayLength(&mask);
+    let struct_count = mask[n - 1u];
+    indirect.x = (struct_count + 255u) / 256u;
+    indirect.y = 1u;
+    indirect.z = 1u;
+}