From 4e08d2b8fd5811dee2220f19a0d3c33a60e0d938 Mon Sep 17 00:00:00 2001 From: Matthew Gordon Date: Mon, 23 Dec 2024 11:39:56 -0400 Subject: [PATCH] Break up shader into multiple files --- Cargo.lock | 22 +++-- Cargo.toml | 1 + rust-toolchain.toml | 2 + src/app/dem_renderer/dem_renderer.wgsl | 93 ++------------------- src/app/dem_renderer/mod.rs | 7 +- src/app/dem_renderer/pack_unpack.wgsl | 3 + src/app/dem_renderer/ray_intersection.wgsl | 96 ++++++++++++++++++++++ 7 files changed, 129 insertions(+), 95 deletions(-) create mode 100644 rust-toolchain.toml create mode 100644 src/app/dem_renderer/pack_unpack.wgsl create mode 100644 src/app/dem_renderer/ray_intersection.wgsl diff --git a/Cargo.lock b/Cargo.lock index b968e55..64dc1cc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "ab_glyph" @@ -1450,9 +1450,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.89" +version = "1.0.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f139b0662de085916d1fb67d2b4169d1addddda1919e696f3252b740b629986e" +checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" dependencies = [ "unicode-ident", ] @@ -1501,6 +1501,7 @@ dependencies = [ "wasm-bindgen-test", "web-sys", "wgpu", + "wgsl-shader-assembler", "winit", ] @@ -1804,9 +1805,9 @@ checksum = "6637bab7722d379c8b41ba849228d680cc12d0a45ba1fa2b48f2a30577a06731" [[package]] name = "syn" -version = "2.0.87" +version = "2.0.90" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25aa4ce346d03a6dcd68dd8b4010bcb74e54e62c90c573f394c46eae99aba32d" +checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31" dependencies = [ "proc-macro2", "quote", @@ -2331,6 +2332,17 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wgsl-shader-assembler" +version = "0.1.0" +source = "git+https://git.gordon.earth/matthew/wgsl-shader-assembler.git#2ce14e25d8f1685bd3486257cba4025587fd0fb9" +dependencies = [ + "naga", + "quote", + "regex", + "syn", +] + [[package]] name = "winapi-util" version = "0.1.9" diff --git a/Cargo.toml b/Cargo.toml index 42e6299..b44a9ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ bytemuck = { version = "1.19.0", features = ["derive"] } glam = "0.29.2" log = "0.4.22" tiff = "0.9.1" +wgsl-shader-assembler = { git = "https://git.gordon.earth/matthew/wgsl-shader-assembler.git" } [target.'cfg(target_arch = "x86_64")'.dependencies] winit = { version = "0.30.3", features = ["rwh_06"] } diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000..271800c --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,2 @@ +[toolchain] +channel = "nightly" \ No newline at end of file diff --git a/src/app/dem_renderer/dem_renderer.wgsl b/src/app/dem_renderer/dem_renderer.wgsl index 7560e05..6faf98b 100644 --- a/src/app/dem_renderer/dem_renderer.wgsl +++ b/src/app/dem_renderer/dem_renderer.wgsl @@ -1,7 +1,7 @@ struct VertexOutput { @location(0) world_space_position: vec4, @builtin(position) position: vec4, -}; +} struct Uniforms { transform: mat4x4, @@ -20,90 +20,7 @@ var dem_texture: texture_2d; @group(0) @binding(2) var dembvh_texture: texture_2d; -struct BoundingNode { - index: vec2, - level: u32 -} - -struct AABB { - min_corner: vec3, - max_corner: vec3 -} - -struct Ray { origin: vec3, direction: vec3 }; - -fn intersect_ray_with_aabb(ray: Ray, aabb: AABB) -> bool { - let t1 = (aabb.min_corner - ray.origin) / ray.direction; - let t2 = (aabb.max_corner - ray.origin) / ray.direction; - let t_mins = min(t1, t2); - let t_min = min(t_mins.x, min(t_mins.y, t_mins.z)); - let t_maxs = max(t1, t2); - let t_max = max(t_maxs.x, max(t_maxs.y, t_maxs.z)); - return t_min <= t_max; -} - -fn scale_u16(value: u32, range: vec2) -> f32 { - return range.x + (f32(value - 1) / 65534.0) * (range.y - range.x); -} - -fn get_xy_min_for_node(node: BoundingNode) -> vec2 { - return uniforms.dem_min_corner.xy + uniforms.dem_cell_size * vec2(node.index) * exp2(f32(node.level)); -} - -fn get_xy_max_for_node(node: BoundingNode) -> vec2 { - return uniforms.dem_min_corner.xy + uniforms.dem_cell_size * vec2(node.index + 1) * exp2(f32(node.level)); -} - -struct NodeStack { - stack: array, - count: u32 -} - -fn push_node_stack(stack: ptr, node: BoundingNode) { - (*stack).stack[(*stack).count] = node; - (*stack).count++; -} - -fn pop_node_stack(stack: ptr) -> BoundingNode { - (*stack).count--; - return (*stack).stack[(*stack).count]; -} - -fn intersect_ray_with_node(ray: Ray, root_node: BoundingNode) -> bool { - var node_stack: NodeStack; - node_stack.count = 0u; - push_node_stack(&node_stack, root_node); - while node_stack.count > 0 { - let node = pop_node_stack(&node_stack); - let minmax_z = textureLoad(dembvh_texture, node.index, i32(node.level)).rg; - if minmax_z.r == 0 { - return false; - } - let min_z = scale_u16(minmax_z.r, uniforms.dem_z_range); - let max_z = scale_u16(minmax_z.g, uniforms.dem_z_range); - var aabb: AABB ; - aabb.min_corner = vec3(get_xy_min_for_node(node), min_z); - aabb.max_corner = vec3(get_xy_max_for_node(node), min_z); - if intersect_ray_with_aabb(ray, aabb) { - if node.level == 0 { - return true; - } else { - let next_index = node.index * 2; - var next_node: BoundingNode; - next_node.index = next_index; - next_node.level = node.level - 1; - push_node_stack(&node_stack, next_node); - next_node.index = next_index + vec2(1, 0); - push_node_stack(&node_stack, next_node); - next_node.index = next_index + vec2(0, 1); - push_node_stack(&node_stack, next_node); - next_node.index = next_index + vec2(1, 1); - push_node_stack(&node_stack, next_node); - } - } - } - return false; -} +//#include ray_intersection.wgsl @vertex fn vs_main( @@ -125,7 +42,11 @@ fn fs_solid(vertex: VertexOutput) -> @location(0) vec4 { root_node.index = vec2(0); root_node.level = textureNumLevels(dembvh_texture) - 1; - if intersect_ray_with_node(ray, root_node) { + if intersect_ray_with_node(dembvh_texture, + uniforms.dem_min_corner, + uniforms.dem_cell_size, + ray, + root_node) { return vec4(1.0); } else { discard; diff --git a/src/app/dem_renderer/mod.rs b/src/app/dem_renderer/mod.rs index 7cdc364..0eee1f5 100644 --- a/src/app/dem_renderer/mod.rs +++ b/src/app/dem_renderer/mod.rs @@ -1,3 +1,5 @@ +use wgsl_shader_assembler::wgsl_module; + use super::raster::{Dem, DemBvh}; use { bytemuck::{Pod, Zeroable}, @@ -250,10 +252,7 @@ impl DemRenderer { }], }]; - let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: None, - source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("dem_renderer.wgsl"))), - }); + let shader = device.create_shader_module(wgsl_module!("dem_renderer.wgsl")); let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor { label: Some("DemRendererPipeline"), diff --git a/src/app/dem_renderer/pack_unpack.wgsl b/src/app/dem_renderer/pack_unpack.wgsl new file mode 100644 index 0000000..16bd91a --- /dev/null +++ b/src/app/dem_renderer/pack_unpack.wgsl @@ -0,0 +1,3 @@ +fn scale_u16(value: u32, range: vec2) -> f32 { + return range.x + (f32(value - 1) / 65534.0) * (range.y - range.x); +} diff --git a/src/app/dem_renderer/ray_intersection.wgsl b/src/app/dem_renderer/ray_intersection.wgsl new file mode 100644 index 0000000..397d68f --- /dev/null +++ b/src/app/dem_renderer/ray_intersection.wgsl @@ -0,0 +1,96 @@ +//#include pack_unpack.wgsl + +struct BoundingNode { + index: vec2, + level: u32 +} + +struct AABB { + min_corner: vec3, + max_corner: vec3 +} + +struct Ray { origin: vec3, direction: vec3 }; + +fn intersect_ray_with_aabb(ray: Ray, aabb: AABB) -> bool { + let t1 = (aabb.min_corner - ray.origin) / ray.direction; + let t2 = (aabb.max_corner - ray.origin) / ray.direction; + let t_mins = min(t1, t2); + let t_min = min(t_mins.x, min(t_mins.y, t_mins.z)); + let t_maxs = max(t1, t2); + let t_max = max(t_maxs.x, max(t_maxs.y, t_maxs.z)); + return t_min <= t_max; +} + +fn get_xy_min_for_node(dem_min_corner: vec2, + dem_cell_size: vec2, + node: BoundingNode) -> vec2 { + return dem_min_corner.xy + dem_cell_size * vec2(node.index) * exp2(f32(node.level)); +} + +fn get_xy_max_for_node(dem_min_corner: vec2, + dem_cell_size: vec2, + node: BoundingNode) -> vec2 { + return dem_min_corner.xy + dem_cell_size * vec2(node.index + 1) * exp2(f32(node.level)); +} + +struct NodeStack { + stack: array, + count: u32 +} + +fn push_node_stack(stack: ptr, node: BoundingNode) { + (*stack).stack[(*stack).count] = node; + (*stack).count++; +} + +fn pop_node_stack(stack: ptr) -> BoundingNode { + (*stack).count--; + return (*stack).stack[(*stack).count]; +} + +fn intersect_ray_with_node(tree_texture: texture_2d, + dem_min_corner: vec2, + dem_cell_size: vec2, + ray: Ray, + root_node: BoundingNode) -> bool { + var node_stack: NodeStack; + node_stack.count = 0u; + push_node_stack(&node_stack, root_node); + while node_stack.count > 0 { + let node = pop_node_stack(&node_stack); + let minmax_z = textureLoad(tree_texture, node.index, i32(node.level)).rg; + if minmax_z.r == 0 { + return false; + } + let min_z = scale_u16(minmax_z.r, uniforms.dem_z_range); + let max_z = scale_u16(minmax_z.g, uniforms.dem_z_range); + var aabb: AABB ; + aabb.min_corner = vec3(get_xy_min_for_node(dem_min_corner, + dem_cell_size, + node), + min_z); + aabb.max_corner = vec3(get_xy_max_for_node(dem_min_corner, + dem_cell_size, + node), + max_z); + if intersect_ray_with_aabb(ray, aabb) { + if node.level == 0 { + return true; + } else { + let next_index = node.index * 2; + var next_node: BoundingNode; + next_node.index = next_index; + next_node.level = node.level - 1; + push_node_stack(&node_stack, next_node); + next_node.index = next_index + vec2(1, 0); + push_node_stack(&node_stack, next_node); + next_node.index = next_index + vec2(0, 1); + push_node_stack(&node_stack, next_node); + next_node.index = next_index + vec2(1, 1); + push_node_stack(&node_stack, next_node); + } + } + } + return false; +}