Compute example

bevy-app/crates/viewer/Cargo.toml

@@ -13,6 +13,7 @@ path = "examples/standard-material.rs"
 
 [dependencies]
 bevy = { version = "0.10.0", features = ["spirv_shader_passthrough"] }
+rmp-serde="1.1.1"
 
 bevy-rust-gpu = { git = "https://github.com/bevy-rust-gpu/bevy-rust-gpu", tag = "v0.5.0" }
 rust-gpu-bridge = { git = "https://github.com/bevy-rust-gpu/rust-gpu-bridge", features = ["glam"], tag = "v0.5.0" }

bevy-app/crates/viewer/examples/compute.rs

@@ -0,0 +1,294 @@
+//! A compute shader that simulates Conway's Game of Life.
+//!
+//! Compute shaders use the GPU for computing arbitrary information, that may be independent of what
+//! is rendered to the screen.
+
+use bevy::{
+    asset::{AssetLoader, LoadedAsset},
+    prelude::{
+        default, AddAsset, App, AssetServer, Assets, Camera2dBundle, ClearColor, Color, Commands,
+        Deref, FromWorld, Handle, Image, IntoSystemConfig, Plugin, PluginGroup, Res, ResMut,
+        Resource, Shader, Vec2, World, info,
+    },
+    render::{
+        extract_resource::{ExtractResource, ExtractResourcePlugin},
+        render_asset::RenderAssets,
+        render_graph::{self, RenderGraph},
+        render_resource::{
+            BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout,
+            BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingResource, BindingType,
+            CachedComputePipelineId, CachedPipelineState, ComputePassDescriptor,
+            ComputePipelineDescriptor, Extent3d, PipelineCache, ShaderStages, StorageTextureAccess,
+            TextureDimension, TextureFormat, TextureUsages, TextureViewDimension,
+        },
+        renderer::{RenderContext, RenderDevice},
+        RenderApp, RenderSet,
+    },
+    sprite::{Sprite, SpriteBundle},
+    window::{Window, WindowPlugin},
+    DefaultPlugins,
+};
+use bevy_rust_gpu::RustGpuBuilderOutput;
+use std::borrow::Cow;
+
+const SIZE: (u32, u32) = (1280, 720);
+const WORKGROUP_SIZE: u32 = 8;
+
+fn main() {
+    App::new()
+        .insert_resource(ClearColor(Color::BLACK))
+        .add_plugins(DefaultPlugins.set(WindowPlugin {
+            primary_window: Some(Window {
+                // uncomment for unthrottled FPS
+                // present_mode: bevy::window::PresentMode::AutoNoVsync,
+                ..default()
+            }),
+            ..default()
+        }))
+        .init_asset_loader::<RustGpuMsgpackLoader>()
+        .add_plugin(GameOfLifeComputePlugin)
+        .add_startup_system(setup)
+        .run();
+}
+
+fn setup(mut commands: Commands, mut images: ResMut<Assets<Image>>) {
+    let mut image = Image::new_fill(
+        Extent3d {
+            width: SIZE.0,
+            height: SIZE.1,
+            depth_or_array_layers: 1,
+        },
+        TextureDimension::D2,
+        &[0, 0, 0, 255],
+        TextureFormat::Rgba8Unorm,
+    );
+    image.texture_descriptor.usage =
+        TextureUsages::COPY_DST | TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING;
+    let image = images.add(image);
+
+    commands.spawn(SpriteBundle {
+        sprite: Sprite {
+            custom_size: Some(Vec2::new(SIZE.0 as f32, SIZE.1 as f32)),
+            ..default()
+        },
+        texture: image.clone(),
+        ..default()
+    });
+    commands.spawn(Camera2dBundle::default());
+
+    commands.insert_resource(GameOfLifeImage(image));
+}
+
+pub struct GameOfLifeComputePlugin;
+
+impl Plugin for GameOfLifeComputePlugin {
+    fn build(&self, app: &mut App) {
+        // Extract the game of life image resource from the main world into the render world
+        // for operation on by the compute shader and display on the sprite.
+        app.add_plugin(ExtractResourcePlugin::<GameOfLifeImage>::default());
+        let render_app = app.sub_app_mut(RenderApp);
+        render_app
+            .init_resource::<GameOfLifePipeline>()
+            .add_system(queue_bind_group.in_set(RenderSet::Queue));
+
+        let mut render_graph = render_app.world.resource_mut::<RenderGraph>();
+        render_graph.add_node("game_of_life", GameOfLifeNode::default());
+        render_graph.add_node_edge(
+            "game_of_life",
+            bevy::render::main_graph::node::CAMERA_DRIVER,
+        );
+    }
+}
+
+#[derive(Resource, Clone, Deref, ExtractResource)]
+struct GameOfLifeImage(Handle<Image>);
+
+#[derive(Resource)]
+struct GameOfLifeImageBindGroup(BindGroup);
+
+fn queue_bind_group(
+    mut commands: Commands,
+    pipeline: Res<GameOfLifePipeline>,
+    gpu_images: Res<RenderAssets<Image>>,
+    game_of_life_image: Res<GameOfLifeImage>,
+    render_device: Res<RenderDevice>,
+) {
+    let view = gpu_images.get(&game_of_life_image.0).unwrap();
+    let bind_group = render_device.create_bind_group(&BindGroupDescriptor {
+        label: None,
+        layout: &pipeline.texture_bind_group_layout,
+        entries: &[BindGroupEntry {
+            binding: 0,
+            resource: BindingResource::TextureView(&view.texture_view),
+        }],
+    });
+    commands.insert_resource(GameOfLifeImageBindGroup(bind_group));
+}
+
+#[derive(Resource)]
+pub struct GameOfLifePipeline {
+    texture_bind_group_layout: BindGroupLayout,
+    init_pipeline: CachedComputePipelineId,
+    update_pipeline: CachedComputePipelineId,
+}
+
+impl FromWorld for GameOfLifePipeline {
+    fn from_world(world: &mut World) -> Self {
+        let texture_bind_group_layout =
+            world
+                .resource::<RenderDevice>()
+                .create_bind_group_layout(&BindGroupLayoutDescriptor {
+                    label: None,
+                    entries: &[BindGroupLayoutEntry {
+                        binding: 0,
+                        visibility: ShaderStages::COMPUTE,
+                        ty: BindingType::StorageTexture {
+                            access: StorageTextureAccess::ReadWrite,
+                            format: TextureFormat::Rgba8Unorm,
+                            view_dimension: TextureViewDimension::D2,
+                        },
+                        count: None,
+                    }],
+                });
+
+        let shader = world
+            .resource::<AssetServer>()
+            .load("rust-gpu/shader.rust-gpu.msgpack");
+
+        // let shader = world
+        //     .resource::<AssetServer>()
+        //     .load("gol.wgsl");
+        let pipeline_cache = world.resource::<PipelineCache>();
+        let init_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
+            label: Some("init_compute_primes".into()),
+            layout: vec![texture_bind_group_layout.clone()],
+            push_constant_ranges: Vec::new(),
+            shader: shader.clone(),
+            shader_defs: vec![],
+            entry_point: Cow::from("init"),
+        });
+        let update_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
+            label: Some("update_compute_primes".into()),
+            layout: vec![texture_bind_group_layout.clone()],
+            push_constant_ranges: Vec::new(),
+            shader,
+            shader_defs: vec![],
+            entry_point: Cow::from("update"),
+        });
+
+        GameOfLifePipeline {
+            texture_bind_group_layout,
+            init_pipeline,
+            update_pipeline,
+        }
+    }
+}
+
+enum GameOfLifeState {
+    Loading,
+    Init,
+    Update,
+}
+
+struct GameOfLifeNode {
+    state: GameOfLifeState,
+}
+
+impl Default for GameOfLifeNode {
+    fn default() -> Self {
+        Self {
+            state: GameOfLifeState::Loading,
+        }
+    }
+}
+
+impl render_graph::Node for GameOfLifeNode {
+    fn update(&mut self, world: &mut World) {
+        let pipeline = world.resource::<GameOfLifePipeline>();
+        let pipeline_cache = world.resource::<PipelineCache>();
+
+        // if the corresponding pipeline has loaded, transition to the next stage
+        match self.state {
+            GameOfLifeState::Loading => {
+                if let CachedPipelineState::Ok(_) =
+                    pipeline_cache.get_compute_pipeline_state(pipeline.init_pipeline)
+                {
+                    self.state = GameOfLifeState::Init;
+                }
+            }
+            GameOfLifeState::Init => {
+                if let CachedPipelineState::Ok(_) =
+                    pipeline_cache.get_compute_pipeline_state(pipeline.update_pipeline)
+                {
+                    self.state = GameOfLifeState::Update;
+                }
+            }
+            GameOfLifeState::Update => {}
+        }
+    }
+
+    fn run(
+        &self,
+        _graph: &mut render_graph::RenderGraphContext,
+        render_context: &mut RenderContext,
+        world: &World,
+    ) -> Result<(), render_graph::NodeRunError> {
+        let texture_bind_group = &world.resource::<GameOfLifeImageBindGroup>().0;
+        let pipeline_cache = world.resource::<PipelineCache>();
+        let pipeline = world.resource::<GameOfLifePipeline>();
+
+        let mut pass = render_context
+            .command_encoder()
+            .begin_compute_pass(&ComputePassDescriptor::default());
+
+        pass.set_bind_group(0, texture_bind_group, &[]);
+
+        // select the pipeline based on the current state
+        match self.state {
+            GameOfLifeState::Loading => {}
+            GameOfLifeState::Init => {
+                let init_pipeline = pipeline_cache
+                    .get_compute_pipeline(pipeline.init_pipeline)
+                    .unwrap();
+                pass.set_pipeline(init_pipeline);
+                pass.dispatch_workgroups(SIZE.0 / WORKGROUP_SIZE, SIZE.1 / WORKGROUP_SIZE, 1);
+            }
+            GameOfLifeState::Update => {
+                let update_pipeline = pipeline_cache
+                    .get_compute_pipeline(pipeline.update_pipeline)
+                    .unwrap();
+                pass.set_pipeline(update_pipeline);
+                pass.dispatch_workgroups(SIZE.0 / WORKGROUP_SIZE, SIZE.1 / WORKGROUP_SIZE, 1);
+            }
+        }
+
+        Ok(())
+    }
+}
+
+#[derive(Default)]
+struct RustGpuMsgpackLoader;
+
+impl AssetLoader for RustGpuMsgpackLoader {
+    fn load<'a>(
+        &'a self,
+        bytes: &'a [u8],
+        load_context: &'a mut bevy::asset::LoadContext,
+    ) -> bevy::utils::BoxedFuture<'a, Result<(), bevy::asset::Error>> {
+        Box::pin(async move {
+            let spirv_bytes: RustGpuBuilderOutput = rmp_serde::from_slice(bytes).unwrap();
+            let shader = match spirv_bytes.modules {
+                bevy_rust_gpu::RustGpuBuilderModules::Single(s) => {
+                    Shader::from_spirv(s)
+                },
+                bevy_rust_gpu::RustGpuBuilderModules::Multi(_) => todo!(),
+            };
+            load_context.set_default_asset(LoadedAsset::new(shader));
+            Ok(())
+        })
+    }
+
+    fn extensions(&self) -> &[&str] {
+        &["rust-gpu.msgpack"]
+    }
+}

rust-gpu/crates/shader/src/lib.rs

@@ -29,7 +29,7 @@ use rust_gpu_sdf::{
 };
 use spirv_std::{
     arch::{ddx, ddy},
-    glam::{Mat3, Vec2, Vec3, Vec4, Vec4Swizzles},
+    glam::{IVec4, Mat3, UVec3, Vec2, Vec3, Vec4, Vec4Swizzles},
     spirv,
 };
 
@@ -453,3 +453,46 @@ pub fn fragment_sdf_3d(
     }
     */
 }
+
+pub fn collatz(mut n: u32) -> Option<u32> {
+    let mut i = 0;
+    if n == 0 {
+        return None;
+    }
+    while n != 1 {
+        n = if n % 2 == 0 {
+            n / 2
+        } else {
+            // Overflow? (i.e. 3*n + 1 > 0xffff_ffff)
+            if n >= 0x5555_5555 {
+                return None;
+            }
+            // TODO: Use this instead when/if checked add/mul can work: n.checked_mul(3)?.checked_add(1)?
+            3 * n + 1
+        };
+        i += 1;
+    }
+    Some(i)
+}
+
+#[spirv(compute(threads(64)))]
+pub fn init(
+    #[spirv(global_invocation_id)] id: UVec3,
+    #[spirv(num_workgroups)] num: UVec3,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] texture: &mut [Vec4],
+) {
+    let index = (id.y * num.x) as usize + id.y as usize;
+    texture[index] = Vec4::new(1.0, 0.2, 0.3, 0.4);
+}
+
+#[spirv(compute(threads(64)))]
+pub fn update(
+    #[spirv(global_invocation_id)] id: UVec3,
+    #[spirv(num_workgroups)] num: UVec3,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] texture: &mut [Vec4],
+) {
+    let index = (id.y * num.x) as usize + id.y as usize;
+    let pixel = &mut texture[index];
+
+    pixel.x = (pixel.x + 0.01).fract();
+}