diff --git a/assets/trace-compute.wgsl b/assets/trace-compute.wgsl new file mode 100644 index 0000000..a7c6e4f --- /dev/null +++ b/assets/trace-compute.wgsl @@ -0,0 +1,226 @@ + +@group(0) @binding(0) var input: texture_storage_2d; + +@group(0) @binding(1) var output: texture_storage_2d; + +@group(0) @binding(2) var config: TracerUniforms; +@group(0) @binding(3) var skybox_texture: texture_cube; +@group(0) @binding(4) var skybox_sampler: sampler; + +struct View { + view_proj: mat4x4, + view: mat4x4, + projection: mat4x4, + inv_view_proj: mat4x4, + inv_view: mat4x4, // equiv to Unity's _CameraToWorld + inv_projection: mat4x4, // equic to Unity's _CameraInverseProjection +}; + +struct TracerUniforms { + sky_color: vec4, + world_from_clip: mat4x4, + world_position: vec3, +} + +struct Object { +// 0: shpere +// 1: plane +// 2: cube + type: u32, + position: vec3, + rotation: vec4, + scale: vec3 +} + +@compute @workgroup_size(8, 8, 1) +fn init(@builtin(global_invocation_id) invocation_id: vec3, @builtin(num_workgroups) num_workgroups: vec3) { + let location = vec2(i32(invocation_id.x), i32(invocation_id.y)); + + let color = vec4(0.0); + + textureStore(output, location, color); +} + + + +@compute @workgroup_size(8, 8, 1) +fn update(@builtin(global_invocation_id) invocation_id: vec3) { + let size = textureDimensions(output); + + let loc = vec2(f32(invocation_id.x), f32(invocation_id.y)) / vec2(size.xy); + let ndc = loc * 2.0f - 1.0f; + + var ray = createCameraRay(ndc); + + let hit_data = trace(ray); + + + var final_color = hit_data.color; + + let step_factor = f32(hit_data.steps)/100.0; + final_color = final_color * step_factor; + + + var color = vec4(final_color, 1.0); + if hit_data.distance == 100.0 { + color = config.sky_color; + // color = textureSampleLevel(skybox_texture, skybox_sampler, vec3(0.0, 0.0, 1.0), 0, 0u); + } + let location = vec2(i32(invocation_id.x), i32(invocation_id.y)); + + textureStore(output, location, color); +} + +fn debug_matrix(uv: vec2) -> vec4{ + let ndc = uv * 2.0f - 1.0f; + var color = vec3(0.0, 0.0, 0.0); + + if uv.y < 0.5 { + if uv.x < 0.5 { + color = debugColor(config.world_from_clip[0][3] * 10); + }else{ + color = debugColor(config.world_from_clip[1][3] * 10); + } + }else{ + if uv.x < 0.5 { + color = debugColor(config.world_from_clip[3][2] * 10); + }else{ + color = debugColor(config.world_from_clip[3][3] * 10); + } + } + + + return vec4(color, 1.0); +} +fn debug(uv: vec2) -> vec4{ + let ndc = uv * 2.0f - 1.0f; + + let near_clip = vec4(ndc, 0.0, 1.0); + let far_clip = vec4(ndc, 1.0, 1.0); + + // project into world space + let near_world4 = config.world_from_clip * near_clip; + let far_world4 = config.world_from_clip * far_clip; + + //Add epsilon to protect from divide by zero + let inv_w_near = 1.0 / (near_world4.w + 1e-6); + let inv_w_far = 1.0 / (far_world4.w + 1e-6); + + let near_world = near_world4.xyz * inv_w_near; + let far_world = far_world4.xyz * inv_w_far; + + let origin = config.world_position; + + let direction = normalize(near_world - origin); + + var ray = createRay(origin, direction); + var color = vec3(0.0, 0.0, 0.0); + if uv.y < 0.5 { + if uv.x < 0.5 { + color = debugColor(ray.direction.x); + }else{ + color = debugColor(ray.direction.y); + } + }else{ + if uv.x < 0.5 { + color = debugColor(ray.direction.z); + }else{ + // color = debugColor(ray.direction.z); + color = debugColor(origin.x * 0.1); + } + } + + return vec4(color, 1.0); +} + +fn debugColor(v: f32) -> vec3{ + return vec3(v * 0.5 + 0.5); +} + +struct Ray { + origin: vec3, + direction: vec3, + energy: vec3, +} + +fn createRay(origin: vec3, direction: vec3) -> Ray +{ + var ray: Ray; + ray.origin = origin; + ray.direction = direction; + ray.energy = vec3(1.0f, 1.0f, 1.0f); + return ray; +} + +fn createCameraRay(ndc: vec2) -> Ray { + let uv = vec2(ndc.x, -ndc.y); + + // clip points at near and far + let near_clip = vec4(uv, 0.0, 1.0); + let far_clip = vec4(uv, 1.0, 1.0); + + // project into world space + let near_world4 = config.world_from_clip * near_clip; + let far_world4 = config.world_from_clip * far_clip; + + //Add epsilon to protect from divide by zero + let inv_w_near = 1.0 / (near_world4.w + 1e-6); + let inv_w_far = 1.0 / (far_world4.w + 1e-6); + + let near_world = near_world4.xyz * inv_w_near; + let far_world = far_world4.xyz * inv_w_far; + + let origin = config.world_position; + + let direction = normalize(near_world - origin); + + return createRay(origin, direction); +} + +struct Hit { + distance: f32, + hit_pos: vec3, + normal: vec3, + color: vec3, + direction: vec3, + steps: i32, +}; + +fn distance_field(p: vec3) -> f32 { + // Simple sphere centered at (0, 0, 0) with radius 1.0 + let sphere_center = vec3(0.0, 0.0, 0.0); + let sphere_radius = 1.0; + + // SDF for a sphere: length(p - center) - radius + let d = length(p - sphere_center) - sphere_radius; + return d; +} + +fn trace(ray: Ray) -> Hit { + var total_distance: f32 = 0.0; + let max_distance: f32 = 100.0; + let min_hit_distance: f32 = 0.001; + const max_steps: i32 = 100; + + for (var i: i32 = 0; i < max_steps; i = i + 1) { + let current_pos = ray.origin + ray.direction * total_distance; + let distance_to_scene = distance_field(current_pos); + + // Check for a hit + if (distance_to_scene < min_hit_distance) { + // A hit occurred! + return Hit(total_distance, current_pos, vec3(0.0), vec3(1.0, 0.0, 0.0), ray.direction, i); // Return red color for now + } + + // Check if we marched too far + if (total_distance > max_distance) { + break; + } + + // Advance the ray + total_distance = total_distance + distance_to_scene; + } + + // No hit found + return Hit(max_distance, vec3(0.0), vec3(0.0), vec3(0.0, 0.0, 0.0), ray.direction, max_steps); // Return black for background +} diff --git a/assets/trace.wgsl b/assets/trace.wgsl index 539843e..8f932df 100644 --- a/assets/trace.wgsl +++ b/assets/trace.wgsl @@ -1,43 +1,14 @@ +#import bevy_sprite::mesh2d_vertex_output::VertexOutput -@group(0) @binding(0) var input: texture_storage_2d; +@group(2) @binding(0) var sky_color: vec4; +@group(2) @binding(1) var view: TracerView; +@group(2) @binding(2) var skybox_texture: texture_2d; +@group(2) @binding(3) var skybox_sampler: sampler; -@group(0) @binding(1) var output: texture_storage_2d; +@fragment +fn fragment(mesh: VertexOutput) -> @location(0) vec4 { -@group(0) @binding(2) var config: TracerUniforms; -@group(0) @binding(3) var skybox_texture: texture_cube; -@group(0) @binding(4) var skybox_sampler: sampler; - -struct View { - view_proj: mat4x4, - view: mat4x4, - projection: mat4x4, - inv_view_proj: mat4x4, - inv_view: mat4x4, // equiv to Unity's _CameraToWorld - inv_projection: mat4x4, // equic to Unity's _CameraInverseProjection -}; - -struct TracerUniforms { - sky_color: vec4, - world_from_clip: mat4x4, - world_position: vec3, -} - -@compute @workgroup_size(8, 8, 1) -fn init(@builtin(global_invocation_id) invocation_id: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - let location = vec2(i32(invocation_id.x), i32(invocation_id.y)); - - let color = vec4(0.0); - - textureStore(output, location, color); -} - - - -@compute @workgroup_size(8, 8, 1) -fn update(@builtin(global_invocation_id) invocation_id: vec3) { - let size = textureDimensions(output); - - let loc = vec2(f32(invocation_id.x), f32(invocation_id.y)) / vec2(size.xy); + let loc = mesh.uv; let ndc = loc * 2.0f - 1.0f; var ray = createCameraRay(ndc); @@ -45,92 +16,18 @@ fn update(@builtin(global_invocation_id) invocation_id: vec3) { let hit_data = trace(ray); - var final_color = hit_data.color; + var final_color = hit_data.color; - let step_factor = f32(hit_data.steps)/100.0; - final_color = final_color * step_factor; + let step_factor = f32(hit_data.steps)/100.0; + final_color = final_color * step_factor; - var color = vec4(final_color, 1.0); - if hit_data.distance == 100.0 { - color = config.sky_color; - // color = textureSampleLevel(skybox_texture, skybox_sampler, vec3(0.0, 0.0, 1.0), 0, 0u); - } - let location = vec2(i32(invocation_id.x), i32(invocation_id.y)); - - textureStore(output, location, color); -} - -fn debug_matrix(uv: vec2) -> vec4{ - let ndc = uv * 2.0f - 1.0f; - var color = vec3(0.0, 0.0, 0.0); - - if uv.y < 0.5 { - if uv.x < 0.5 { - color = debugColor(config.world_from_clip[0][3] * 10); - }else{ - color = debugColor(config.world_from_clip[1][3] * 10); - } - }else{ - if uv.x < 0.5 { - color = debugColor(config.world_from_clip[3][2] * 10); - }else{ - color = debugColor(config.world_from_clip[3][3] * 10); - } - } - - - return vec4(color, 1.0); -} -fn debug(uv: vec2) -> vec4{ - let ndc = uv * 2.0f - 1.0f; - - let near_clip = vec4(ndc, 0.0, 1.0); - let far_clip = vec4(ndc, 1.0, 1.0); - - // project into world space - let near_world4 = config.world_from_clip * near_clip; - let far_world4 = config.world_from_clip * far_clip; - - //Add epsilon to protect from divide by zero - let inv_w_near = 1.0 / (near_world4.w + 1e-6); - let inv_w_far = 1.0 / (far_world4.w + 1e-6); - - let near_world = near_world4.xyz * inv_w_near; - let far_world = far_world4.xyz * inv_w_far; - - let origin = config.world_position; - - let direction = normalize(near_world - origin); - - var ray = createRay(origin, direction); - var color = vec3(0.0, 0.0, 0.0); - if uv.y < 0.5 { - if uv.x < 0.5 { - color = debugColor(ray.direction.x); - }else{ - color = debugColor(ray.direction.y); - } - }else{ - if uv.x < 0.5 { - color = debugColor(ray.direction.z); - }else{ - // color = debugColor(ray.direction.z); - color = debugColor(origin.x * 0.1); - } - } - - return vec4(color, 1.0); -} - -fn debugColor(v: f32) -> vec3{ - return vec3(v * 0.5 + 0.5); -} - -struct Ray { - origin: vec3, - direction: vec3, - energy: vec3, + var color = vec4(final_color, 1.0); + if hit_data.distance == 100.0 { + color = sky_color; + // color = textureSampleLevel(skybox_texture, skybox_sampler, vec3(0.0, 0.0, 1.0), 0, 0u); + } + return color; } fn createRay(origin: vec3, direction: vec3) -> Ray @@ -150,8 +47,8 @@ fn createCameraRay(ndc: vec2) -> Ray { let far_clip = vec4(uv, 1.0, 1.0); // project into world space - let near_world4 = config.world_from_clip * near_clip; - let far_world4 = config.world_from_clip * far_clip; + let near_world4 = view.world_from_clip * near_clip; + let far_world4 = view.world_from_clip * far_clip; //Add epsilon to protect from divide by zero let inv_w_near = 1.0 / (near_world4.w + 1e-6); @@ -160,22 +57,13 @@ fn createCameraRay(ndc: vec2) -> Ray { let near_world = near_world4.xyz * inv_w_near; let far_world = far_world4.xyz * inv_w_far; - let origin = config.world_position; + let origin = view.world_position; let direction = normalize(near_world - origin); return createRay(origin, direction); } -struct Hit { - distance: f32, - hit_pos: vec3, - normal: vec3, - color: vec3, - direction: vec3, - steps: i32, -}; - fn distance_field(p: vec3) -> f32 { // Simple sphere centered at (0, 0, 0) with radius 1.0 let sphere_center = vec3(0.0, 0.0, 0.0); @@ -214,3 +102,33 @@ fn trace(ray: Ray) -> Hit { // No hit found return Hit(max_distance, vec3(0.0), vec3(0.0), vec3(0.0, 0.0, 0.0), ray.direction, max_steps); // Return black for background } + +struct TracerView { + world_from_clip: mat4x4, + world_position: vec3, +} + +struct Object { +// 0: shpere +// 1: plane +// 2: cube + object_type: u32, + position: vec3, + rotation: vec4, + scale: vec3 +} + +struct Ray { + origin: vec3, + direction: vec3, + energy: vec3, +} + +struct Hit { + distance: f32, + hit_pos: vec3, + normal: vec3, + color: vec3, + direction: vec3, + steps: i32, +}; diff --git a/src/app.rs b/src/app.rs index 501131d..bce9dab 100644 --- a/src/app.rs +++ b/src/app.rs @@ -12,8 +12,8 @@ use bevy::{ use iyes_perf_ui::prelude::*; use crate::{ - components::rt::RTCamera, - render::pipeline::{TracerPipelinePlugin, TracerRenderTextures, TracerUniforms}, + components::rt::{RTCamera, RTDisplay}, + render::{tracer::TracerPlugin, tracer_material::TracerMaterial}, }; pub struct Blackhole; @@ -32,18 +32,12 @@ struct SkyboxAsset(Handle); impl Plugin for Blackhole { fn build(&self, app: &mut App) { - app.register_type::(); - app.init_state::(); app.add_systems(Startup, setup) .add_systems(Update, asset_load_check.run_if(in_state(AssetLoad::Loading))) - .add_systems(Update, prepare_skybox.run_if(in_state(AssetLoad::Init))) + // .add_systems(Update, prepare_skybox.run_if(in_state(AssetLoad::Init))) .add_systems(Last, asset_init.run_if(in_state(AssetLoad::Init))); - app.add_plugins(TracerPipelinePlugin); - app.insert_resource(TracerUniforms { - sky_color: LinearRgba::rgb(0.1, 0.0, 0.01), - ..default() - }); + app.add_plugins(TracerPlugin); //Perf UI app.add_plugins(bevy::diagnostic::FrameTimeDiagnosticsPlugin::default()) @@ -55,10 +49,10 @@ impl Plugin for Blackhole { fn setup( mut commands: Commands, - mut images: ResMut>, - window: Single<&Window, With>, asset_server: Res, mut load_state: ResMut>, + mut materials: ResMut>, + mut meshes: ResMut>, ) { commands.spawn(( PerfUiRoot::default(), @@ -68,58 +62,16 @@ fn setup( PerfUiEntryFrameTimeWorst::default(), )); - let size = window.physical_size(); - - let extent = Extent3d { - width: size.x, - height: size.y, - ..Default::default() - }; - - const PIXEL_FORMAT: TextureFormat = TextureFormat::Rgba32Float; - const PIXEL_SIZE: usize = 16; - let mut image = Image::new_fill( - extent, - TextureDimension::D2, - &[255; PIXEL_SIZE], - PIXEL_FORMAT, - RenderAssetUsages::RENDER_WORLD, - ); - image.texture_descriptor.usage = TextureUsages::TEXTURE_BINDING - | TextureUsages::STORAGE_BINDING - | TextureUsages::COPY_DST - | TextureUsages::RENDER_ATTACHMENT; - - let render0 = images.add(image.clone()); - let render1 = images.add(image); - - let mut skybox_render_image = Image::new_fill( - Extent3d { - width: 256, - height: 1536, - ..Default::default() - }, - TextureDimension::D2, - &[255; PIXEL_SIZE], - PIXEL_FORMAT, - RenderAssetUsages::RENDER_WORLD, - ); - skybox_render_image.reinterpret_stacked_2d_as_array(6); - skybox_render_image.texture_descriptor.usage = TextureUsages::TEXTURE_BINDING | TextureUsages::RENDER_ATTACHMENT; - skybox_render_image.texture_view_descriptor = Some(TextureViewDescriptor { - dimension: Some(TextureViewDimension::Cube), - ..default() - }); - let skybox_render_image_handle = images.add(skybox_render_image); - let skybox_asset = asset_server.load("sky-array.png"); commands.spawn(( - Name::new("Render Sprite"), - Sprite { - image: render0.clone(), - custom_size: Some(size.as_vec2()), + Name::new("Render Display"), + MeshMaterial2d(materials.add(TracerMaterial { + sky_color: LinearRgba::rgb(0.1, 0.0, 0.01), + skybox: Some(skybox_asset.clone()), ..default() - }, + })), + RTDisplay, + Mesh2d(meshes.add(Rectangle::from_size(vec2(1920.0, 1080.0)))), Transform::from_translation(Vec3::ZERO), )); @@ -128,11 +80,6 @@ fn setup( commands .spawn(( Camera3d::default(), - // Camera { order: -1, ..default() }, - // Projection::Perspective(PerspectiveProjection { - // aspect_ratio: size.x as f32 / size.y as f32, - // ..default() - // }), RTCamera, RenderLayers::layer(1), Transform::from_xyz(0.0, 5.0, 20.0).looking_at(Vec3::ZERO, Vec3::Y), @@ -140,11 +87,6 @@ fn setup( )) .insert(Camera { order: -1, ..default() }); - commands.insert_resource(TracerRenderTextures { - main: render0, - secondary: render1, - skybox: skybox_render_image_handle, - }); commands.insert_resource(SkyboxAsset(skybox_asset)); load_state.set(AssetLoad::Loading); @@ -162,11 +104,7 @@ fn asset_load_check( } } -fn prepare_skybox( - tracer_textures: Res, - skybox: Res, - mut image_assets: ResMut>, -) { +fn prepare_skybox(skybox: Res, mut image_assets: ResMut>) { let mut skybox_image = image_assets .get(skybox.0.id()) .expect("Skybox asset image does not exist") @@ -176,7 +114,7 @@ fn prepare_skybox( dimension: Some(TextureViewDimension::Cube), ..default() }); - image_assets.insert(tracer_textures.skybox.id(), skybox_image.clone()); + image_assets.insert(skybox.0.id(), skybox_image); } fn asset_init(mut load_state: ResMut>) { diff --git a/src/components/rt.rs b/src/components/rt.rs index f6474ff..e371116 100644 --- a/src/components/rt.rs +++ b/src/components/rt.rs @@ -1,4 +1,7 @@ use bevy::prelude::*; -#[derive(Component)] +#[derive(Component, Debug)] pub struct RTCamera; + +#[derive(Component, Debug)] +pub struct RTDisplay; diff --git a/src/render/mod.rs b/src/render/mod.rs index 07caabe..cb5aaee 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -1,2 +1,2 @@ -pub mod node; -pub mod pipeline; +pub mod tracer; +pub mod tracer_material; diff --git a/src/render/node.rs b/src/render/node.rs deleted file mode 100644 index 9bc9d73..0000000 --- a/src/render/node.rs +++ /dev/null @@ -1,112 +0,0 @@ -use bevy::{ - prelude::*, - render::{ - render_graph::{self}, - render_resource::{CachedPipelineState, ComputePassDescriptor, PipelineCache, PipelineCacheError}, - renderer::RenderContext, - }, -}; - -use crate::render::pipeline::{TracerImageBindGroups, TracerPipeline, TracerRenderTextures}; -use crate::{SHADER_ASSET_PATH, WORKGROUP_SIZE}; - -pub enum TracerState { - Loading, - Init, - Update(usize), -} - -pub struct TracerNode { - state: TracerState, -} - -impl Default for TracerNode { - fn default() -> Self { - Self { - state: TracerState::Loading, - } - } -} - -impl render_graph::Node for TracerNode { - fn update(&mut self, world: &mut World) { - let pipeline = world.resource::(); - let pipeline_cache = world.resource::(); - - // if the corresponding pipeline has loaded, transition to the next stage - match self.state { - TracerState::Loading => { - match pipeline_cache.get_compute_pipeline_state(pipeline.init_pipeline) { - CachedPipelineState::Ok(_) => { - self.state = TracerState::Init; - } - // If the shader hasn't loaded yet, just wait. - CachedPipelineState::Err(PipelineCacheError::ShaderNotLoaded(_)) => (), - CachedPipelineState::Err(err) => { - panic!("Initializing assets/{SHADER_ASSET_PATH}:\n{err}") - } - _ => (), - }; - - // let tex = world.resource::(); - // let asset_server = world.resource::(); - // let load_state = asset_server.get_load_state(tex.skybox.id()).expect(); - // if load_state.is_loaded() && shader_loaded { - // self.state = TracerState::Init; - // } - } - TracerState::Init => { - if let CachedPipelineState::Ok(_) = pipeline_cache.get_compute_pipeline_state(pipeline.update_pipeline) - { - self.state = TracerState::Update(1); - } - } - TracerState::Update(0) => { - self.state = TracerState::Update(1); - } - TracerState::Update(1) => { - self.state = TracerState::Update(0); - } - TracerState::Update(_) => unreachable!(), - } - } - - fn run( - &self, - _graph: &mut render_graph::RenderGraphContext, - render_context: &mut RenderContext, - world: &World, - ) -> Result<(), render_graph::NodeRunError> { - let bind_groups_res = world.get_resource::(); - if bind_groups_res.is_none() { - return Ok(()); - } - let bind_groups = &bind_groups_res.unwrap().0; - let pipeline_cache = world.resource::(); - let pipeline = world.resource::(); - - let mut pass = render_context - .command_encoder() - .begin_compute_pass(&ComputePassDescriptor::default()); - - // select the pipeline based on the current state - match self.state { - TracerState::Loading => {} - TracerState::Init => { - let init_pipeline = pipeline_cache.get_compute_pipeline(pipeline.init_pipeline).unwrap(); - pass.set_bind_group(0, &bind_groups[0], &[]); - pass.set_pipeline(init_pipeline); - pass.dispatch_workgroups(1920 / WORKGROUP_SIZE, 1080 / WORKGROUP_SIZE, 1); - } - TracerState::Update(index) => { - if let Some(update_pipeline) = pipeline_cache.get_compute_pipeline(pipeline.update_pipeline) { - pass.set_bind_group(0, &bind_groups[index], &[]); - pass.set_pipeline(update_pipeline); - pass.dispatch_workgroups(1920 / WORKGROUP_SIZE, 1080 / WORKGROUP_SIZE, 1); - } - } - } - - Ok(()) - } -} diff --git a/src/render/pipeline.rs b/src/render/pipeline.rs deleted file mode 100644 index 0545bf3..0000000 --- a/src/render/pipeline.rs +++ /dev/null @@ -1,236 +0,0 @@ -use std::borrow::Cow; - -use bevy::{ - prelude::*, - render::{ - Render, RenderApp, RenderSet, - extract_resource::{ExtractResource, ExtractResourcePlugin}, - render_asset::RenderAssets, - render_graph::{RenderGraph, RenderLabel}, - render_resource::{ - BindGroup, BindGroupEntries, BindGroupLayout, BindGroupLayoutEntries, CachedComputePipelineId, - ComputePipelineDescriptor, PipelineCache, SamplerBindingType, ShaderStages, ShaderType, - StorageTextureAccess, TextureFormat, TextureSampleType, UniformBuffer, - binding_types::{sampler, texture_cube, texture_storage_2d, uniform_buffer}, - }, - renderer::{RenderDevice, RenderQueue}, - texture::GpuImage, - }, -}; - -use crate::{SHADER_ASSET_PATH, app::AssetLoad, components::rt::RTCamera, render::node::TracerNode}; - -#[derive(Debug, Hash, PartialEq, Eq, Clone, RenderLabel)] -pub struct TracerLabel; - -#[derive(Resource, Reflect, ExtractResource, Clone)] -#[reflect(Resource)] -pub struct TracerRenderTextures { - pub main: Handle, - pub secondary: Handle, - pub skybox: Handle, -} - -#[derive(Resource, Clone, ExtractResource, ShaderType, Default)] -pub struct TracerUniforms { - pub sky_color: LinearRgba, - pub world_from_clip: Mat4, - pub world_position: Vec3, -} - -pub struct TracerPipelinePlugin; - -impl Plugin for TracerPipelinePlugin { - fn build(&self, app: &mut App) { - app.add_plugins(( - ExtractResourcePlugin::::default(), - ExtractResourcePlugin::::default(), - )); - app.init_resource::() - .add_systems(Update, switch_textures.run_if(in_state(AssetLoad::Ready))); - - app.add_systems(First, update_tracer_uniforms.run_if(in_state(AssetLoad::Ready))); - let render_app = app.sub_app_mut(RenderApp); - - // render_app.add_systems(Startup, init_pipeline); - render_app.add_systems(Render, prepare_bind_groups.in_set(RenderSet::PrepareBindGroups)); - - let mut render_graph = render_app.world_mut().resource_mut::(); - render_graph.add_node(TracerLabel, TracerNode::default()); - render_graph.add_node_edge(TracerLabel, bevy::render::graph::CameraDriverLabel); - } - - fn finish(&self, app: &mut App) { - let render_app = app.sub_app_mut(RenderApp); - render_app.init_resource::(); - } -} - -fn switch_textures(images: Res, mut sprite: Single<&mut Sprite>) { - if sprite.image == images.main { - sprite.image = images.secondary.clone(); - } else { - sprite.image = images.main.clone(); - } -} - -#[derive(Resource)] -pub struct TracerPipeline { - pub texture_bind_group_layout: BindGroupLayout, - pub init_pipeline: CachedComputePipelineId, - pub update_pipeline: CachedComputePipelineId, -} - -impl FromWorld for TracerPipeline { - fn from_world(world: &mut World) -> Self { - let render_device = world.resource::(); - - let texture_bind_group_layout = render_device.create_bind_group_layout( - "TracerImages", - &BindGroupLayoutEntries::sequential( - ShaderStages::COMPUTE, - ( - texture_storage_2d(TextureFormat::Rgba32Float, StorageTextureAccess::ReadOnly), - texture_storage_2d(TextureFormat::Rgba32Float, StorageTextureAccess::WriteOnly), - uniform_buffer::(false), - texture_cube(TextureSampleType::Float { filterable: true }), - sampler(SamplerBindingType::Filtering), - ), - ), - ); - let shader = world.load_asset(SHADER_ASSET_PATH); - let pipeline_cache = world.resource::(); - let init_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { - layout: vec![texture_bind_group_layout.clone()], - shader: shader.clone(), - entry_point: Cow::from("init"), - label: None, - zero_initialize_workgroup_memory: false, - push_constant_ranges: Default::default(), - shader_defs: Default::default(), - }); - - let update_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { - layout: vec![texture_bind_group_layout.clone()], - shader, - entry_point: Cow::from("update"), - label: None, - zero_initialize_workgroup_memory: false, - push_constant_ranges: Default::default(), - shader_defs: Default::default(), - }); - - return TracerPipeline { - texture_bind_group_layout, - init_pipeline, - update_pipeline, - }; - } -} - -#[allow(dead_code)] //Pending bevy update for RenderStartup schedule -fn init_pipeline( - mut commands: Commands, - render_device: Res, - asset_server: Res, - pipeline_cache: Res, -) { - let texture_bind_group_layout = render_device.create_bind_group_layout( - "TracerImages", - &BindGroupLayoutEntries::sequential( - ShaderStages::COMPUTE, - ( - texture_storage_2d(TextureFormat::Rgba32Float, StorageTextureAccess::ReadOnly), - texture_storage_2d(TextureFormat::Rgba32Float, StorageTextureAccess::WriteOnly), - uniform_buffer::(false), - texture_cube(TextureSampleType::Float { filterable: true }), - sampler(SamplerBindingType::Filtering), - ), - ), - ); - let shader = asset_server.load(SHADER_ASSET_PATH); - let init_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { - layout: vec![texture_bind_group_layout.clone()], - shader: shader.clone(), - entry_point: Cow::from("init"), - label: None, - zero_initialize_workgroup_memory: false, - push_constant_ranges: Default::default(), - shader_defs: Default::default(), - }); - - let update_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { - layout: vec![texture_bind_group_layout.clone()], - shader, - entry_point: Cow::from("update"), - label: None, - zero_initialize_workgroup_memory: false, - push_constant_ranges: Default::default(), - shader_defs: Default::default(), - }); - - commands.insert_resource(TracerPipeline { - texture_bind_group_layout, - init_pipeline, - update_pipeline, - }); -} - -#[derive(Resource)] -pub struct TracerImageBindGroups(pub [BindGroup; 2]); - -fn update_tracer_uniforms( - mut tracer_uniforms: ResMut, - rt_camera: Single<(&GlobalTransform, &Camera), With>, -) { - let (transform, cam) = rt_camera.into_inner(); - - let clip_from_view = cam.clip_from_view(); - let world_from_clip = transform.compute_matrix() * clip_from_view.inverse(); - - tracer_uniforms.world_from_clip = world_from_clip; - tracer_uniforms.world_position = transform.translation(); -} - -fn prepare_bind_groups( - mut commands: Commands, - pipeline: Res, - gpu_images: Res>, - tracer_images: Res, - tracer_uniforms: Res, - render_device: Res, - queue: Res, -) { - let view_a = gpu_images.get(&tracer_images.main).unwrap(); - let view_b = gpu_images.get(&tracer_images.secondary).unwrap(); - let skybox = gpu_images.get(&tracer_images.skybox).unwrap(); - - // Uniform buffer is used here to demonstrate how to set up a uniform in a compute shader - // Alternatives such as storage buffers or push constants may be more suitable for your use case - let mut uniform_buffer = UniformBuffer::from(tracer_uniforms.into_inner()); - uniform_buffer.write_buffer(&render_device, &queue); - - let bind_group_0 = render_device.create_bind_group( - None, - &pipeline.texture_bind_group_layout, - &BindGroupEntries::sequential(( - &view_a.texture_view, - &view_b.texture_view, - &uniform_buffer, - &skybox.texture_view, - &skybox.sampler, - )), - ); - let bind_group_1 = render_device.create_bind_group( - None, - &pipeline.texture_bind_group_layout, - &BindGroupEntries::sequential(( - &view_b.texture_view, - &view_a.texture_view, - &uniform_buffer, - &skybox.texture_view, - &skybox.sampler, - )), - ); - commands.insert_resource(TracerImageBindGroups([bind_group_0, bind_group_1])); -} diff --git a/src/render/tracer.rs b/src/render/tracer.rs new file mode 100644 index 0000000..c0afbb7 --- /dev/null +++ b/src/render/tracer.rs @@ -0,0 +1,35 @@ +use bevy::{prelude::*, sprite::Material2dPlugin}; + +use crate::{ + app::AssetLoad, + components::rt::{RTCamera, RTDisplay}, + render::tracer_material::{TracerMaterial, TracerView}, +}; + +pub struct TracerPlugin; + +impl Plugin for TracerPlugin { + fn build(&self, app: &mut App) { + app.add_plugins(Material2dPlugin::::default()); + app.add_systems(First, update_tracer_uniforms.run_if(in_state(AssetLoad::Ready))); + } +} + +fn update_tracer_uniforms( + rt_camera: Single<(&GlobalTransform, &Camera), With>, + display: Single<&MeshMaterial2d, With>, + mut materials: ResMut>, +) { + let (transform, cam) = rt_camera.into_inner(); + + let clip_from_view = cam.clip_from_view(); + let world_from_clip = transform.compute_matrix() * clip_from_view.inverse(); + + let mat = materials + .get_mut(display.0.id()) + .expect("Tracer Materials doesn't exist"); + mat.view = TracerView { + world_from_clip: world_from_clip, + world_position: transform.translation(), + }; +} diff --git a/src/render/tracer_material.rs b/src/render/tracer_material.rs new file mode 100644 index 0000000..db36742 --- /dev/null +++ b/src/render/tracer_material.rs @@ -0,0 +1,35 @@ +use bevy::{ + prelude::*, + reflect::Reflect, + render::render_resource::{AsBindGroup, ShaderType}, + sprite::Material2d, +}; + +use crate::SHADER_ASSET_PATH; + +#[derive(Asset, TypePath, AsBindGroup, Debug, Clone, Default)] +pub struct TracerMaterial { + #[uniform(0)] + pub sky_color: LinearRgba, + #[uniform(1)] + pub view: TracerView, + #[texture(2)] + #[sampler(3)] + pub skybox: Option>, +} + +#[derive(Debug, ShaderType, Clone, Reflect, Default)] +pub struct TracerView { + pub world_from_clip: Mat4, + pub world_position: Vec3, +} + +impl Material2d for TracerMaterial { + fn fragment_shader() -> bevy::render::render_resource::ShaderRef { + SHADER_ASSET_PATH.into() + } + + fn alpha_mode(&self) -> bevy::sprite::AlphaMode2d { + bevy::sprite::AlphaMode2d::Opaque + } +}