compute-shader mesh generation example

Cargo.toml

@@ -3020,6 +3020,17 @@ description = "A very simple compute shader that writes to a buffer that is read
 category = "Shaders"
 wasm = false
 
+[[example]]
+name = "compute_mesh"
+path = "examples/shader/compute_mesh.rs"
+doc-scrape-examples = true
+
+[package.metadata.example.compute_mesh]
+name = "Compute Shader Mesh"
+description = "A compute shader that generates a mesh that is controlled by a Handle"
+category = "Shaders"
+wasm = false
+
 [[example]]
 name = "array_texture"
 path = "examples/shader/array_texture.rs"

assets/shaders/compute_mesh.wgsl

@@ -0,0 +1,75 @@
+// This shader is used for the compute_mesh example
+// The actual work it does is not important for the example and
+// has been hardcoded to return a cube mesh
+
+struct FirstIndex {
+    vertex: u32,
+    vertex_index: u32,
+}
+
+@group(0) @binding(0) var<uniform> first_index: FirstIndex;
+@group(0) @binding(1) var<storage, read_write> vertex_data: array<f32>;
+@group(0) @binding(2) var<storage, read_write> index_data: array<u32>;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    for (var i = 0u; i < 192; i++) {
+        // buffer is bigger than just our mesh, so we use the first_index.vertex
+        // to write to the correct range
+        vertex_data[i + first_index.vertex] = vertices[i];
+    }
+    for (var i = 0u; i < 36; i++) {
+        // buffer is bigger than just our mesh, so we use the first_index.vertex_index
+        // to write to the correct range
+        index_data[i + first_index.vertex_index] = u32(indices[i]);
+    }
+}
+
+// hardcoded compute shader data.
+const half_size = vec3(2.);
+const min = -half_size;
+const max = half_size;
+
+// Suppose Y-up right hand, and camera look from +Z to -Z
+const vertices = array(
+    // xyz, normal.xyz, uv.xy
+    // Front
+    min.x, min.y, max.z, 0.0, 0.0, 1.0, 0.0, 0.0,
+    max.x, min.y, max.z, 0.0, 0.0, 1.0, 1.0, 0.0,
+    max.x, max.y, max.z, 0.0, 0.0, 1.0, 1.0, 1.0,
+    min.x, max.y, max.z, 0.0, 0.0, 1.0, 0.0, 1.0,
+    // Back
+    min.x, max.y, min.z, 0.0, 0.0, -1.0, 1.0, 0.0,
+    max.x, max.y, min.z, 0.0, 0.0, -1.0, 0.0, 0.0,
+    max.x, min.y, min.z, 0.0, 0.0, -1.0, 0.0, 1.0,
+    min.x, min.y, min.z, 0.0, 0.0, -1.0, 1.0, 1.0,
+    // Right
+    max.x, min.y, min.z, 1.0, 0.0, 0.0, 0.0, 0.0,
+    max.x, max.y, min.z, 1.0, 0.0, 0.0, 1.0, 0.0,
+    max.x, max.y, max.z, 1.0, 0.0, 0.0, 1.0, 1.0,
+    max.x, min.y, max.z, 1.0, 0.0, 0.0, 0.0, 1.0,
+    // Left
+    min.x, min.y, max.z, -1.0, 0.0, 0.0, 1.0, 0.0,
+    min.x, max.y, max.z, -1.0, 0.0, 0.0, 0.0, 0.0,
+    min.x, max.y, min.z, -1.0, 0.0, 0.0, 0.0, 1.0,
+    min.x, min.y, min.z, -1.0, 0.0, 0.0, 1.0, 1.0,
+    // Top
+    max.x, max.y, min.z, 0.0, 1.0, 0.0, 1.0, 0.0,
+    min.x, max.y, min.z, 0.0, 1.0, 0.0, 0.0, 0.0,
+    min.x, max.y, max.z, 0.0, 1.0, 0.0, 0.0, 1.0,
+    max.x, max.y, max.z, 0.0, 1.0, 0.0, 1.0, 1.0,
+    // Bottom
+    max.x, min.y, max.z, 0.0, -1.0, 0.0, 0.0, 0.0,
+    min.x, min.y, max.z, 0.0, -1.0, 0.0, 1.0, 0.0,
+    min.x, min.y, min.z, 0.0, -1.0, 0.0, 1.0, 1.0,
+    max.x, min.y, min.z, 0.0, -1.0, 0.0, 0.0, 1.0
+);
+
+const indices = array(
+    0, 1, 2, 2, 3, 0, // front
+    4, 5, 6, 6, 7, 4, // back
+    8, 9, 10, 10, 11, 8, // right
+    12, 13, 14, 14, 15, 12, // left
+    16, 17, 18, 18, 19, 16, // top
+    20, 21, 22, 22, 23, 20, // bottom
+);

examples/README.md

@@ -467,6 +467,7 @@ Example | Description
 [Animated](../examples/shader/animate_shader.rs) | A shader that uses dynamic data like the time since startup
 [Array Texture](../examples/shader/array_texture.rs) | A shader that shows how to reuse the core bevy PBR shading functionality in a custom material that obtains the base color from an array texture.
 [Compute - Game of Life](../examples/shader/compute_shader_game_of_life.rs) | A compute shader that simulates Conway's Game of Life
+[Compute Shader Mesh](../examples/shader/compute_mesh.rs) | A compute shader that generates a mesh that is controlled by a Handle
 [Custom Render Phase](../examples/shader_advanced/custom_render_phase.rs) | Shows how to make a complete render phase
 [Custom Vertex Attribute](../examples/shader_advanced/custom_vertex_attribute.rs) | A shader that reads a mesh's custom vertex attribute
 [Custom phase item](../examples/shader_advanced/custom_phase_item.rs) | Demonstrates how to enqueue custom draw commands in a render phase

examples/shader/compute_mesh.rs

@@ -0,0 +1,297 @@
+//! This example shows how to initialize an empty mesh with a Handle
+//! and a render-world only usage. That buffer is then filled by a
+//! compute shader on the GPU without transferring data back
+//! to the CPU.
+//!
+//! The `mesh_allocator` is used to get references to the relevant slabs
+//! that contain the mesh data we're interested in.
+//!
+//! This example does not remove the `GenerateMesh` component after
+//! generating the mesh.
+
+use bevy::{
+    asset::RenderAssetUsages,
+    color::palettes::tailwind::RED_400,
+    mesh::Indices,
+    prelude::*,
+    render::{
+        extract_component::{ExtractComponent, ExtractComponentPlugin},
+        mesh::allocator::MeshAllocator,
+        render_graph::{self, RenderGraph, RenderLabel},
+        render_resource::{
+            binding_types::{storage_buffer, uniform_buffer},
+            *,
+        },
+        renderer::{RenderContext, RenderQueue},
+        Render, RenderApp, RenderStartup,
+    },
+};
+
+/// This example uses a shader source file from the assets subdirectory
+const SHADER_ASSET_PATH: &str = "shaders/compute_mesh.wgsl";
+
+fn main() {
+    App::new()
+        .add_plugins((
+            DefaultPlugins,
+            ComputeShaderMeshGeneratorPlugin,
+            ExtractComponentPlugin::<GenerateMesh>::default(),
+        ))
+        .insert_resource(ClearColor(Color::BLACK))
+        .add_systems(Startup, setup)
+        .run();
+}
+
+// We need a plugin to organize all the systems and render node required for this example
+struct ComputeShaderMeshGeneratorPlugin;
+impl Plugin for ComputeShaderMeshGeneratorPlugin {
+    fn build(&self, app: &mut App) {
+        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
+            return;
+        };
+
+        render_app
+            .init_resource::<Chunks>()
+            .add_systems(
+                RenderStartup,
+                (init_compute_pipeline, add_compute_render_graph_node),
+            )
+            .add_systems(Render, prepare_chunks);
+    }
+    fn finish(&self, app: &mut App) {
+        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
+            return;
+        };
+        render_app
+            .world_mut()
+            .resource_mut::<MeshAllocator>()
+            .extra_buffer_usages = BufferUsages::STORAGE;
+    }
+}
+
+/// Holds a handle to the empty mesh that should be filled
+/// by the compute shader.
+#[derive(Component, ExtractComponent, Clone)]
+struct GenerateMesh(Handle<Mesh>);
+
+fn setup(
+    mut commands: Commands,
+    mut meshes: ResMut<Assets<Mesh>>,
+    mut materials: ResMut<Assets<StandardMaterial>>,
+) {
+    // a truly empty mesh will error if used in Mesh3d
+    // so we set up the data to be what we want the compute shader to output
+    // We're using 36 indices, 24 vertices which is directly taken from
+    // the Bevy Cuboid mesh implementation
+    let empty_mesh = {
+        let mut mesh = Mesh::new(
+            PrimitiveTopology::TriangleList,
+            RenderAssetUsages::RENDER_WORLD,
+        )
+        .with_inserted_attribute(Mesh::ATTRIBUTE_POSITION, vec![[0.; 3]; 24])
+        .with_inserted_attribute(Mesh::ATTRIBUTE_NORMAL, vec![[0.; 3]; 24])
+        .with_inserted_attribute(Mesh::ATTRIBUTE_UV_0, vec![[0.; 2]; 24])
+        .with_inserted_indices(Indices::U32(vec![0; 36]));
+
+        mesh.asset_usage = RenderAssetUsages::RENDER_WORLD;
+        mesh
+    };
+
+    let handle = meshes.add(empty_mesh);
+
+    // we spawn two "users" of the mesh handle,
+    // but only insert `GenerateMesh` on one of them
+    // to show that the mesh handle works as usual
+    commands.spawn((
+        GenerateMesh(handle.clone()),
+        Mesh3d(handle.clone()),
+        MeshMaterial3d(materials.add(StandardMaterial {
+            base_color: RED_400.into(),
+            ..default()
+        })),
+        Transform::from_xyz(-2.5, 1., 0.),
+    ));
+
+    commands.spawn((
+        Mesh3d(handle),
+        MeshMaterial3d(materials.add(StandardMaterial {
+            base_color: RED_400.into(),
+            ..default()
+        })),
+        Transform::from_xyz(2.5, 1., 0.),
+    ));
+
+    // some additional scene elements.
+    // This mesh specifically is here so that we don't assume
+    // mesh_allocator offsets that would only work if we had
+    // one mesh in the scene.
+    commands.spawn((
+        Mesh3d(meshes.add(Circle::new(4.0))),
+        MeshMaterial3d(materials.add(Color::WHITE)),
+        Transform::from_rotation(Quat::from_rotation_x(-std::f32::consts::FRAC_PI_2)),
+    ));
+    commands.spawn((
+        PointLight {
+            shadows_enabled: true,
+            ..default()
+        },
+        Transform::from_xyz(4.0, 8.0, 4.0),
+    ));
+    // camera
+    commands.spawn((
+        Camera3d::default(),
+        Transform::from_xyz(-2.5, 4.5, 9.0).looking_at(Vec3::ZERO, Vec3::Y),
+    ));
+}
+
+fn add_compute_render_graph_node(mut render_graph: ResMut<RenderGraph>) {
+    // Add the compute node as a top-level node to the render graph. This means it will only execute
+    // once per frame. Normally, adding a node would use the `RenderGraphApp::add_render_graph_node`
+    // method, but it does not allow adding as a top-level node.
+    render_graph.add_node(ComputeNodeLabel, ComputeNode::default());
+}
+
+/// This is called "Chunks" because this example originated
+/// from a use case of generating chunks of landscape or voxels
+#[derive(Resource, Default)]
+struct Chunks(Vec<AssetId<Mesh>>);
+
+fn prepare_chunks(meshes_to_generate: Query<&GenerateMesh>, mut chunks: ResMut<Chunks>) {
+    // get the AssetId for each Handle<Mesh>
+    // which we'll use later to get the relevant buffers
+    // from the mesh_allocator
+    let chunk_data: Vec<AssetId<Mesh>> = meshes_to_generate
+        .iter()
+        .map(|gmesh| gmesh.0.id())
+        .collect();
+    chunks.0 = chunk_data;
+}
+
+#[derive(Resource)]
+struct ComputePipeline {
+    layout: BindGroupLayoutDescriptor,
+    pipeline: CachedComputePipelineId,
+}
+
+// init only happens once
+fn init_compute_pipeline(
+    mut commands: Commands,
+    asset_server: Res<AssetServer>,
+    pipeline_cache: Res<PipelineCache>,
+) {
+    let layout = BindGroupLayoutDescriptor::new(
+        "",
+        &BindGroupLayoutEntries::sequential(
+            ShaderStages::COMPUTE,
+            (
+                // offsets
+                uniform_buffer::<FirstIndex>(false),
+                // vertices
+                storage_buffer::<Vec<u32>>(false),
+                // indices
+                storage_buffer::<Vec<u32>>(false),
+            ),
+        ),
+    );
+    let shader = asset_server.load(SHADER_ASSET_PATH);
+    let pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
+        label: Some("Mesh generation compute shader".into()),
+        layout: vec![layout.clone()],
+        shader: shader.clone(),
+        ..default()
+    });
+    commands.insert_resource(ComputePipeline { layout, pipeline });
+}
+
+/// Label to identify the node in the render graph
+#[derive(Debug, Hash, PartialEq, Eq, Clone, RenderLabel)]
+struct ComputeNodeLabel;
+
+/// The node that will execute the compute shader
+#[derive(Default)]
+struct ComputeNode {}
+
+// A uniform that holds the vertex and index offsets
+// for the vertex/index mesh_allocator buffer slabs
+#[derive(ShaderType)]
+struct FirstIndex {
+    vertex: u32,
+    vertex_index: u32,
+}
+
+impl render_graph::Node for ComputeNode {
+    fn run(
+        &self,
+        _graph: &mut render_graph::RenderGraphContext,
+        render_context: &mut RenderContext,
+        world: &World,
+    ) -> Result<(), render_graph::NodeRunError> {
+        let Some(chunks) = world.get_resource::<Chunks>() else {
+            info!("no chunks");
+            return Ok(());
+        };
+        let mesh_allocator = world.resource::<MeshAllocator>();
+
+        for mesh_id in &chunks.0 {
+            let pipeline_cache = world.resource::<PipelineCache>();
+            let pipeline = world.resource::<ComputePipeline>();
+
+            if let Some(init_pipeline) = pipeline_cache.get_compute_pipeline(pipeline.pipeline) {
+                // the mesh_allocator holds slabs of meshes, so the buffers we get here
+                // can contain more data than just the mesh we're asking for.
+                // That's why there is a range field.
+                // You should *not* touch data in these buffers that is outside of the range.
+                let vertex_buffer_slice = mesh_allocator.mesh_vertex_slice(mesh_id).unwrap();
+                let index_buffer_slice = mesh_allocator.mesh_index_slice(mesh_id).unwrap();
+
+                let first = FirstIndex {
+                    // there are 8 vertex data values (pos, normal, uv) per vertex
+                    // and the vertex_buffer_slice.range.start is in "vertex elements"
+                    // which includes all of that data, so each index is worth 8 indices
+                    // to our shader code.
+                    vertex: vertex_buffer_slice.range.start * 8,
+                    // but each vertex index is a single value, so the index of the
+                    // vertex indices is exactly what the value is
+                    vertex_index: index_buffer_slice.range.start,
+                };
+
+                let mut uniforms = UniformBuffer::from(first);
+                uniforms.write_buffer(
+                    render_context.render_device(),
+                    world.resource::<RenderQueue>(),
+                );
+
+                // pass in the full mesh_allocator slabs as well as the first index
+                // offsets for the vertex and index buffers
+                let bind_group = render_context.render_device().create_bind_group(
+                    None,
+                    &pipeline_cache.get_bind_group_layout(&pipeline.layout),
+                    &BindGroupEntries::sequential((
+                        &uniforms,
+                        vertex_buffer_slice.buffer.as_entire_buffer_binding(),
+                        index_buffer_slice.buffer.as_entire_buffer_binding(),
+                    )),
+                );
+
+                let mut pass =
+                    render_context
+                        .command_encoder()
+                        .begin_compute_pass(&ComputePassDescriptor {
+                            label: Some("Mesh generation compute pass"),
+                            ..default()
+                        });
+                pass.push_debug_group("compute_mesh");
+
+                pass.set_bind_group(0, &bind_group, &[]);
+                pass.set_pipeline(init_pipeline);
+                // we only dispatch 1,1,1 workgroup here, but a real compute shader
+                // would take advantage of more workgroups
+                pass.dispatch_workgroups(1, 1, 1);
+
+                pass.pop_debug_group();
+            }
+        }
+
+        Ok(())
+    }
+}