render: proper number generator

blade-render/shader.wgsl

@@ -6,7 +6,7 @@ struct Parameters {
     depth: f32,
     cam_orientation: vec4<f32>,
     fov: vec2<f32>,
-    random_seed: u32,
+    frame_index: u32,
     debug_mode: u32,
 };
 
@@ -55,25 +55,96 @@ fn qrot(q: vec4<f32>, v: vec3<f32>) -> vec3<f32> {
     return v + 2.0*cross(q.xyz, cross(q.xyz,v) + q.w*v);
 }
 
-// https://en.wikipedia.org/wiki/Halton_sequence#Implementation_in_pseudocode
-fn halton(base: i32, start_index: i32) -> f32
-{
-    var result = 0.0f;
-    var f = 1.0f;
-    var index = start_index;
-    while (index > 0)
-    {
-        f /= f32(base);
-        result += f * f32(index % base);
-        index /= base;
-    }
-    return result;
+struct RandomState {
+    seed: u32,
+    index: u32,
+}
+
+// 32 bit Jenkins hash
+fn hash_jenkins(value: u32) -> u32 {
+    var a = value;
+    // http://burtleburtle.net/bob/hash/integer.html
+    a = (a + 0x7ed55d16u) + (a << 12u);
+    a = (a ^ 0xc761c23cu) ^ (a >> 19u);
+    a = (a + 0x165667b1u) + (a << 5u);
+    a = (a + 0xd3a2646cu) ^ (a << 9u);
+    a = (a + 0xfd7046c5u) + (a << 3u);
+    a = (a ^ 0xb55a4f09u) ^ (a >> 16u);
+    return a;
+}
+
+fn random_init(pixel_index: u32, frame_index: u32) -> RandomState {
+    var rs: RandomState;
+    rs.seed = hash_jenkins(pixel_index) + frame_index;
+    rs.index = 0u;
+    return rs;
+}
+
+fn rot32(x: u32, bits: u32) -> u32 {
+    return (x << bits) | (x >> (32u - bits));
+}
+
+// https://en.wikipedia.org/wiki/MurmurHash
+fn murmur3(rng: ptr<function, RandomState>) -> u32 {
+    let c1 = 0xcc9e2d51u;
+    let c2 = 0x1b873593u;
+    let r1 = 15u;
+    let r2 = 13u;
+    let m = 5u;
+    let n = 0xe6546b64u;
+
+    var hash = (*rng).seed;
+    (*rng).index += 1u;
+    var k = (*rng).index;
+    k *= c1;
+    k = rot32(k, r1);
+    k *= c2;
+
+    hash ^= k;
+    hash = rot32(hash, r2) * m + n;
+
+    hash ^= 4u;
+    hash ^= (hash >> 16u);
+    hash *= 0x85ebca6bu;
+    hash ^= (hash >> 13u);
+    hash *= 0xc2b2ae35u;
+    hash ^= (hash >> 16u);
+
+    return hash;
+}
+
+fn random_gen(rng: ptr<function, RandomState>) -> f32 {
+    let v = murmur3(rng);
+    let one = bitcast<u32>(1.0);
+    let mask = (1u << 23u) - 1u;
+    return bitcast<f32>((mask & v) | one) - 1.0;
+}
+
+fn sample_disk(rand: vec2<f32>) -> vec2<f32> {
+    let angle = 2.0 * PI * rand.x;
+    return vec2<f32>(cos(angle), sin(angle)) * sqrt(rand.y);
+}
+
+fn square(v: f32) -> f32 {
+    return v * v;
+}
+
+fn sample_uniform_hemisphere(rng: ptr<function, RandomState>) -> vec3<f32> {
+    // See (6-8) in https://mathworld.wolfram.com/SpherePointPicking.html
+    let r = random_gen(rng);
+    let h = random_gen(rng);
+    let tangential = sample_disk(vec2<f32>(r, 1.0 - square(h)));
+    return vec3<f32>(tangential.xy, h);
 }
 
-fn compute_hit_color(ri: RayIntersection, ray_dir: vec3<f32>, hit_world: vec3<f32>, random: u32) -> vec4<f32> {
+fn get_environment_radiance(dir: vec3<f32>) -> vec3<f32> {
+    return select(vec3<f32>(0.0), vec3<f32>(1.0), dir.y > 0.0);
+}
+
+fn compute_hit_color(ri: RayIntersection, ray_dir: vec3<f32>, hit_world: vec3<f32>, rng: ptr<function, RandomState>) -> vec3<f32> {
     let entry = hit_entries[ri.instance_custom_index + ri.geometry_index];
     if (parameters.debug_mode == DEBUG_MODE_DEPTH) {
-        return vec4<f32>(ri.t / parameters.depth);
+        return vec3<f32>(ri.t / parameters.depth);
     }
 
     var indices = ri.primitive_index * 3u + vec3<u32>(0u, 1u, 2u);
@@ -90,50 +161,46 @@ fn compute_hit_color(ri: RayIntersection, ray_dir: vec3<f32>, hit_world: vec3<f3
     );
 
     let barycentrics = vec3<f32>(1.0 - ri.barycentrics.x - ri.barycentrics.y, ri.barycentrics);
-    let pos_object = mat3x3(vertices[0].pos, vertices[1].pos, vertices[2].pos) * barycentrics;
-    let pos_world = ri.object_to_world * vec4<f32>(pos_object, 1.0);
+    //let pos_object = mat3x3(vertices[0].pos, vertices[1].pos, vertices[2].pos) * barycentrics;
+    //let pos_world = ri.object_to_world * vec4<f32>(pos_object, 1.0);
     let tex_coords = mat3x2(vertices[0].tex_coords, vertices[1].tex_coords, vertices[2].tex_coords) * barycentrics;
     let normal_rough = mat3x2(unpack2x16snorm(vertices[0].normal), unpack2x16snorm(vertices[1].normal), unpack2x16snorm(vertices[2].normal)) * barycentrics;
     let normal_object = vec3<f32>(normal_rough, sqrt(max(0.0, 1.0 - dot(normal_rough, normal_rough))));
     let object_to_world_rot = normalize(unpack4x8snorm(entry.rotation));
     let normal_world = qrot(object_to_world_rot, normal_object);
 
     // Note: this line allows to check the correctness of data passed in
-    //return vec4<f32>(abs(pos_world - hit_world) * 1000000.0, 1.0);
+    //return abs(pos_world - hit_world) * 1000000.0;
     if (parameters.debug_mode == DEBUG_MODE_NORMAL) {
-        return vec4<f32>(normal_world, 1.0);
+        return normal_world;
     }
 
     let pre_tangent = select(
         vec3<f32>(0.0, 0.0, 1.0),
         vec3<f32>(0.0, 1.0, 0.0),
-        abs(dot(normal_world, vec3<f32>(0.0, 1.0, 0.0)))<0.8);
+        abs(dot(normal_world, vec3<f32>(0.0, 1.0, 0.0))) < abs(dot(normal_world, vec3<f32>(0.0, 0.0, 1.0))));
     let bitangent = normalize(cross(normal_world, pre_tangent));
     let tangent = normalize(cross(bitangent, normal_world));
 
-    let halton_base = 4013;
-    let num_diffuse_samples = 1;
-    let origin = hit_world + 0.1 * normal_world;
+    let steradians_in_hemisphere = 2.0 * PI;
+    let lambert_brdf = 1.0 / PI;
+    let num_diffuse_samples = 10;
     var color = vec3<f32>(0.0);
     var rq: ray_query;
     for (var i = 0i; i < num_diffuse_samples; i += 1) {
-        //TODO: proper random generator
-        let r = random * u32(i+1);
-        let h1 = halton(halton_base, i32(r >> 16u));
-        let h2 = halton(halton_base, i32(r & 0xFFFFu));
-        let alpha = h1 * PI * 2.0;
-        let gamma = h2 * PI * 0.5;
-        let light_dir_tbn = vec3<f32>(sin(alpha) * cos(gamma), cos(alpha) * cos(gamma), sin(gamma));
+        let light_dir_tbn = sample_uniform_hemisphere(rng);
         let light_dir = light_dir_tbn.x * tangent + light_dir_tbn.y * bitangent + light_dir_tbn.z * normal_world;
-        if (light_dir.y <= 0.0) {
+        let lambert_term = light_dir_tbn.z;
+        let estimate_color = get_environment_radiance(light_dir) * lambert_term;
+        if (dot(estimate_color, estimate_color) < 0.01) {
             continue;
         }
 
-        rayQueryInitialize(&rq, acc_struct, RayDesc(RAY_FLAG_TERMINATE_ON_FIRST_HIT, 0xFFu, 0.1, parameters.depth, origin, light_dir));
+        rayQueryInitialize(&rq, acc_struct, RayDesc(RAY_FLAG_TERMINATE_ON_FIRST_HIT, 0xFFu, 0.5, parameters.depth, hit_world, light_dir));
         rayQueryProceed(&rq);
         let intersection = rayQueryGetCommittedIntersection(&rq);
         if (intersection.kind == RAY_QUERY_INTERSECTION_NONE) {
-            color += vec3<f32>(1.0);
+            color += estimate_color * steradians_in_hemisphere * lambert_brdf;
         }
     }
     color /= f32(num_diffuse_samples);
@@ -142,7 +209,8 @@ fn compute_hit_color(ri: RayIntersection, ray_dir: vec3<f32>, hit_world: vec3<f3
     let lod = 0.0; //TODO: this is actually complicated
     let base_color_sample = textureSampleLevel(textures[entry.base_color_texture], sampler_linear, tex_coords, lod);
 
-    return base_color_factor * base_color_sample * vec4<f32>(color, 1.0);
+    //return (base_color_factor * base_color_sample).xyz * color;
+    return color;
 }
 
 @compute @workgroup_size(8, 8)
@@ -154,24 +222,27 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
         return;
     }
 
-    let random = parameters.random_seed ^ (global_id.y << 16u) ^ global_id.x;
+    let global_index = global_id.y * target_size.x + global_id.x;
+    var rng = random_init(global_index, parameters.frame_index);
     let local_dir = vec3<f32>(ndc * tan(parameters.fov), 1.0);
     let world_dir = normalize(qrot(parameters.cam_orientation, local_dir));
 
     var rq: ray_query;
     var ray_pos = parameters.cam_position;
     var ray_dir = world_dir;
-    rayQueryInitialize(&rq, acc_struct, RayDesc(0x10u, 0xFFu, 0.1, parameters.depth, ray_pos, ray_dir));
+    rayQueryInitialize(&rq, acc_struct, RayDesc(0x10u, 0xFFu, 0.0, parameters.depth, ray_pos, ray_dir));
     var iterations = 0u;
     while (rayQueryProceed(&rq)) {iterations += 1u;}
     let intersection = rayQueryGetCommittedIntersection(&rq);
 
-    var color = vec4<f32>(0.0);
-    if (intersection.kind != RAY_QUERY_INTERSECTION_NONE) {
+    var color = vec3<f32>(0.0);
+    if (intersection.kind == RAY_QUERY_INTERSECTION_NONE) {
+        color = get_environment_radiance(ray_dir);
+    } else {
         let expected = ray_pos + intersection.t * ray_dir;
-        color = compute_hit_color(intersection, ray_dir, expected, random);
+        color = compute_hit_color(intersection, ray_dir, expected, &rng);
     }
-    textureStore(output, global_id.xy, color);
+    textureStore(output, global_id.xy, vec4<f32>(color, 1.0));
 }
 
 var input: texture_2d<f32>;

blade-render/src/renderer.rs

@@ -1,4 +1,4 @@
-use std::{mem, ptr, time};
+use std::{mem, ptr};
 
 const TARGET_FORMAT: blade::TextureFormat = blade::TextureFormat::Rgba16Float;
 const MAX_RESOURCES: u32 = 1000;
@@ -36,7 +36,7 @@ pub struct Renderer {
     samplers: Samplers,
     is_tlas_dirty: bool,
     screen_size: blade::Extent,
-    init_time: time::Instant,
+    frame_index: u32,
 }
 
 #[repr(C)]
@@ -46,7 +46,7 @@ struct Parameters {
     depth: f32,
     cam_orientation: [f32; 4],
     fov: [f32; 2],
-    random_seed: u32,
+    frame_index: u32,
     debug_mode: u32,
 }
 
@@ -290,7 +290,7 @@ impl Renderer {
             samplers,
             is_tlas_dirty: true,
             screen_size,
-            init_time: time::Instant::now(),
+            frame_index: 0,
         }
     }
 
@@ -333,6 +333,7 @@ impl Renderer {
         gpu: &blade::Context,
         temp_buffers: &mut Vec<blade::Buffer>,
     ) {
+        self.frame_index += 1;
         if self.is_tlas_dirty {
             self.is_tlas_dirty = false;
             if self.acceleration_structure != blade::AccelerationStructure::default() {
@@ -447,7 +448,12 @@ impl Renderer {
         }
     }
 
-    pub fn ray_trace(&self, command_encoder: &mut blade::CommandEncoder, camera: &super::Camera, debug_mode: DebugMode) {
+    pub fn ray_trace(
+        &self,
+        command_encoder: &mut blade::CommandEncoder,
+        camera: &super::Camera,
+        debug_mode: DebugMode,
+    ) {
         assert!(!self.is_tlas_dirty);
 
         let mut pass = command_encoder.compute();
@@ -468,7 +474,7 @@ impl Renderer {
                     depth: camera.depth,
                     cam_orientation: camera.rot.into(),
                     fov: [fov_x, camera.fov_y],
-                    random_seed: self.init_time.elapsed().subsec_micros(),
+                    frame_index: self.frame_index,
                     debug_mode: debug_mode as u32,
                 },
                 acc_struct: self.acceleration_structure,