From 7d3afc9002fa0f792c7e2b90fd85b6bd5bf2bdff Mon Sep 17 00:00:00 2001 From: Amatsugu Date: Mon, 1 Sep 2025 13:16:45 -0400 Subject: [PATCH] translate ray tracer code --- assets/trace.wgsl | 233 +++++++++++++++++++++++++++++++++-------- src/render/pipeline.rs | 6 ++ 2 files changed, 194 insertions(+), 45 deletions(-) diff --git a/assets/trace.wgsl b/assets/trace.wgsl index 4bfc9e1..d4dc470 100644 --- a/assets/trace.wgsl +++ b/assets/trace.wgsl @@ -5,70 +5,213 @@ @group(0) @binding(2) var config: TracerUniforms; +@group(0) @binding(3) var view: ViewUniform; + +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 ViewUniform { + clip_from_world: mat4x4, + unjittered_clip_from_world: mat4x4, + world_from_clip: mat4x4, + world_from_view: mat4x4, + view_from_world: mat4x4, + clip_from_view: mat4x4, + view_from_clip: mat4x4, + world_position: vec3, + exposure: f32, + viewport: vec4, + frustum: array, 6>, + color_grading: vec4, // simplified for example + mip_bias: f32, + frame_count: u32, +}; + struct TracerUniforms { sky_color: vec4, } -fn hash(value: u32) -> u32 { - var state = value; - state = state ^ 2747636419u; - state = state * 2654435769u; - state = state ^ (state >> 16u); - state = state * 2654435769u; - state = state ^ (state >> 16u); - state = state * 2654435769u; - return state; -} - -fn randomFloat(value: u32) -> f32 { - return f32(hash(value)) / 4294967295.0; -} - @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 randomNumber = randomFloat((invocation_id.y << 16u) | invocation_id.x); - let alive = randomNumber > 0.9; - // Use alpha channel to keep track of cell's state - let color = vec4(config.sky_color.rgb, f32(alive)); + let color = vec4(0.0); textureStore(output, location, color); } -fn is_alive(location: vec2, offset_x: i32, offset_y: i32) -> i32 { - let value: vec4 = textureLoad(input, location + vec2(offset_x, offset_y)); - return i32(value.a); -} -fn count_alive(location: vec2) -> i32 { - return is_alive(location, -1, -1) + - is_alive(location, -1, 0) + - is_alive(location, -1, 1) + - is_alive(location, 0, -1) + - is_alive(location, 0, 1) + - is_alive(location, 1, -1) + - is_alive(location, 1, 0) + - is_alive(location, 1, 1); -} @compute @workgroup_size(8, 8, 1) fn update(@builtin(global_invocation_id) invocation_id: vec3) { -let location = vec2(i32(invocation_id.x), i32(invocation_id.y)); + let size = textureDimensions(output); - let n_alive = count_alive(location); + let loc = vec2(f32(invocation_id.x), f32(invocation_id.y)) / vec2(size.xy); + let ndc = loc * 2.0f - 1.0f; - var alive: bool; + var ray = createCameraRay(ndc); - if (n_alive == 3) { - alive = true; - } else if (n_alive == 2) { - let currently_alive = is_alive(location, 0, 0); - alive = bool(currently_alive); - } else { - alive = false; - } - let color = vec4(config.sky_color.rgb , f32(alive)); + var result = vec3(0.0f); + var hit = trace(ray); + result += ray.energy * shade(&ray, hit); + + let color = vec4(result , 1.0); + let location = vec2(i32(invocation_id.x), i32(invocation_id.y)); textureStore(output, location, color); } + + +struct Ray { + origin: vec3, + direction: vec3, + energy: vec3, +} + + +struct RayHit { + distance: f32, + position: vec3, + normal: vec3, + albedo: vec3, + specular: vec3 +} + +struct Sphere +{ + position: vec3, + radius: f32, + albedo: vec3, + specular: vec3 +} + +fn createRayHit() -> RayHit { + var hit: RayHit; + hit.position = vec3(0.0, 0.0, 0.0); + hit.distance = -1.0f; // A negative number to represent infinity + hit.normal = vec3(0.0, 0.0, 0.0); + hit.albedo = vec3(0.0, 0.0, 0.0); + hit.specular = vec3(0.0, 0.0, 0.0); + return hit; +} + + +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 origin = (view.inv_view * vec4(0.0, 0.0, 0.0, 1.0)).xyz; + + // let direction_view = (view.inv_projection * vec4(uv, 0.0, 1.0)).xyz; + // let direction = (view.inv_view * vec4(direction_view, 0.0)).xyz; + + let origin = view.world_position; + let target_point = view.world_from_clip * vec4(ndc, 0.0f, 1.0f); + let direction_point = target_point.xyz / target_point.w; + let direction = normalize(direction_point - origin); + + return createRay(origin, direction); +} + +fn createSphere(position: vec3, radius: f32) -> Sphere +{ + var s: Sphere; + s.position = position; + s.radius = radius; + s.albedo = vec3(0.8f, 0.8f, 0.8f); + s.specular = vec3(0.6f, 0.6f, 0.6f); + return s; +} + +fn intersectSphere(ray: Ray, bestHit: ptr, sphereIndex: u32) +{ + //Sphere sphere = _Spheres[sphereIndex]; + var sphere = createSphere(vec3(0.0), 1.0f); + var d = ray.origin - sphere.position; + var p1 = -dot(ray.direction, d); + var p2sqr = p1 * p1 - dot(d, d) + sphere.radius * sphere.radius; + if p2sqr < 0 { + return; + } + var p2 = sqrt(p2sqr); + // var t = p1 - p2 > 0 ? p1 - p2 : p1 + p2; + var t = 0f; + if p1 - p2 > 0 { + t = p1 - p2; + } else { + t = p1 + p2; + } + if t > 0 && t < (*bestHit).distance + { + (*bestHit).position = ray.origin + t * ray.direction; + (*bestHit).normal = normalize((*bestHit).position - sphere.position); + (*bestHit).albedo = sphere.albedo; + (*bestHit).specular = sphere.specular; + (*bestHit).distance = t; + } +} + +fn intersectGroundPlane(ray: Ray, bestHit: ptr) +{ + var t = -ray.origin.y / ray.direction.y; + if t > 0 && t < (*bestHit).distance + { + (*bestHit).distance = t; + (*bestHit).position = ray.origin + t * ray.direction; + (*bestHit).normal = vec3(0.0f, 1.0f, 0.0f); + (*bestHit).albedo = vec3(0.8f); + (*bestHit).specular = vec3(0.3f); + } +} + +fn trace(ray: Ray) -> RayHit +{ + var bestHit = createRayHit(); + intersectGroundPlane(ray, &bestHit); + intersectSphere(ray, &bestHit, 0); + return bestHit; +} + + +fn shade(ray: ptr, hit: RayHit) -> vec3 +{ + if hit.distance > -1.0f + { + (*ray).origin = hit.position + hit.normal * 0.001f; + (*ray).direction = reflect((*ray).direction, hit.normal); + (*ray).energy *= hit.specular; + + //Shadows + // var shadow = false; + // Ray shadowRay = createRay(hit.position + hit.normal * 0.001f, -1 * _DirectionalLight.xyz); + // RayHit shadowHit = Trace(shadowRay); + // if (shadowHit.distance != 1.#INF) + // { + // return float3(0.0f, 0.0f, 0.0f); + // } + + // return saturate(dot(hit.normal, _DirectionalLight.xyz) * -1) * _DirectionalLight.w * hit.albedo; + return hit.albedo; + } + else + { + (*ray).energy = vec3(0.0f); + return vec3(0.1f); + // var theta = acos(ray.direction.y) / -PI; + // var phi = atan2(ray.direction.x, -ray.direction.z) / -PI * .5f; + // return _SkyboxTexture.SampleLevel(sampler_SkyboxTexture, float2(phi, theta), 0).xyz; + } +} diff --git a/src/render/pipeline.rs b/src/render/pipeline.rs index 218bc2e..ae143ed 100644 --- a/src/render/pipeline.rs +++ b/src/render/pipeline.rs @@ -15,6 +15,7 @@ use bevy::{ }, renderer::{RenderDevice, RenderQueue}, texture::GpuImage, + view::{ViewUniform, ViewUniforms}, }, }; @@ -85,6 +86,7 @@ impl FromWorld for TracerPipeline { texture_storage_2d(TextureFormat::Rgba32Float, StorageTextureAccess::ReadOnly), texture_storage_2d(TextureFormat::Rgba32Float, StorageTextureAccess::WriteOnly), uniform_buffer::(false), + uniform_buffer::(false), ), ), ); @@ -133,6 +135,7 @@ fn init_pipeline( texture_storage_2d(TextureFormat::Rgba32Float, StorageTextureAccess::ReadOnly), texture_storage_2d(TextureFormat::Rgba32Float, StorageTextureAccess::WriteOnly), uniform_buffer::(false), + uniform_buffer::(false), ), ), ); @@ -175,10 +178,13 @@ fn prepare_bind_groups( tracer_uniforms: Res, render_device: Res, queue: Res, + view_uniforms: Res, ) { let view_a = gpu_images.get(&tracer_images.0).unwrap(); let view_b = gpu_images.get(&tracer_images.1).unwrap(); + //Todo: Insert View Uniforms + // 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());