Skip to content

Commit

Permalink
Upgraded to wgpu 0.11
Browse files Browse the repository at this point in the history 
Lots of changes to conform to new wgpu API

Also a bit of refactoring
  • Loading branch information
@tombh
tombh committed Nov 7, 2021
1 parent b24e4cd commit ad16899
Show file tree
Hide file tree
Showing 10 changed files with 733 additions and 843 deletions.
841 changes: 334 additions & 507 deletions Cargo.lock
View file

Large diffs are not rendered by default.

10 changes: 6 additions & 4 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ name = "wrach"
version = "0.0.0"
authors = []
edition = "2018"
resolver = "2"

[features]
default = ["use-compiled-tools"]
Expand All @@ -12,9 +13,10 @@ use-compiled-tools = ["spirv-builder/use-compiled-tools"]
[dependencies]
shaders = { path = "shaders" }
futures = { version = "0.3", default-features = false, features = ["std", "executor"] }
wgpu = "0.6.0"
winit = { version = "0.23", features = ["web-sys"] }
wgpu-subscriber = "0.1.0"
wgpu = { version = "0.11.0", features = ["spirv"] }
wgpu-hal = "=0.11.2"
winit = { version = "0.25", features = ["web-sys"] }
bytemuck = { version = "1.6.3", features = ["derive"] }

[build-dependencies]
spirv-std = { git = "https://github.com/EmbarkStudios/rust-gpu", rev = "04146858", features = [ "glam" ] }
Expand All @@ -38,5 +40,5 @@ opt-level = 3

# HACK(eddyb) don't optimize the shader crate, to avoid `spirv-opt` taking
# a long time (10 minutes if itself was optimized, over an hour otherwise).
[profile.release.package."ray-tracing-shaders"]
[profile.release.package."shaders"]
opt-level = 0
4 changes: 2 additions & 2 deletions README.md
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
# Ray Tracing using rust-gpu
# Wrach

Implementation of ray tracing using `rust-gpu`. The biolerplate setup code [`rust-gpu-shadertoys`](https://github.com/LykenSol/rust-gpu-shadertoys).
A game like Noita, but using the GPU
49 changes: 0 additions & 49 deletions default.nix
View file

This file was deleted.

1 change: 1 addition & 0 deletions rustfmt.toml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
edition = "2018"
1 change: 1 addition & 0 deletions shaders/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,4 @@ crate-type = ["dylib"]

[dependencies]
spirv-std = { git = "https://github.com/EmbarkStudios/rust-gpu", rev = "04146858", features = [ "glam" ] }
bytemuck = { version = "1.6.3", features = ["derive"] }
28 changes: 28 additions & 0 deletions shaders/src/compute.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
pub fn collatz(mut n: u32) -> Option<u32> {
let mut i = 0;
if n == 0 {
return None;
}
while n != 1 {
n = if n % 2 == 0 {
n / 2
} else {
// Overflow? (i.e. 3*n + 1 > 0xffff_ffff)
if n >= 0x5555_5555 {
return None;
}
// TODO: Use this instead when/if checked add/mul can work: n.checked_mul(3)?.checked_add(1)?
3 * n + 1
};
i += 1;
}
Some(i)
}

// Work around https://github.com/EmbarkStudios/rust-gpu/issues/677
pub fn unwrap_or_max(option: Option<u32>) -> u32 {
match option {
Some(inner) => inner,
None => u32::MAX,
}
}
209 changes: 25 additions & 184 deletions shaders/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,21 +9,19 @@
// HACK(eddyb) can't easily see warnings otherwise from `spirv-builder` builds.
#![deny(warnings)]

pub use spirv_std::glam;

#[cfg(not(target_arch = "spirv"))]
use spirv_std::macros::spirv;

// Note: This cfg is incorrect on its surface, it really should be "are we compiling with std", but
// we tie #[no_std] above to the same condition, so it's fine.
#[cfg(target_arch = "spirv")]
use spirv_std::num_traits::Float;

use core::f32::consts::PI;
use bytemuck::{Pod, Zeroable};
use glam::{vec2, UVec3, Vec2, Vec4};
pub use spirv_std::glam;
use vertex::fs;

use glam::{const_vec3, vec2, vec3, Vec2, Vec3, Vec4};
mod compute;
mod vertex;

#[derive(Copy, Clone)]
#[derive(Copy, Clone, Pod, Zeroable)]
#[repr(C)]
pub struct ShaderConstants {
pub width: u32,
pub height: u32,
Expand All @@ -34,179 +32,26 @@ pub struct ShaderConstants {
pub drag_start_y: f32,
pub drag_end_x: f32,
pub drag_end_y: f32,
pub mouse_left_pressed: u32,
pub mouse_left_clicked: u32,
}

pub fn saturate(x: f32) -> f32 {
x.max(0.0).min(1.0)
}

pub fn pow(v: Vec3, power: f32) -> Vec3 {
vec3(v.x.powf(power), v.y.powf(power), v.z.powf(power))
}

pub fn exp(v: Vec3) -> Vec3 {
vec3(v.x.exp(), v.y.exp(), v.z.exp())
}

/// Based on: <https://seblagarde.wordpress.com/2014/12/01/inverse-trigonometric-functions-gpu-optimization-for-amd-gcn-architecture/>
pub fn acos_approx(v: f32) -> f32 {
let x = v.abs();
let mut res = -0.155972 * x + 1.56467; // p(x)
res *= (1.0f32 - x).sqrt();

if v >= 0.0 {
res
} else {
PI - res
}
}

pub fn smoothstep(edge0: f32, edge1: f32, x: f32) -> f32 {
// Scale, bias and saturate x to 0..1 range
let x = saturate((x - edge0) / (edge1 - edge0));
// Evaluate polynomial
x * x * (3.0 - 2.0 * x)
}

const DEPOLARIZATION_FACTOR: f32 = 0.035;
const MIE_COEFFICIENT: f32 = 0.005;
const MIE_DIRECTIONAL_G: f32 = 0.8;
const MIE_K_COEFFICIENT: Vec3 = const_vec3!([0.686, 0.678, 0.666]);
const MIE_V: f32 = 4.0;
const MIE_ZENITH_LENGTH: f32 = 1.25e3;
const NUM_MOLECULES: f32 = 2.542e25f32;
const PRIMARIES: Vec3 = const_vec3!([6.8e-7f32, 5.5e-7f32, 4.5e-7f32]);
const RAYLEIGH: f32 = 1.0;
const RAYLEIGH_ZENITH_LENGTH: f32 = 8.4e3;
const REFRACTIVE_INDEX: f32 = 1.0003;
const SUN_ANGULAR_DIAMETER_DEGREES: f32 = 0.0093333;
const SUN_INTENSITY_FACTOR: f32 = 1000.0;
const SUN_INTENSITY_FALLOFF_STEEPNESS: f32 = 1.5;
const TURBIDITY: f32 = 2.0;

pub fn tonemap(col: Vec3) -> Vec3 {
// see https://www.desmos.com/calculator/0eo9pzo1at
const A: f32 = 2.35;
const B: f32 = 2.8826666;
const C: f32 = 789.7459;
const D: f32 = 0.935;

let z = pow(col, A);
z / (pow(z, D) * B + Vec3::splat(C))
}

fn total_rayleigh(lambda: Vec3) -> Vec3 {
(8.0 * PI.powf(3.0)
* (REFRACTIVE_INDEX.powf(2.0) - 1.0).powf(2.0)
* (6.0 + 3.0 * DEPOLARIZATION_FACTOR))
/ (3.0 * NUM_MOLECULES * pow(lambda, 4.0) * (6.0 - 7.0 * DEPOLARIZATION_FACTOR))
}
/// Bit mask of the pressed buttons (0 = Left, 1 = Middle, 2 = Right).
pub mouse_button_pressed: u32,

fn total_mie(lambda: Vec3, k: Vec3, t: f32) -> Vec3 {
let c = 0.2 * t * 10e-18;
0.434 * c * PI * pow((2.0 * PI) / lambda, MIE_V - 2.0) * k
/// The last time each mouse button (Left, Middle or Right) was pressed,
/// or `f32::NEG_INFINITY` for buttons which haven't been pressed yet.
///
/// If this is the first frame after the press of some button, that button's
/// entry in `mouse_button_press_time` will exactly equal `time`.
pub mouse_button_press_time: [f32; 3],
}

fn rayleigh_phase(cos_theta: f32) -> f32 {
(3.0 / (16.0 * PI)) * (1.0 + cos_theta.powf(2.0))
}

fn henyey_greenstein_phase(cos_theta: f32, g: f32) -> f32 {
(1.0 / (4.0 * PI)) * ((1.0 - g.powf(2.0)) / (1.0 - 2.0 * g * cos_theta + g.powf(2.0)).powf(1.5))
}

fn sun_intensity(zenith_angle_cos: f32) -> f32 {
let cutoff_angle = PI / 1.95; // Earth shadow hack
SUN_INTENSITY_FACTOR
* 0.0f32.max(
1.0 - (-((cutoff_angle - acos_approx(zenith_angle_cos))
/ SUN_INTENSITY_FALLOFF_STEEPNESS))
.exp(),
)
}

fn sky(dir: Vec3, sun_position: Vec3) -> Vec3 {
let up = vec3(0.0, 1.0, 0.0);
let sunfade = 1.0 - (1.0 - saturate(sun_position.y / 450000.0).exp());
let rayleigh_coefficient = RAYLEIGH - (1.0 * (1.0 - sunfade));
let beta_r = total_rayleigh(PRIMARIES) * rayleigh_coefficient;

// Mie coefficient
let beta_m = total_mie(PRIMARIES, MIE_K_COEFFICIENT, TURBIDITY) * MIE_COEFFICIENT;

// Optical length, cutoff angle at 90 to avoid singularity
let zenith_angle = acos_approx(up.dot(dir).max(0.0));
let denom = (zenith_angle).cos() + 0.15 * (93.885 - ((zenith_angle * 180.0) / PI)).powf(-1.253);

let s_r = RAYLEIGH_ZENITH_LENGTH / denom;
let s_m = MIE_ZENITH_LENGTH / denom;

// Combined extinction factor
let fex = exp(-(beta_r * s_r + beta_m * s_m));

// In-scattering
let sun_direction = sun_position.normalize();
let cos_theta = dir.dot(sun_direction);
let beta_r_theta = beta_r * rayleigh_phase(cos_theta * 0.5 + 0.5);

let beta_m_theta = beta_m * henyey_greenstein_phase(cos_theta, MIE_DIRECTIONAL_G);
let sun_e = sun_intensity(sun_direction.dot(up));
let mut lin = pow(
sun_e * ((beta_r_theta + beta_m_theta) / (beta_r + beta_m)) * (Vec3::splat(1.0) - fex),
1.5,
);

lin *= Vec3::splat(1.0).lerp(
pow(
sun_e * ((beta_r_theta + beta_m_theta) / (beta_r + beta_m)) * fex,
0.5,
),
saturate((1.0 - up.dot(sun_direction)).powf(5.0)),
);

// Composition + solar disc
let sun_angular_diameter_cos = SUN_ANGULAR_DIAMETER_DEGREES.cos();
let sundisk = smoothstep(
sun_angular_diameter_cos,
sun_angular_diameter_cos + 0.00002,
cos_theta,
);
let mut l0 = 0.1 * fex;
l0 += sun_e * 19000.0 * fex * sundisk;

lin + l0
}

fn get_ray_dir(uv: Vec2, pos: Vec3, look_at_pos: Vec3) -> Vec3 {
let forward = (look_at_pos - pos).normalize();
let right = vec3(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(
constants.width as f32 - constants.cursor_x,
constants.height as f32 - constants.cursor_y
)
) / constants.height as f32;
uv.y = -uv.y;

// hard-code information because we can't bind buffers at the moment
let eye_pos = vec3(0.0, 0.0997, 0.2);
let sun_pos = vec3(0.0, 75.0, -1000.0);
let dir = get_ray_dir(uv, eye_pos, sun_pos);

// evaluate Preetham sky model
let color = sky(dir, sun_pos);

// Tonemapping
let color = color.max(Vec3::splat(0.0)).min(Vec3::splat(1024.0));

tonemap(color).extend(1.0)
// LocalSize/numthreads of (x = 64, y = 1, z = 1)
#[spirv(compute(threads(64)))]
pub fn main_cs(
#[spirv(global_invocation_id)] id: UVec3,
#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] prime_indices: &mut [u32],
) {
let index = id.x as usize;
prime_indices[index] = compute::unwrap_or_max(compute::collatz(prime_indices[index]));
}

#[spirv(fragment)]
Expand All @@ -220,15 +65,11 @@ pub fn main_fs(
}

#[spirv(vertex)]
pub fn main_vs(
#[spirv(vertex_index)] vert_idx: i32,
#[spirv(position)] builtin_pos: &mut Vec4
) {
pub fn main_vs(#[spirv(vertex_index)] vert_idx: i32, #[spirv(position)] builtin_pos: &mut Vec4) {
// Create a "full screen triangle" by mapping the vertex index.
// ported from https://www.saschawillems.de/blog/2016/08/13/vulkan-tutorial-on-rendering-a-fullscreen-quad-without-buffers/
let uv = vec2(((vert_idx << 1) & 2) as f32, (vert_idx & 2) as f32);
let pos = 2.0 * uv - Vec2::ONE;

*builtin_pos = pos.extend(0.0).extend(1.0);
}

Loading

0 comments on commit ad16899

Please sign in to comment.