Amatsugu/black-hole
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

migrate to a fragment shader

Browse Source
This commit is contained in:
Amatsugu
2026-02-19 12:11:44 -05:00
parent dd7232d5e7
commit b62acf3277
9 changed files with 367 additions and 560 deletions
Download Patch File Download Diff File Expand all files Collapse all files

226
assets/trace-compute.wgsl Normal file
View File

@@ -0,0 +1,226 @@
@group(0) @binding(0) var input: texture_storage_2d<rgba32float, read>;
@group(0) @binding(1) var output: texture_storage_2d<rgba32float, write>;
@group(0) @binding(2) var<uniform> config: TracerUniforms;
@group(0) @binding(3) var skybox_texture: texture_cube<f32>;
@group(0) @binding(4) var skybox_sampler: sampler;
struct View {
view_proj: mat4x4<f32>,
view: mat4x4<f32>,
projection: mat4x4<f32>,
inv_view_proj: mat4x4<f32>,
inv_view: mat4x4<f32>, // equiv to Unity's _CameraToWorld
inv_projection: mat4x4<f32>, // equic to Unity's _CameraInverseProjection
};
struct TracerUniforms {
sky_color: vec4<f32>,
world_from_clip: mat4x4<f32>,
world_position: vec3<f32>,
}
struct Object {
// 0: shpere
// 1: plane
// 2: cube
type: u32,
position: vec3<f32>,
rotation: vec4<f32>,
scale: vec3<f32>
}
@compute @workgroup_size(8, 8, 1)
fn init(@builtin(global_invocation_id) invocation_id: vec3<u32>, @builtin(num_workgroups) num_workgroups: vec3<u32>) {
let location = vec2<i32>(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<u32>) {
let size = textureDimensions(output);
let loc = vec2<f32>(f32(invocation_id.x), f32(invocation_id.y)) / vec2<f32>(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<f32>(final_color, 1.0);
if hit_data.distance == 100.0 {
color = config.sky_color;
// color = textureSampleLevel(skybox_texture, skybox_sampler, vec3<f32>(0.0, 0.0, 1.0), 0, 0u);
}
let location = vec2<i32>(i32(invocation_id.x), i32(invocation_id.y));
textureStore(output, location, color);
}
fn debug_matrix(uv: vec2<f32>) -> vec4<f32>{
let ndc = uv * 2.0f - 1.0f;
var color = vec3<f32>(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<f32>(color, 1.0);
}
fn debug(uv: vec2<f32>) -> vec4<f32>{
let ndc = uv * 2.0f - 1.0f;
let near_clip = vec4<f32>(ndc, 0.0, 1.0);
let far_clip = vec4<f32>(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<f32>(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<f32>(color, 1.0);
}
fn debugColor(v: f32) -> vec3<f32>{
return vec3<f32>(v * 0.5 + 0.5);
}
struct Ray {
origin: vec3<f32>,
direction: vec3<f32>,
energy: vec3<f32>,
}
fn createRay(origin: vec3<f32>, direction: vec3<f32>) -> Ray
{
var ray: Ray;
ray.origin = origin;
ray.direction = direction;
ray.energy = vec3<f32>(1.0f, 1.0f, 1.0f);
return ray;
}
fn createCameraRay(ndc: vec2<f32>) -> Ray {
let uv = vec2<f32>(ndc.x, -ndc.y);
// clip points at near and far
let near_clip = vec4<f32>(uv, 0.0, 1.0);
let far_clip = vec4<f32>(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<f32>,
normal: vec3<f32>,
color: vec3<f32>,
direction: vec3<f32>,
steps: i32,
};
fn distance_field(p: vec3<f32>) -> f32 {
// Simple sphere centered at (0, 0, 0) with radius 1.0
let sphere_center = vec3<f32>(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<f32>(0.0), vec3<f32>(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<f32>(0.0), vec3<f32>(0.0), vec3<f32>(0.0, 0.0, 0.0), ray.direction, max_steps); // Return black for background
}

182
assets/trace.wgsl
View File

@@ -1,43 +1,14 @@
#import bevy_sprite::mesh2d_vertex_output::VertexOutput
@group(0) @binding(0) var input: texture_storage_2d<rgba32float, read>;
@group(2) @binding(0) var<uniform> sky_color: vec4<f32>;
@group(2) @binding(1) var<uniform> view: TracerView;
@group(2) @binding(2) var skybox_texture: texture_2d<f32>;
@group(2) @binding(3) var skybox_sampler: sampler;
@group(0) @binding(1) var output: texture_storage_2d<rgba32float, write>;
@fragment
fn fragment(mesh: VertexOutput) -> @location(0) vec4<f32> {
@group(0) @binding(2) var<uniform> config: TracerUniforms;
@group(0) @binding(3) var skybox_texture: texture_cube<f32>;
@group(0) @binding(4) var skybox_sampler: sampler;
struct View {
view_proj: mat4x4<f32>,
view: mat4x4<f32>,
projection: mat4x4<f32>,
inv_view_proj: mat4x4<f32>,
inv_view: mat4x4<f32>, // equiv to Unity's _CameraToWorld
inv_projection: mat4x4<f32>, // equic to Unity's _CameraInverseProjection
};
struct TracerUniforms {
sky_color: vec4<f32>,
world_from_clip: mat4x4<f32>,
world_position: vec3<f32>,
}
@compute @workgroup_size(8, 8, 1)
fn init(@builtin(global_invocation_id) invocation_id: vec3<u32>, @builtin(num_workgroups) num_workgroups: vec3<u32>) {
let location = vec2<i32>(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<u32>) {
let size = textureDimensions(output);
let loc = vec2<f32>(f32(invocation_id.x), f32(invocation_id.y)) / vec2<f32>(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<u32>) {
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<f32>(final_color, 1.0);
if hit_data.distance == 100.0 {
color = config.sky_color;
// color = textureSampleLevel(skybox_texture, skybox_sampler, vec3<f32>(0.0, 0.0, 1.0), 0, 0u);
}
let location = vec2<i32>(i32(invocation_id.x), i32(invocation_id.y));
textureStore(output, location, color);
}
fn debug_matrix(uv: vec2<f32>) -> vec4<f32>{
let ndc = uv * 2.0f - 1.0f;
var color = vec3<f32>(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<f32>(color, 1.0);
}
fn debug(uv: vec2<f32>) -> vec4<f32>{
let ndc = uv * 2.0f - 1.0f;
let near_clip = vec4<f32>(ndc, 0.0, 1.0);
let far_clip = vec4<f32>(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<f32>(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<f32>(color, 1.0);
}
fn debugColor(v: f32) -> vec3<f32>{
return vec3<f32>(v * 0.5 + 0.5);
}
struct Ray {
origin: vec3<f32>,
direction: vec3<f32>,
energy: vec3<f32>,
var color = vec4<f32>(final_color, 1.0);
if hit_data.distance == 100.0 {
color = sky_color;
// color = textureSampleLevel(skybox_texture, skybox_sampler, vec3<f32>(0.0, 0.0, 1.0), 0, 0u);
}
return color;
}
fn createRay(origin: vec3<f32>, direction: vec3<f32>) -> Ray
@@ -150,8 +47,8 @@ fn createCameraRay(ndc: vec2<f32>) -> Ray {
let far_clip = vec4<f32>(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<f32>) -> 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<f32>,
normal: vec3<f32>,
color: vec3<f32>,
direction: vec3<f32>,
steps: i32,
};
fn distance_field(p: vec3<f32>) -> f32 {
// Simple sphere centered at (0, 0, 0) with radius 1.0
let sphere_center = vec3<f32>(0.0, 0.0, 0.0);
@@ -214,3 +102,33 @@ fn trace(ray: Ray) -> Hit {
// No hit found
return Hit(max_distance, vec3<f32>(0.0), vec3<f32>(0.0), vec3<f32>(0.0, 0.0, 0.0), ray.direction, max_steps); // Return black for background
}
struct TracerView {
world_from_clip: mat4x4<f32>,
world_position: vec3<f32>,
}
struct Object {
// 0: shpere
// 1: plane
// 2: cube
object_type: u32,
position: vec3<f32>,
rotation: vec4<f32>,
scale: vec3<f32>
}
struct Ray {
origin: vec3<f32>,
direction: vec3<f32>,
energy: vec3<f32>,
}
struct Hit {
distance: f32,
hit_pos: vec3<f32>,
normal: vec3<f32>,
color: vec3<f32>,
direction: vec3<f32>,
steps: i32,
};

92
src/app.rs
View File

@@ -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<Image>);
impl Plugin for Blackhole {
fn build(&self, app: &mut App) {
app.register_type::<TracerRenderTextures>();
app.init_state::<AssetLoad>();
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<Assets<Image>>,
window: Single<&Window, With<PrimaryWindow>>,
asset_server: Res<AssetServer>,
mut load_state: ResMut<NextState<AssetLoad>>,
mut materials: ResMut<Assets<TracerMaterial>>,
mut meshes: ResMut<Assets<Mesh>>,
) {
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<TracerRenderTextures>,
skybox: Res<SkyboxAsset>,
mut image_assets: ResMut<Assets<Image>>,
) {
fn prepare_skybox(skybox: Res<SkyboxAsset>, mut image_assets: ResMut<Assets<Image>>) {
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<NextState<AssetLoad>>) {

5
src/components/rt.rs
View File

@@ -1,4 +1,7 @@
use bevy::prelude::*;
#[derive(Component)]
#[derive(Component, Debug)]
pub struct RTCamera;
#[derive(Component, Debug)]
pub struct RTDisplay;

4
src/render/mod.rs
View File

@@ -1,2 +1,2 @@
pub mod node;
pub mod pipeline;
pub mod tracer;
pub mod tracer_material;

112
src/render/node.rs
View File

@@ -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::<TracerPipeline>();
let pipeline_cache = world.resource::<PipelineCache>();
// 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::<TracerRenderTextures>();
// let asset_server = world.resource::<AssetServer>();
// 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::<TracerImageBindGroups>();
if bind_groups_res.is_none() {
return Ok(());
}
let bind_groups = &bind_groups_res.unwrap().0;
let pipeline_cache = world.resource::<PipelineCache>();
let pipeline = world.resource::<TracerPipeline>();
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(())
}
}

236
src/render/pipeline.rs
View File

@@ -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<Image>,
pub secondary: Handle<Image>,
pub skybox: Handle<Image>,
}
#[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::<TracerRenderTextures>::default(),
ExtractResourcePlugin::<TracerUniforms>::default(),
));
app.init_resource::<TracerUniforms>()
.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::<RenderGraph>();
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::<TracerPipeline>();
}
}
fn switch_textures(images: Res<TracerRenderTextures>, 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::<RenderDevice>();
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::<TracerUniforms>(false),
texture_cube(TextureSampleType::Float { filterable: true }),
sampler(SamplerBindingType::Filtering),
),
),
);
let shader = world.load_asset(SHADER_ASSET_PATH);
let pipeline_cache = world.resource::<PipelineCache>();
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<RenderDevice>,
asset_server: Res<AssetServer>,
pipeline_cache: Res<PipelineCache>,
) {
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::<TracerUniforms>(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<TracerUniforms>,
rt_camera: Single<(&GlobalTransform, &Camera), With<RTCamera>>,
) {
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<TracerPipeline>,
gpu_images: Res<RenderAssets<GpuImage>>,
tracer_images: Res<TracerRenderTextures>,
tracer_uniforms: Res<TracerUniforms>,
render_device: Res<RenderDevice>,
queue: Res<RenderQueue>,
) {
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]));
}

35
src/render/tracer.rs Normal file
View File

@@ -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::<TracerMaterial>::default());
app.add_systems(First, update_tracer_uniforms.run_if(in_state(AssetLoad::Ready)));
}
}
fn update_tracer_uniforms(
rt_camera: Single<(&GlobalTransform, &Camera), With<RTCamera>>,
display: Single<&MeshMaterial2d<TracerMaterial>, With<RTDisplay>>,
mut materials: ResMut<Assets<TracerMaterial>>,
) {
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(),
};
}

35
src/render/tracer_material.rs Normal file
View File

@@ -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<Handle<Image>>,
}
#[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
}
}