Merge pull request #30 from eliasstepanik/codex/multithread-rebuild_dirty_chunks-und-enqueue_visible_chunks

Parallelize chunk systems

client/assets/shaders/greedy_meshing.wgsl

@@ -21,7 +21,7 @@ struct Vertex {
 @group(0) @binding(1) var<uniform> params: Params;
 @group(0) @binding(2) var<storage, read_write> vertices: array<Vertex>;
 @group(0) @binding(3) var<storage, read_write> indices: array<u32>;
-@group(0) @binding(4) var<storage, read_write> counts: atomic<u32>;
+@group(0) @binding(4) var<storage, read_write> counts: array<atomic<u32>, 2>;
 
 const N: u32 = 16u;
 
@@ -86,16 +86,21 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
                         }
 
                         var height: u32 = 1u;
-                        outer: loop {
+                        loop {
                             if v0 + height >= N {
                                 break;
                             }
+                            var can_expand: bool = true;
                             for (var du: u32 = 0u; du < width; du = du + 1u) {
                                 let idx = (u0 + du) * N + v0 + height;
                                 if !mask[idx] || visited[idx] {
-                                    break outer;
+                                    can_expand = false;
+                                    break;
                                 }
                             }
+                            if !can_expand {
+                                break;
+                            }
                             height = height + 1u;
                         }
 
@@ -108,11 +113,11 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
                         // Compute base world-space position.
                         var base = params.origin;
                         if axis == 0u {
-                            base = base + vec3<f32>(f32(slice) + (dir > 0 ? 1.0 : 0.0), f32(u0), f32(v0)) * params.step;
+                            base = base + vec3<f32>(f32(slice) + select(0.0, 1.0, dir > 0), f32(u0), f32(v0)) * params.step;
                         } else if axis == 1u {
-                            base = base + vec3<f32>(f32(v0), f32(slice) + (dir > 0 ? 1.0 : 0.0), f32(u0)) * params.step;
+                            base = base + vec3<f32>(f32(v0), f32(slice) + select(0.0, 1.0, dir > 0), f32(u0)) * params.step;
                         } else {
-                            base = base + vec3<f32>(f32(u0), f32(v0), f32(slice) + (dir > 0 ? 1.0 : 0.0)) * params.step;
+                            base = base + vec3<f32>(f32(u0), f32(v0), f32(slice) + select(0.0, 1.0, dir > 0)) * params.step;
                         }
 
                         let size = vec2<f32>(f32(width) * params.step, f32(height) * params.step);
@@ -141,10 +146,10 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
                         let p3 = base + v_unit * size.y;
 
                         let vi = atomicAdd(&counts[0], 4u);
-                        vertices[vi] = Vertex(pos: p0, normal: normal, uv: vec2<f32>(0.0, 1.0));
-                        vertices[vi + 1u] = Vertex(pos: p1, normal: normal, uv: vec2<f32>(1.0, 1.0));
-                        vertices[vi + 2u] = Vertex(pos: p2, normal: normal, uv: vec2<f32>(1.0, 0.0));
-                        vertices[vi + 3u] = Vertex(pos: p3, normal: normal, uv: vec2<f32>(0.0, 0.0));
+                        vertices[vi] = Vertex(p0, normal, vec2<f32>(0.0, 1.0));
+                        vertices[vi + 1u] = Vertex(p1, normal, vec2<f32>(1.0, 1.0));
+                        vertices[vi + 2u] = Vertex(p2, normal, vec2<f32>(1.0, 0.0));
+                        vertices[vi + 3u] = Vertex(p3, normal, vec2<f32>(0.0, 0.0));
 
                         let ii = atomicAdd(&counts[1], 6u);
                         if dir > 0 {

client/src/plugins/environment/systems/voxels/meshing_gpu.rs

@@ -1,7 +1,7 @@
 use bevy::prelude::*;
 use bevy_app_compute::prelude::*;
 
-use super::structure::{MeshBufferPool, SparseVoxelOctree};
+use super::structure::{CHUNK_SIZE, MeshBufferPool, SparseVoxelOctree};
 
 #[repr(C)]
 #[derive(ShaderType, Copy, Clone, Default)]
@@ -22,6 +22,11 @@ pub struct VertexGpu {
     pub uv: Vec2,
 }
 
+const MAX_VOXELS: usize = (CHUNK_SIZE as usize) * (CHUNK_SIZE as usize) * (CHUNK_SIZE as usize);
+const MAX_QUADS: usize = MAX_VOXELS * 6;
+const MAX_VERTICES: usize = MAX_QUADS * 4;
+const MAX_INDICES: usize = MAX_QUADS * 6;
+
 #[derive(TypePath)]
 struct GreedyMeshingShader;
 
@@ -37,11 +42,16 @@ pub struct GpuMeshingWorker;
 
 impl ComputeWorker for GpuMeshingWorker {
     fn build(world: &mut World) -> AppComputeWorker<Self> {
+        // Allocate large temporary arrays on the heap to avoid stack overflows
+        let voxels = Box::new([0u32; MAX_VOXELS]);
+        let vertices = Box::new([VertexGpu::default(); MAX_VERTICES]);
+        let indices = Box::new([0u32; MAX_INDICES]);
+
         AppComputeWorkerBuilder::new(world)
-            .add_storage("voxels", &[0u32; 1])
+            .add_storage("voxels", voxels.as_ref())
             .add_uniform("params", &Params::default())
-            .add_storage("vertices", &[VertexGpu::default(); 1])
-            .add_storage("indices", &[0u32; 1])
+            .add_storage("vertices", vertices.as_ref())
+            .add_storage("indices", indices.as_ref())
             .add_storage("counts", &[0u32; 2])
             .add_pass::<GreedyMeshingShader>(
                 [1, 1, 1],

client/src/plugins/environment/systems/voxels/queue_systems.rs

@@ -1,23 +1,26 @@
-use bevy::prelude::*;
 use crate::plugins::environment::systems::voxels::helper::world_to_chunk;
 use crate::plugins::environment::systems::voxels::structure::*;
-
-
+use bevy::prelude::*;
+use rayon::prelude::*;
 
 /// enqueue chunks that *should* be visible but are not yet spawned
 /// enqueue chunks that *should* be visible but are not yet spawned
 pub fn enqueue_visible_chunks(
-    mut queue      : ResMut<ChunkQueue>,
-    spawned        : Res<SpawnedChunks>,
-    mut prev_cam   : ResMut<PrevCameraChunk>,
-    cfg            : Res<ChunkCullingCfg>,
-    cam_q          : Query<&GlobalTransform, With<Camera>>,
-    tree_q         : Query<&SparseVoxelOctree>,
+    mut queue: ResMut<ChunkQueue>,
+    spawned: Res<SpawnedChunks>,
+    mut prev_cam: ResMut<PrevCameraChunk>,
+    cfg: Res<ChunkCullingCfg>,
+    cam_q: Query<&GlobalTransform, With<Camera>>,
+    tree_q: Query<&SparseVoxelOctree>,
 ) {
-    let Ok(tree)     = tree_q.get_single() else { return };
-    let Ok(cam_tf)  = cam_q.get_single() else { return };
-    let cam_pos  = cam_tf.translation();
-    let centre   = world_to_chunk(tree, cam_pos);
+    let Ok(tree) = tree_q.get_single() else {
+        return;
+    };
+    let Ok(cam_tf) = cam_q.get_single() else {
+        return;
+    };
+    let cam_pos = cam_tf.translation();
+    let centre = world_to_chunk(tree, cam_pos);
 
     if prev_cam.0 == Some(centre) {
         return;
@@ -28,14 +31,18 @@ pub fn enqueue_visible_chunks(
 
     let mut keys: Vec<(ChunkKey, i32)> = tree
         .occupied_chunks
-        .iter()
+        .par_iter()
         .filter_map(|key| {
             let dx = key.0 - centre.0;
             let dy = key.1 - centre.1;
             let dz = key.2 - centre.2;
-            if dx.abs() > r || dy.abs() > r || dz.abs() > r { return None; }
-            if spawned.0.contains_key(key) { return None; }
-            Some((*key, dx*dx + dy*dy + dz*dz))
+            if dx.abs() > r || dy.abs() > r || dz.abs() > r {
+                return None;
+            }
+            if spawned.0.contains_key(key) {
+                return None;
+            }
+            Some((*key, dx * dx + dy * dy + dz * dz))
         })
         .collect();
 
@@ -51,15 +58,19 @@ pub fn enqueue_visible_chunks(
 
 /// move a limited number of keys from the queue into the octree’s dirty set
 pub fn process_chunk_queue(
-    mut queue   : ResMut<ChunkQueue>,
-    budget      : Res<ChunkBudget>,
-    mut tree_q  : Query<&mut SparseVoxelOctree>,
+    mut queue: ResMut<ChunkQueue>,
+    budget: Res<ChunkBudget>,
+    mut tree_q: Query<&mut SparseVoxelOctree>,
 ) {
-    let Ok(mut tree) = tree_q.get_single_mut() else { return };
+    let Ok(mut tree) = tree_q.get_single_mut() else {
+        return;
+    };
     for _ in 0..budget.per_frame {
         if let Some(key) = queue.keys.pop_front() {
             queue.set.remove(&key);
             tree.dirty_chunks.insert(key);
-        } else { break; }
+        } else {
+            break;
+        }
     }
-}
+}

client/src/plugins/environment/systems/voxels/render_chunks.rs

@@ -1,12 +1,13 @@
 use crate::plugins::big_space::big_space_plugin::RootGrid;
+use crate::plugins::environment::systems::voxels::atlas::VoxelTextureAtlas;
 use crate::plugins::environment::systems::voxels::meshing::mesh_chunk;
 use crate::plugins::environment::systems::voxels::structure::*;
-use crate::plugins::environment::systems::voxels::atlas::VoxelTextureAtlas;
 use bevy::pbr::wireframe::Wireframe;
 use bevy::prelude::*;
 use bevy::render::mesh::Mesh;
 use big_space::prelude::GridCell;
 use itertools::Itertools;
+use rayon::prelude::*;
 use std::collections::HashMap;
 use std::fmt::format;
 
@@ -41,38 +42,46 @@ pub fn rebuild_dirty_chunks(
         }
 
         //------------------------------------------------ collect voxel data
-        let mut bufs = Vec::new();
-        for key in tree.dirty_chunks.iter().copied() {
-            let lod = existing.get(&key).map(|v| v.3).unwrap_or(0);
-            let mut buf = [[[None; CHUNK_SIZE as usize]; CHUNK_SIZE as usize]; CHUNK_SIZE as usize];
+        let tree_ref = &*tree;
+        let bufs: Vec<_> = tree
+            .dirty_chunks
+            .par_iter()
+            .copied()
+            .map(|key| {
+                let lod = existing.get(&key).map(|v| v.3).unwrap_or(0);
+                let mut buf =
+                    [[[None; CHUNK_SIZE as usize]; CHUNK_SIZE as usize]; CHUNK_SIZE as usize];
 
-            let half = tree.size * 0.5;
-            let step = tree.get_spacing_at_depth(tree.max_depth);
-            let origin = Vec3::new(
-                key.0 as f32 * CHUNK_SIZE as f32 * step - half,
-                key.1 as f32 * CHUNK_SIZE as f32 * step - half,
-                key.2 as f32 * CHUNK_SIZE as f32 * step - half,
-            );
+                let half = tree_ref.size * 0.5;
+                let step = tree_ref.get_spacing_at_depth(tree_ref.max_depth);
+                let origin = Vec3::new(
+                    key.0 as f32 * CHUNK_SIZE as f32 * step - half,
+                    key.1 as f32 * CHUNK_SIZE as f32 * step - half,
+                    key.2 as f32 * CHUNK_SIZE as f32 * step - half,
+                );
 
-            let mult = 1 << lod;
-            for gx in (0..CHUNK_SIZE).step_by(mult as usize) {
-                for gy in (0..CHUNK_SIZE).step_by(mult as usize) {
-                    for gz in (0..CHUNK_SIZE).step_by(mult as usize) {
-                        let center = origin
-                            + Vec3::new(
-                                (gx + mult / 2) as f32 * step,
-                                (gy + mult / 2) as f32 * step,
-                                (gz + mult / 2) as f32 * step,
-                            );
-                        if let Some(v) = tree.get_voxel_at_world_coords(center) {
-                            for lx in 0..mult {
-                                for ly in 0..mult {
-                                    for lz in 0..mult {
-                                        let ix = gx + lx;
-                                        let iy = gy + ly;
-                                        let iz = gz + lz;
-                                        if ix < CHUNK_SIZE && iy < CHUNK_SIZE && iz < CHUNK_SIZE {
-                                            buf[ix as usize][iy as usize][iz as usize] = Some(*v);
+                let mult = 1 << lod;
+                for gx in (0..CHUNK_SIZE).step_by(mult as usize) {
+                    for gy in (0..CHUNK_SIZE).step_by(mult as usize) {
+                        for gz in (0..CHUNK_SIZE).step_by(mult as usize) {
+                            let center = origin
+                                + Vec3::new(
+                                    (gx + mult / 2) as f32 * step,
+                                    (gy + mult / 2) as f32 * step,
+                                    (gz + mult / 2) as f32 * step,
+                                );
+                            if let Some(v) = tree_ref.get_voxel_at_world_coords(center) {
+                                for lx in 0..mult {
+                                    for ly in 0..mult {
+                                        for lz in 0..mult {
+                                            let ix = gx + lx;
+                                            let iy = gy + ly;
+                                            let iz = gz + lz;
+                                            if ix < CHUNK_SIZE && iy < CHUNK_SIZE && iz < CHUNK_SIZE
+                                            {
+                                                buf[ix as usize][iy as usize][iz as usize] =
+                                                    Some(*v);
+                                            }
                                         }
                                     }
                                 }
@@ -80,10 +89,10 @@ pub fn rebuild_dirty_chunks(
                         }
                     }
                 }
-            }
 
-            bufs.push((key, buf, origin, step, lod));
-        }
+                (key, buf, origin, step, lod)
+            })
+            .collect();
 
         //------------------------------------------------ create / update
         for (key, buf, origin, step, lod) in bufs {