wrgpgpu

Wren's library for GPGPU compute shaders, based on WGPU.

Minimal example:

This minimal example dispatches a compute shader that squares some values

use wrgpgpu::{include_wgsl, Bind, BindGroup, Device, ShaderArgs, StorageBufferBind};

#[tokio::main]
async fn main() {
	let device = Device::auto_high_performance().await;
    // Create the shader...
	let shader = device.create_shader::<BindGroup<StorageBufferBind<[f32; 256]>>>(ShaderArgs {
		label: "Square",
		shader: include_wgsl!("simple.wgsl"),
		entrypoint: "square",
	});

    // Create and upload a buffer of data to work on...
	let mut n = 0.0;
	let buffer = StorageBufferBind::new_init(
		&device,
		[0.0; 256].map(|_| {
			n += 1.0;
			n
		}),
	);
	let bind_group = device.bind(&buffer);

    // Dispatch the compute shader, this assumes a workgroup size of (256, 1, 1) in the shader...
	device.dispatch(&shader, &bind_group, (1, 1, 1));

    // Wait until it is complete
	while !device.is_complete() {
		tokio::time::sleep(Duration::from_millis(1)).await;
	}

    // Download the squared data
	let data = buffer.download(&device);

    // do something with the data...
}

And the corresponding gpu code:

@group(0) @binding(0)
var<storage, read_write> buffer: array<f32, 256>;

@compute @workgroup_size(256)
fn square(@builtin(global_invocation_id) global_id: vec3<u32>) {
	buffer[global_id.x] = buffer[global_id.x]*buffer[global_id.x];
}