Merge pull request #31 from davids91/feature/compute-shader-migration

Resolves #20 - Migrated GPU implementation from fragment shader to compute shaders

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "shocovox-rs"
-version = "0.2.2"
+version = "0.3.0"
 edition = "2021"
 authors = ["Dávid Tóth <[email protected]>"]
 license = "MIT OR Apache-2.0"

assets/shaders/viewport_render.wgsl

@@ -671,36 +671,45 @@ struct Viewport {
     fov: f32,
 }
 
-@group(2) @binding(0)
+@group(0) @binding(1)
+var output_texture: texture_storage_2d<rgba8unorm, read_write>;
+
+@group(0) @binding(2)
 var<uniform> viewport: Viewport;
 
-@group(2) @binding(1)
+@group(0) @binding(3)
 var<uniform> octreeMetaData: OctreeMetaData;
 
-@group(2) @binding(2)
+@group(0) @binding(4)
 var<storage, read_write> nodes: array<SizedNode>;
 
-@group(2) @binding(3)
+@group(0) @binding(5)
 var<storage, read_write> voxels: array<Voxelement>;
 
-@fragment
-fn fragment(mesh: VertexOutput) -> @location(0) vec4<f32> {
-    let viewport_ = viewport; //Read only once from global RAM
+@compute @workgroup_size(8, 8, 1)
+fn update(
+    @builtin(global_invocation_id) invocation_id: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>,
+) {
+    let pixel_location = vec2u(invocation_id.xy);
+    let pixel_location_normalized = vec2f(
+        f32(invocation_id.x) / f32(num_workgroups.x * 8),
+        f32(invocation_id.y) / f32(num_workgroups.y * 8)
+    );
     let viewport_up_direction = vec3f(0., 1., 0.);
     let viewport_right_direction = normalize(cross(
-        viewport_up_direction, viewport_.direction
+        viewport_up_direction, viewport.direction
     ));
-    let 
-    viewport_bottom_left = viewport_.origin 
-        + (viewport_.direction * viewport_.fov)
-        - (viewport_right_direction * (viewport_.size.x / 2.))
-        - (viewport_up_direction * (viewport_.size.y / 2.))
+    let viewport_bottom_left = viewport.origin 
+        + (viewport.direction * viewport.fov)
+        - (viewport_right_direction * (viewport.size.x / 2.))
+        - (viewport_up_direction * (viewport.size.y / 2.))
         ;
     let ray_endpoint = viewport_bottom_left
-        + viewport_right_direction * viewport_.size.x * mesh.uv.x
-        + viewport_up_direction * viewport_.size.y * (1. - mesh.uv.y)
+        + viewport_right_direction * viewport.size.x * f32(pixel_location_normalized.x)
+        + viewport_up_direction * viewport.size.y * (1. - f32(pixel_location_normalized.y))
         ;
-    var ray = Line(ray_endpoint, normalize(ray_endpoint - viewport_.origin));
+    var ray = Line(ray_endpoint, normalize(ray_endpoint - viewport.origin));
 
     var ray_result = get_by_ray(ray);
     var rgb_result = vec3f(0.5,0.5,0.5);
@@ -711,7 +720,8 @@ fn fragment(mesh: VertexOutput) -> @location(0) vec4<f32> {
         let result_with_lights = ray_result.albedo.rgb * diffuse_light_strength;
         rgb_result = result_with_lights.rgb;
     }
-    return vec4<f32>(rgb_result, 1.0);
+
+    textureStore(output_texture, pixel_location, vec4f(rgb_result, 1.));
 }
 
 // Note: should be const

examples/bevy_wgpu_render.rs

@@ -1,14 +1,27 @@
 #[cfg(feature = "bevy_wgpu")]
-use bevy::prelude::*;
+use bevy::{prelude::*, window::WindowPlugin};
+
+#[cfg(feature = "bevy_wgpu")]
+use shocovox_rs::octree::{
+    raytracing::{ShocoVoxRenderPlugin, ShocoVoxViewingGlass, Viewport},
+    V3c,
+};
+
+#[cfg(feature = "bevy_wgpu")]
+const DISPLAY_RESOLUTION: [u32; 2] = [1024, 768];
+
 #[cfg(feature = "bevy_wgpu")]
-use shocovox_rs::{octree::raytracing::OctreeViewMaterial, octree::V3c};
+const ARRAY_DIMENSION: u32 = 64;
 
 #[cfg(feature = "bevy_wgpu")]
 fn main() {
     App::new()
+        .insert_resource(ClearColor(Color::BLACK))
         .add_plugins((
-            DefaultPlugins,
-            MaterialPlugin::<OctreeViewMaterial>::default(),
+            DefaultPlugins.set(WindowPlugin::default()),
+            ShocoVoxRenderPlugin {
+                resolution: DISPLAY_RESOLUTION,
+            },
         ))
         .add_systems(Startup, setup)
         .add_systems(Update, rotate_camera)
@@ -17,38 +30,12 @@ fn main() {
 }
 
 #[cfg(feature = "bevy_wgpu")]
-const ARRAY_DIMENSION: u32 = 64;
-
-#[cfg(feature = "bevy_wgpu")]
-fn setup(
-    mut commands: Commands,
-    mut meshes: ResMut<Assets<Mesh>>,
-    mut materials: ResMut<Assets<OctreeViewMaterial>>,
-) {
-    use shocovox_rs::octree::{raytracing::Viewport, types::Albedo};
-
-    commands.spawn(PointLightBundle {
-        point_light: PointLight {
-            intensity: 3000.0,
-            ..Default::default()
-        },
-        transform: Transform::from_xyz(-3.0, 2.0, -1.0),
-        ..Default::default()
-    });
-    commands.spawn(PointLightBundle {
-        point_light: PointLight {
-            intensity: 3000.0,
-            ..Default::default()
-        },
-        transform: Transform::from_xyz(3.0, 2.0, 1.0),
-        ..Default::default()
-    });
-
-    commands.spawn(Camera3dBundle {
-        transform: Transform::from_xyz(5.0, 5.0, 7.0).looking_at(Vec3::new(4., 1., 0.0), Vec3::Y),
-        ..Default::default()
-    });
-
+fn setup(mut commands: Commands, images: ResMut<Assets<Image>>) {
+    let origin = Vec3::new(
+        ARRAY_DIMENSION as f32 * 2.,
+        ARRAY_DIMENSION as f32 / 2.,
+        ARRAY_DIMENSION as f32 * -2.,
+    );
     commands.spawn(DomePosition { yaw: 0. });
 
     // fill octree with data
@@ -95,27 +82,26 @@ fn setup(
             }
         }
     }
-    let quad_size = 10.;
-    let mesh_handle = meshes.add(Mesh::from(Rectangle {
-        half_size: Vec2::new(quad_size, quad_size) / 2.,
-    }));
-    let origin = Vec3::new(
-        ARRAY_DIMENSION as f32 * 2.,
-        ARRAY_DIMENSION as f32 / 2.,
-        ARRAY_DIMENSION as f32 * -2.,
+    let viewing_glass = tree.create_bevy_view(
+        &Viewport {
+            direction: (Vec3::new(0., 0., 0.) - origin).normalize(),
+            origin,
+            size: Vec2::new(10., 10.),
+            fov: 3.,
+        },
+        DISPLAY_RESOLUTION,
+        images,
     );
-    let material_handle = materials.add(tree.create_bevy_material_view(&Viewport {
-        direction: (Vec3::new(0., 0., 0.) - origin).normalize(),
-        origin,
-        size: Vec2::new(10., 10.),
-        fov: 3.,
-    }));
-    commands.spawn(MaterialMeshBundle {
-        mesh: mesh_handle.clone(),
-        material: material_handle.clone(),
-        transform: Transform::from_xyz(quad_size + 0.5, 0.0, 0.0),
-        ..Default::default()
+    commands.spawn(SpriteBundle {
+        sprite: Sprite {
+            custom_size: Some(Vec2::new(1024., 768.)),
+            ..default()
+        },
+        texture: viewing_glass.output_texture.clone(),
+        ..default()
     });
+    commands.spawn(Camera2dBundle::default());
+    commands.insert_resource(viewing_glass);
 }
 
 #[cfg(feature = "bevy_wgpu")]
@@ -127,10 +113,10 @@ struct DomePosition {
 #[cfg(feature = "bevy_wgpu")]
 fn rotate_camera(
     mut angles_query: Query<&mut DomePosition>,
-    mut mats: ResMut<Assets<OctreeViewMaterial>>,
+    mut viewing_glass: ResMut<ShocoVoxViewingGlass>,
 ) {
     let angle = {
-        let addition = ARRAY_DIMENSION as f32 / 1024.;
+        let addition = ARRAY_DIMENSION as f32 / 10.;
         let angle = angles_query.single().yaw + addition;
         if angle < 360. {
             angle
@@ -140,31 +126,27 @@ fn rotate_camera(
     };
     angles_query.single_mut().yaw = angle;
 
-    for (_mat_handle, mat) in mats.as_mut().iter_mut() {
-        let radius = ARRAY_DIMENSION as f32 * 1.3;
-        mat.viewport.origin = Vec3::new(
-            ARRAY_DIMENSION as f32 / 2. + angle.sin() * radius,
-            ARRAY_DIMENSION as f32 / 2.,
-            ARRAY_DIMENSION as f32 / 2. + angle.cos() * radius,
-        );
-        mat.viewport.direction = (Vec3::new(
-            ARRAY_DIMENSION as f32 / 2.,
-            ARRAY_DIMENSION as f32 / 2.,
-            ARRAY_DIMENSION as f32 / 2.,
-        ) - mat.viewport.origin)
-            .normalize();
-    }
+    let radius = ARRAY_DIMENSION as f32 * 1.3;
+    viewing_glass.viewport.origin = Vec3::new(
+        ARRAY_DIMENSION as f32 / 2. + angle.sin() * radius,
+        ARRAY_DIMENSION as f32 / 2.,
+        ARRAY_DIMENSION as f32 / 2. + angle.cos() * radius,
+    );
+    viewing_glass.viewport.direction = (Vec3::new(
+        ARRAY_DIMENSION as f32 / 2.,
+        ARRAY_DIMENSION as f32 / 2.,
+        ARRAY_DIMENSION as f32 / 2.,
+    ) - viewing_glass.viewport.origin)
+        .normalize();
 }
 
 #[cfg(feature = "bevy_wgpu")]
-fn handle_zoom(keys: Res<ButtonInput<KeyCode>>, mut mats: ResMut<Assets<OctreeViewMaterial>>) {
-    for (_mat_handle, mat) in mats.as_mut().iter_mut() {
-        if keys.pressed(KeyCode::ArrowUp) {
-            mat.viewport.size *= 1.1;
-        }
-        if keys.pressed(KeyCode::ArrowDown) {
-            mat.viewport.size *= 0.9;
-        }
+fn handle_zoom(keys: Res<ButtonInput<KeyCode>>, mut viewing_glass: ResMut<ShocoVoxViewingGlass>) {
+    if keys.pressed(KeyCode::ArrowUp) {
+        viewing_glass.viewport.size *= 1.1;
+    }
+    if keys.pressed(KeyCode::ArrowDown) {
+        viewing_glass.viewport.size *= 0.9;
     }
 }
 

src/octree/raytracing/classic_raytracing_on_bevy_wgpu.rs → src/octree/raytracing/bevy/data.rs

@@ -1,22 +1,18 @@
 use crate::object_pool::empty_marker;
+
 use crate::octree::{
-    raytracing::types::{OctreeMetaData, OctreeViewMaterial, SizedNode, Viewport, Voxelement},
+    raytracing::{
+        bevy::create_ouput_texture,
+        bevy::types::{OctreeMetaData, ShocoVoxViewingGlass, SizedNode, Viewport, Voxelement},
+    },
     types::{NodeChildrenArray, NodeContent},
+    Octree, VoxelData,
 };
 
 use bevy::{
-    math::Vec3,
-    pbr::Material,
-    render::{color::Color, render_resource::ShaderRef},
+    asset::Assets, ecs::system::ResMut, math::Vec3, render::color::Color, render::texture::Image,
 };
 
-impl Material for OctreeViewMaterial {
-    fn fragment_shader() -> ShaderRef {
-        "shaders/viewport_render.wgsl".into()
-    }
-}
-
-use crate::octree::{Octree, VoxelData};
 impl<T, const DIM: usize> Octree<T, DIM>
 where
     T: Default + Clone + VoxelData,
@@ -58,7 +54,12 @@ where
         meta
     }
 
-    pub fn create_bevy_material_view(&self, viewport: &Viewport) -> OctreeViewMaterial {
+    pub fn create_bevy_view(
+        &self,
+        viewport: &Viewport,
+        resolution: [u32; 2],
+        images: ResMut<Assets<Image>>,
+    ) -> ShocoVoxViewingGlass {
         let meta = OctreeMetaData {
             octree_size: self.octree_size,
             voxel_matrix_dim: DIM as u32,
@@ -113,7 +114,9 @@ where
             }
             nodes.push(sized_node);
         }
-        OctreeViewMaterial {
+
+        ShocoVoxViewingGlass {
+            output_texture: create_ouput_texture(resolution, images),
             viewport: *viewport,
             meta,
             nodes,