initial commit

Author: Kjolnyr

.gitignore

@@ -0,0 +1,2 @@
+/target
+Cargo.lock

Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "bevy_app_compute"
+version = "0.10.0"
+authors = ["Kjolnyr <[email protected]>"]
+edition = "2021"
+description = "App compute plugin for Bevy"
+repository = "https://github.com/Kjolnyr/bevy_app_compute"
+homepage = "https://github.com/Kjolnyr/bevy_app_compute"
+documentation = "https://docs.rs/bevy_app_compute"
+license = "MIT OR Apache-2.0"
+readme = "README.md"
+categories = ["game-development"]
+
+
+[dependencies]
+bevy = "0.10"
+parking_lot = "0.12.1"
+wgpu = "0.15.1"
+codespan-reporting = "0.11.1"
+futures-lite = "1.13.0"
+
+
+[[example]]
+name = "simple"
+
+[[example]]
+name = "async_heavy_work"

README.md

@@ -0,0 +1,109 @@
+# Bevy App Compute
+
+![MIT/Apache 2.0](https://img.shields.io/badge/license-MIT%2FApache-blue.svg)
+[![Doc](https://docs.rs/bevy_app_compute/badge.svg)](https://docs.rs/bevy_app_compute)
+[![Crate](https://img.shields.io/crates/v/bevy_mod_compute.svg)](https://crates.io/crates/bevy_mod_compute)
+
+
+Dispatch and run compute shaders on bevy from App World .
+
+## Usage
+
+### Setup
+
+Make an empty struct and implement `ComputeShader` on it. the `shader()` fn should point to your shader source code:
+```rust
+struct SimpleComputeShader;
+
+impl ComputeShader for SimpleComputeShader {
+    fn shader() -> ShaderRef {
+        "shaders/simple.wgsl".into()
+    }
+}
+```
+
+Add the plugin to your app:
+
+```rust
+use bevy::prelude::*;
+use bevy_app_compute::AppComputePlugin;
+
+fn main() {
+    App::new()
+        .add_plugin(AppComputePlugin::<SimpleComputeShader>::default());
+}
+```
+
+And then use the `AppCompute<T>` resource to run compute shaders!
+
+```rust
+use bevy::prelude::*;
+use bevy_app_compute::prelude::*;
+
+fn my_system(
+    app_compute: Res<AppCompute<SimpleComputeShader>>,
+) {
+
+    // Create a new worker
+    let Some(mut worker) = app_compute.worker() else { return; };
+
+    // Add some uniforms and storages with default values
+    worker.add_uniform(0, "uni", 5f32);
+    worker.add_storage(0, "storage", vec![0f32; 8]);
+
+    // Create a buffer needed to get data back from the GPU
+    // It has to be linked to a storage.
+    worker.add_staging_buffer("staging", "storage", std::mem::size_of::<f32>() * 8);
+
+    // run the shader
+    worker.run((8, 1, 1));
+
+    // You can read data from your staging buffer now
+    let result = worker.get_data("staging");
+    let value: &[f32] = cast_slice(&result);
+
+    println!("value: {:?}", value);
+}
+```
+
+
+### Asynchronous computation
+
+You can run your compute shaders asynchronously to avoid loosing frame time.
+
+```rust
+use bevy::prelude::*;
+use bevy_app_compute::prelude::*;
+
+fn my_sender_system(
+    mut app_compute: ResMut<AppCompute<MyComputeShader>>
+    ){
+    
+    let Some(mut worker) = app_compute.worker() else { return; };
+
+    worker.add_storage(0, "storage", [0f32; 30]);
+    worker.add_staging_buffer("staging", "storage", std::mem::size_of::<f32>() * 30);
+
+    // queue your worker for later use
+    app_compute.queue(worker, (30, 1, 1));
+}
+
+fn my_receiver_system(
+mut worker_events: EventReader<WorkerEvent<MyComputeShader>>
+) {
+    // An event is fired once a worker has finished processing
+    for ev in &mut worker_events.iter() {
+        let worker = &ev.worker;
+
+        let result = worker.get_data("staging");
+        let value: &[f32] = cast_slice(&result);
+
+        println!("got {} items back!", value.len());
+    }
+}
+```
+
+## Examples
+
+See [examples](https://github.com/kjolnyr/bevy_app_compute/tree/main/examples)
+

assets/shaders/heavy_work.wgsl

@@ -0,0 +1,16 @@
+
+
+@group(0) @binding(0)
+var<storage, read_write> my_storage: array<f32>;
+
+
+fn index_from_coords(coords: vec3<u32>) -> u32 {
+    return coords.z * 65535u * 65535u + coords.y * 65535u + coords.x;
+}
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) invocation_id: vec3<u32>) {
+
+    my_storage[index_from_coords(invocation_id)] = f32(index_from_coords(invocation_id));
+
+}

assets/shaders/simple.wgsl

@@ -0,0 +1,13 @@
+
+@group(0) @binding(0)
+var<uniform> uni: f32;
+
+@group(0) @binding(1)
+var<storage, read_write> my_storage: array<f32>;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) invocation_id: vec3<u32>) {
+
+    my_storage[invocation_id.x] = pow(uni, f32(invocation_id.x));
+
+}

examples/async_heavy_work.rs

@@ -0,0 +1,45 @@
+use bevy::{core::cast_slice, prelude::*, render::render_resource::ShaderRef};
+use bevy_app_compute::prelude::*;
+
+pub struct HeavyComputeShader;
+
+impl ComputeShader for HeavyComputeShader {
+    fn shader() -> ShaderRef {
+        "shaders/heavy_work.wgsl".into()
+    }
+}
+
+fn main() {
+    App::new()
+        .add_plugins(DefaultPlugins)
+        .add_plugin(AppComputePlugin::<HeavyComputeShader>::default())
+        .add_system(test)
+        .add_system(receive_data)
+        .run();
+}
+
+fn test(buttons: Res<Input<MouseButton>>, mut app_compute: ResMut<AppCompute<HeavyComputeShader>>) {
+    if !buttons.just_pressed(MouseButton::Left) {
+        return;
+    };
+
+    let storage_vec = vec![0f32; 30];
+    let Some(mut worker) = app_compute.worker() else { return; };
+
+    worker.add_storage(0, "storage", storage_vec);
+    worker.add_staging_buffer("staging", "storage", std::mem::size_of::<f32>() * 30);
+
+    app_compute.queue(worker, (30, 1, 1));
+}
+
+fn receive_data(mut worker_events: EventReader<WorkerEvent<HeavyComputeShader>>) {
+    for ev in &mut worker_events.iter() {
+        let worker = &ev.worker;
+
+        let result = worker.get_data("staging");
+
+        let value: &[f32] = cast_slice(&result);
+
+        println!("got {} items back", value.len());
+    }
+}

examples/simple.rs

@@ -0,0 +1,41 @@
+use bevy::{core::cast_slice, prelude::*, render::render_resource::ShaderRef};
+use bevy_app_compute::prelude::*;
+
+struct SimpleComputeShader;
+
+impl ComputeShader for SimpleComputeShader {
+    fn shader() -> ShaderRef {
+        "shaders/simple.wgsl".into()
+    }
+}
+
+fn main() {
+    App::new()
+        .add_plugins(DefaultPlugins)
+        .add_plugin(AppComputePlugin::<SimpleComputeShader>::default())
+        .add_system(on_click_compute)
+        .run();
+}
+
+fn on_click_compute(
+    buttons: Res<Input<MouseButton>>,
+    app_compute: Res<AppCompute<SimpleComputeShader>>,
+) {
+    if !buttons.just_pressed(MouseButton::Left) {
+        return;
+    };
+
+    let Some(mut worker) = app_compute.worker() else { return; };
+
+    worker.add_uniform(0, "uni", 5f32);
+    worker.add_storage(0, "storage", vec![0f32; 8]);
+    worker.add_staging_buffer("staging", "storage", std::mem::size_of::<f32>() * 8);
+
+    worker.run((8, 1, 1));
+
+    let result = worker.get_data("staging");
+
+    let value: &[f32] = cast_slice(&result);
+
+    println!("value: {:?}", value); // [1.0, 5.0, 24.999998, 124.999985, 624.9999, 3124.9993, 15624.996, 78124.98]
+}

rust-toolchain.toml

@@ -0,0 +1,2 @@
+[toolchain]
+channel = "nightly"

src/app_compute.rs

@@ -0,0 +1,91 @@
+use std::{future::IntoFuture, marker::PhantomData};
+
+use bevy::{
+    prelude::*,
+    render::{
+        render_resource::ComputePipeline,
+        renderer::{RenderDevice, RenderQueue},
+    },
+    tasks::{AsyncComputeTaskPool, Task},
+};
+use futures_lite::future;
+
+use crate::{worker::AppComputeWorker, ComputeShader, WorkerEvent};
+
+
+// Struct responsible for creating new workers and processing tasks
+// It requires <C> so that we don't mix tasks between different <C> 
+#[derive(Resource)]
+pub struct AppCompute<C: ComputeShader> {
+    render_device: RenderDevice,
+    render_queue: RenderQueue,
+    pub(crate) pipeline: Option<ComputePipeline>,
+    tasks: Vec<Task<AppComputeWorker>>,
+    _phantom: PhantomData<C>,
+}
+
+impl<C: ComputeShader> FromWorld for AppCompute<C> {
+    fn from_world(world: &mut bevy::prelude::World) -> Self {
+        let render_device = world.resource::<RenderDevice>().clone();
+        let render_queue = world.resource::<RenderQueue>().clone();
+
+        Self {
+            render_device,
+            render_queue,
+            pipeline: None,
+            tasks: vec![],
+            _phantom: PhantomData::default(),
+        }
+    }
+}
+
+impl<C: ComputeShader> AppCompute<C> {
+    pub fn worker(&self) -> Option<AppComputeWorker> {
+        if let Some(pipeline) = &self.pipeline {
+            // Probably could avoid cloning with some cursed lifetime rust code
+            Some(AppComputeWorker::new(
+                self.render_device.clone(),
+                self.render_queue.clone(),
+                pipeline.clone(),
+            ))
+        } else {
+            None
+        }
+    }
+
+    // Add a new compute tasks to the queue, this allow running compute shaders without blocking the main thread
+    pub fn queue(&mut self, mut worker: AppComputeWorker, workgroups: (u32, u32, u32)) {
+        let pool = AsyncComputeTaskPool::get();
+
+        let task = pool.spawn(async move {
+            worker.run(workgroups);
+            worker
+        });
+
+        self.tasks.push(task);
+    }
+
+    // Process the tasks and send an event once finished with the data
+    pub fn process_tasks(
+        mut app_compute: ResMut<Self>,
+        mut worker_events: EventWriter<WorkerEvent<C>>,
+    ) {
+        if app_compute.tasks.is_empty() {
+            return;
+        }
+
+        let mut indices_to_remove = vec![];
+
+        for (idx, task) in &mut app_compute.tasks.iter_mut().enumerate() {
+            let Some(worker) = future::block_on(future::poll_once(task.into_future())) else { continue; };
+
+            worker_events.send(WorkerEvent::new(worker));
+
+            indices_to_remove.push(idx);
+        }
+
+        for idx in indices_to_remove {
+            let _ = app_compute.tasks.remove(idx);
+        }
+    }
+}