Flesh out compute example

Just a copy of wgpu-rs's 'hello-compute' example:
https://github.com/gfx-rs/wgpu-rs/tree/v0.6/examples/hello-compute

Author: tombh

Cargo.lock

@@ -23,6 +23,7 @@ wgpu = "0.6.0"
 winit = { version = "0.24", features = ["web-sys"] }
 clap = "3.0.0-beta.2"
 strum = { version = "0.19", default_features = false, features = ["derive"] }
+bytemuck = { version = "1.4", features = ["derive"] }
 
 [build-dependencies]
 spirv-builder = { path = "../../../crates/spirv-builder", default-features = false }

examples/runners/wgpu/Cargo.toml

@@ -1,11 +1,13 @@
 use super::{shader_module, Options};
+use std::convert::TryInto;
+use wgpu::util::DeviceExt;
 
 fn create_device_queue() -> (wgpu::Device, wgpu::Queue) {
     async fn create_device_queue_async() -> (wgpu::Device, wgpu::Queue) {
         let instance = wgpu::Instance::new(wgpu::BackendBit::PRIMARY);
         let adapter = instance
             .request_adapter(&wgpu::RequestAdapterOptions {
-                power_preference: wgpu::PowerPreference::default(),
+                power_preference: wgpu::PowerPreference::Default,
                 compatible_surface: None,
             })
             .await
@@ -32,15 +34,51 @@ fn create_device_queue() -> (wgpu::Device, wgpu::Queue) {
     }
 }
 
-pub fn start(options: &Options) {
+pub async fn start(options: &Options) -> Vec<u32> {
+    let numbers: Vec<u32> = vec![1, 2, 3, 4];
+    let slice_size = numbers.len() * std::mem::size_of::<u32>();
+    let size = slice_size as wgpu::BufferAddress;
+
     let (device, queue) = create_device_queue();
 
-    // Load the shaders from disk
-    let module = device.create_shader_module(shader_module(options.shader));
+    let cs_module = device.create_shader_module(shader_module(options.shader));
+
+    let staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
+        label: None,
+        size,
+        usage: wgpu::BufferUsage::MAP_READ | wgpu::BufferUsage::COPY_DST,
+        mapped_at_creation: false,
+    });
+
+    let storage_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
+        label: Some("Storage Buffer"),
+        contents: bytemuck::cast_slice(&numbers),
+        usage: wgpu::BufferUsage::STORAGE
+            | wgpu::BufferUsage::COPY_DST
+            | wgpu::BufferUsage::COPY_SRC,
+    });
 
     let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
         label: None,
-        entries: &[],
+        entries: &[wgpu::BindGroupLayoutEntry {
+            binding: 0,
+            visibility: wgpu::ShaderStage::COMPUTE,
+            ty: wgpu::BindingType::StorageBuffer {
+                dynamic: false,
+                readonly: false,
+                min_binding_size: wgpu::BufferSize::new(4),
+            },
+            count: None,
+        }],
+    });
+
+    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
+        label: None,
+        layout: &bind_group_layout,
+        entries: &[wgpu::BindGroupEntry {
+            binding: 0,
+            resource: wgpu::BindingResource::Buffer(storage_buffer.slice(..)),
+        }],
     });
 
     let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
@@ -53,26 +91,47 @@ pub fn start(options: &Options) {
         label: None,
         layout: Some(&pipeline_layout),
         compute_stage: wgpu::ProgrammableStageDescriptor {
-            module: &module,
+            module: &cs_module,
             entry_point: "main_cs",
         },
     });
 
-    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
-        label: None,
-        layout: &bind_group_layout,
-        entries: &[],
-    });
-
     let mut encoder =
         device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
-
     {
         let mut cpass = encoder.begin_compute_pass();
-        cpass.set_bind_group(0, &bind_group, &[]);
         cpass.set_pipeline(&compute_pipeline);
-        cpass.dispatch(1, 1, 1);
+        cpass.set_bind_group(0, &bind_group, &[]);
+        cpass.insert_debug_marker("compute collatz iterations");
+        cpass.dispatch(numbers.len() as u32, 1, 1);
     }
+    encoder.copy_buffer_to_buffer(&storage_buffer, 0, &staging_buffer, 0, size);
 
     queue.submit(Some(encoder.finish()));
+
+    // Note that we're not calling `.await` here.
+    let buffer_slice = staging_buffer.slice(..);
+    let buffer_future = buffer_slice.map_async(wgpu::MapMode::Read);
+
+    // Poll the device in a blocking manner so that our future resolves.
+    // In an actual application, `device.poll(...)` should
+    // be called in an event loop or on another thread.
+    device.poll(wgpu::Maintain::Wait);
+
+    if let Ok(()) = buffer_future.await {
+        let data = buffer_slice.get_mapped_range();
+        let result = data
+            .chunks_exact(4)
+            .map(|b| u32::from_ne_bytes(b.try_into().unwrap()))
+            .collect();
+
+        // With the current interface, we have to make sure all mapped views are
+        // dropped before we unmap the buffer.
+        drop(data);
+        staging_buffer.unmap();
+
+        result
+    } else {
+        panic!("failed to run compute on gpu!")
+    }
 }

examples/runners/wgpu/src/compute.rs

@@ -32,11 +32,11 @@ pub struct Options {
 }
 
 #[cfg_attr(target_os = "android", ndk_glue::main(backtrace = "on"))]
-pub fn main() {
+pub async fn main() {
     let options: Options = Options::parse();
 
     if is_compute_shader(options.shader) {
-        compute::start(&options)
+        compute::start(&options).await;
     } else {
         graphics::start(&options);
     }

examples/runners/wgpu/src/lib.rs

@@ -1,3 +1,11 @@
 fn main() {
-    example_runner_wgpu::main()
+    subscriber::initialize_default_subscriber(None);
+    futures::executor::block_on(run());
+}
+
+async fn run() {
+    let times = example_runner_wgpu::main().await;
+    println!("Times: {:?}", times);
+    #[cfg(target_arch = "wasm32")]
+    log::info!("Times: {:?}", times);
 }

examples/runners/wgpu/src/main.rs

@@ -11,6 +11,34 @@ extern crate spirv_std;
 #[macro_use]
 pub extern crate spirv_std_macros;
 
+use spirv_std::storage_class::{Input, StorageBuffer};
+
+// The Collatz Conjecture states that for any integer n:
+// If n is even, n = n/2
+// If n is odd, n = 3n+1
+// And repeat this process for each new n, you will always eventually reach 1.
+// Though the conjecture has not been proven, no counterexample has ever been found.
+// This function returns how many times this recurrence needs to be applied to reach 1.
+pub fn collatz_iterations(mut n: i32) -> i32 {
+    let mut i = 0;
+    while n != 1 {
+        if n.rem_euclid(2) == 0 {
+            n = n / 2;
+        } else {
+            n = 3 * n + 1;
+        }
+        i += 1;
+    }
+    i
+}
+
 #[allow(unused_attributes)]
-#[spirv(gl_compute)]
-pub fn main_cs() {}
+#[spirv(gl_compute(local_size_x = 1))]
+pub fn main_cs(
+    #[spirv(global_invocation_id)] gid: Input<i32>,
+    #[spirv(storage_buffer)] mut storage: StorageBuffer<u32>,
+) {
+    let gid = gid.load();
+    let result = collatz_iterations(gid);
+    storage.store(result as u32)
+}