Sky shader Improvements. (#245)

* Added some constants to the shader

* did something for a certain video i forgot to revert

* removed unused function

* cleanup

* cpu runner now works

* explicit naming

* removed paranthesis

* clippy

* CLIPPY

examples/runners/ash/src/main.rs

@@ -32,6 +32,18 @@ macro_rules! offset_of {
         }
     }};
 }
+
+#[allow(dead_code)]
+struct ShaderConstants {
+    width: u32,
+    height: u32,
+    time: f32,
+}
+
+unsafe fn any_as_u8_slice<T: Sized>(p: &T) -> &[u8] {
+    ::std::slice::from_raw_parts((p as *const T) as *const u8, ::std::mem::size_of::<T>())
+}
+
 /// Helper function for submitting command buffers. Immediately waits for the fence before the command buffer
 /// is executed. That way we can delay the waiting for the fences by 1 frame which is good for performance.
 /// Make sure to create the fence in a signaled state on the first use.
@@ -632,7 +644,14 @@ fn main() {
             .create_shader_module(&shader_info, None)
             .expect("Shader module error");
 
-        let layout_create_info = vk::PipelineLayoutCreateInfo::default();
+        let push_constant_range = vk::PushConstantRange::builder()
+            .offset(0)
+            .size(std::mem::size_of::<ShaderConstants>() as u32)
+            .stage_flags(vk::ShaderStageFlags::all())
+            .build();
+        let layout_create_info = vk::PipelineLayoutCreateInfo::builder()
+            .push_constant_ranges(&[push_constant_range])
+            .build();
 
         let pipeline_layout = base
             .device
@@ -750,6 +769,8 @@ fn main() {
 
         let graphic_pipeline = graphics_pipelines[0];
 
+        let start = std::time::Instant::now();
+
         base.render_loop(events_loop, move |base| {
             let (present_index, _) = base
                 .swapchain_loader
@@ -796,6 +817,19 @@ fn main() {
                     );
                     device.cmd_set_viewport(draw_command_buffer, 0, &viewports);
                     device.cmd_set_scissor(draw_command_buffer, 0, &scissors);
+                    let push_constants = ShaderConstants {
+                        width: 1280, // ash runner currently does not support resizing.
+                        height: 720,
+                        time: start.elapsed().as_secs_f32(),
+                    };
+
+                    device.cmd_push_constants(
+                        draw_command_buffer,
+                        pipeline_layout,
+                        ash::vk::ShaderStageFlags::all(),
+                        0,
+                        any_as_u8_slice(&push_constants),
+                    );
                     device.cmd_draw(draw_command_buffer, 3, 1, 0, 0);
                     // Or draw without the index buffer
                     // device.cmd_draw(draw_command_buffer, 3, 1, 0, 0);

examples/runners/cpu/src/main.rs

@@ -1,6 +1,6 @@
 use minifb::{Key, Window, WindowOptions};
 use rayon::prelude::*;
-use spirv_std::glam::{vec2, Vec4};
+use spirv_std::glam::{vec2, Vec2, Vec4};
 use std::time::Instant;
 
 use sky_shader as shader_module;
@@ -35,6 +35,12 @@ fn main() {
     )
     .expect("Window creation failed");
 
+    let push_constants = shader_module::ShaderConstants {
+        width: WIDTH as u32,
+        height: HEIGHT as u32,
+        time: 0f32,
+    };
+
     // Limit to max ~60 fps update rate
     window.limit_update_rate(Some(std::time::Duration::from_micros(16600)));
 
@@ -48,8 +54,12 @@ fn main() {
                 -((i / WIDTH) as f32 / HEIGHT as f32 * 2.0 - 1.0),
             );
 
+            let frag_coord = (vec2(screen_pos.x(), -screen_pos.y()) + Vec2::one())
+                / Vec2::splat(2.0)
+                * vec2(WIDTH as f32, HEIGHT as f32);
+
             // evaluate the fragment shader for the specific pixel
-            let color = shader_module::fs(screen_pos);
+            let color = shader_module::fs(&push_constants, frag_coord);
 
             color_u32_from_vec4(color)
         })

examples/runners/wgpu/src/graphics.rs

@@ -1,10 +1,17 @@
-use super::{shader_module, Options};
+use super::{shader_module, Options, ShaderConstants};
 use winit::{
     event::{Event, KeyboardInput, VirtualKeyCode, WindowEvent},
     event_loop::{ControlFlow, EventLoop},
     window::Window,
 };
 
+unsafe fn any_as_u32_slice<T: Sized>(p: &T) -> &[u32] {
+    ::std::slice::from_raw_parts(
+        (p as *const T) as *const u32,
+        ::std::mem::size_of::<T>() / 4,
+    )
+}
+
 async fn run(
     options: &Options,
     event_loop: EventLoop<()>,
@@ -31,12 +38,18 @@ async fn run(
         .await
         .expect("Failed to find an appropriate adapter");
 
+    let features = wgpu::Features::PUSH_CONSTANTS;
+    let limits = wgpu::Limits {
+        max_push_constant_size: 256,
+        ..Default::default()
+    };
+
     // Create the logical device and command queue
     let (device, queue) = adapter
         .request_device(
             &wgpu::DeviceDescriptor {
-                features: wgpu::Features::empty(),
-                limits: wgpu::Limits::default(),
+                features,
+                limits,
                 shader_validation: true,
             },
             None,
@@ -50,7 +63,10 @@ async fn run(
     let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
         label: None,
         bind_group_layouts: &[],
-        push_constant_ranges: &[],
+        push_constant_ranges: &[wgpu::PushConstantRange {
+            stages: wgpu::ShaderStage::all(),
+            range: 0..std::mem::size_of::<ShaderConstants>() as u32,
+        }],
     });
 
     let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
@@ -90,6 +106,8 @@ async fn run(
         .as_ref()
         .map(|surface| device.create_swap_chain(&surface, &sc_desc));
 
+    let start = std::time::Instant::now();
+
     event_loop.run(move |event, _, control_flow| {
         // Have the closure take ownership of the resources.
         // `event_loop.run` never returns, therefore we must do this to ensure
@@ -98,6 +116,9 @@ async fn run(
 
         *control_flow = ControlFlow::Wait;
         match event {
+            Event::MainEventsCleared => {
+                window.request_redraw();
+            }
             Event::Resumed => {
                 let s = unsafe { instance.create_surface(&window) };
                 swap_chain = Some(device.create_swap_chain(&s, &sc_desc));
@@ -138,7 +159,15 @@ async fn run(
                             }],
                             depth_stencil_attachment: None,
                         });
+                        let push_constants = ShaderConstants {
+                            width: window.inner_size().width,
+                            height: window.inner_size().height,
+                            time: start.elapsed().as_secs_f32(),
+                        };
                         rpass.set_pipeline(&render_pipeline);
+                        rpass.set_push_constants(wgpu::ShaderStage::all(), 0, unsafe {
+                            any_as_u32_slice(&push_constants)
+                        });
                         rpass.draw(0..3, 0..1);
                     }
 

examples/runners/wgpu/src/lib.rs

@@ -11,6 +11,13 @@ pub enum RustGPUShader {
     Compute,
 }
 
+#[allow(dead_code)]
+pub struct ShaderConstants {
+    width: u32,
+    height: u32,
+    time: f32,
+}
+
 fn shader_module(shader: RustGPUShader) -> wgpu::ShaderModuleSource<'static> {
     match shader {
         RustGPUShader::Simplest => wgpu::include_spirv!(env!("simplest_shader.spv")),

examples/shaders/sky-shader/src/lib.rs

@@ -6,8 +6,8 @@
 #![register_attr(spirv)]
 
 use core::f32::consts::PI;
-use spirv_std::glam::{const_vec3, Mat4, Vec2, Vec3, Vec4};
-use spirv_std::{Input, MathExt, Output};
+use spirv_std::glam::{const_vec3, Vec2, Vec3, Vec4};
+use spirv_std::{Input, MathExt, Output, PushConstant};
 
 const DEPOLARIZATION_FACTOR: f32 = 0.035;
 const MIE_COEFFICIENT: f32 = 0.005;
@@ -25,6 +25,13 @@ const SUN_INTENSITY_FACTOR: f32 = 1000.0;
 const SUN_INTENSITY_FALLOFF_STEEPNESS: f32 = 1.5;
 const TURBIDITY: f32 = 2.0;
 
+#[derive(Copy, Clone)]
+pub struct ShaderConstants {
+    pub width: u32,
+    pub height: u32,
+    pub time: f32,
+}
+
 // TODO: add this to glam? Rust std has it on f32/f64
 fn pow(v: Vec3, power: f32) -> Vec3 {
     Vec3::new(v.x().pow(power), v.y().pow(power), v.z().pow(power))
@@ -149,19 +156,22 @@ fn sky(dir: Vec3, sun_position: Vec3) -> Vec3 {
     lin + l0
 }
 
-pub fn fs(screen_pos: Vec2) -> Vec4 {
+fn get_ray_dir(uv: Vec2, pos: Vec3, look_at_pos: Vec3) -> Vec3 {
+    let forward = (look_at_pos - pos).normalize();
+    let right = Vec3::new(0.0, 1.0, 0.0).cross(forward).normalize();
+    let up = forward.cross(right);
+    (forward + uv.x() * right + uv.y() * up).normalize()
+}
+
+pub fn fs(constants: &ShaderConstants, frag_coord: Vec2) -> Vec4 {
+    let mut uv = (frag_coord - 0.5 * Vec2::new(constants.width as f32, constants.height as f32))
+        / constants.height as f32;
+    uv.set_y(-uv.y());
+
     // hard-code information because we can't bind buffers at the moment
     let eye_pos = Vec3::new(0.0, 0.0997, 0.2);
     let sun_pos = Vec3::new(0.0, 75.0, -1000.0);
-    let clip_to_world = Mat4::from_cols(
-        Vec4::new(-0.5522849, 0.0, 0.0, 0.0),
-        Vec4::new(0.0, 0.4096309, -0.061444636, 0.0),
-        Vec4::new(0.0, 99.99999, 199.99998, 999.99994),
-        Vec4::new(0.0, -0.14834046, -0.98893654, 0.0),
-    );
-
-    let world_pos = clip_to_world.transform_point3(screen_pos.extend(1.0));
-    let dir = (world_pos - eye_pos).normalize();
+    let dir = get_ray_dir(uv, eye_pos, sun_pos);
 
     // evaluate Preetham sky model
     let color = sky(dir, sun_pos);
@@ -174,8 +184,15 @@ pub fn fs(screen_pos: Vec2) -> Vec4 {
 
 #[allow(unused_attributes)]
 #[spirv(fragment)]
-pub fn main_fs(in_pos: Input<Vec2>, mut output: Output<Vec4>) {
-    let color = fs(in_pos.load());
+pub fn main_fs(
+    #[spirv(frag_coord)] in_frag_coord: Input<Vec4>,
+    #[spirv(push_constant)] constants: PushConstant<ShaderConstants>,
+    mut output: Output<Vec4>,
+) {
+    let constants = constants.load();
+
+    let frag_coord = Vec2::new(in_frag_coord.load().x(), in_frag_coord.load().y());
+    let color = fs(&constants, frag_coord);
     output.store(color);
 }
 
@@ -184,7 +201,6 @@ pub fn main_fs(in_pos: Input<Vec2>, mut output: Output<Vec4>) {
 pub fn main_vs(
     #[spirv(vertex_index)] vert_idx: Input<i32>,
     #[spirv(position)] mut builtin_pos: Output<Vec4>,
-    mut out_pos: Output<Vec2>,
 ) {
     let vert_idx = vert_idx.load();
 
@@ -194,7 +210,6 @@ pub fn main_vs(
     let pos = 2.0 * uv - Vec2::one();
 
     builtin_pos.store(pos.extend(0.0).extend(1.0));
-    out_pos.store(pos);
 }
 
 #[cfg(all(not(test), target_arch = "spirv"))]