cornell: Add reflections to raytracer (#245)

And also fix a couple of misc things.

Author: ben-clayton

src/sample/cornell/common.ts

@@ -82,16 +82,21 @@ export default class Common {
   }
 
   /** Updates the uniform buffer data */
-  update(args: { rotateCamera: boolean }) {
+  update(params: { rotateCamera: boolean; aspect: number }) {
     const viewMatrix = mat4.create();
     const mvp = mat4.create();
     const invMVP = mat4.create();
 
-    const aspect = 1; // canvas.width / canvas.height;
     const projectionMatrix = mat4.create();
-    mat4.perspective(projectionMatrix, (2 * Math.PI) / 8, aspect, 0.5, 100);
+    mat4.perspective(
+      projectionMatrix,
+      (2 * Math.PI) / 8,
+      params.aspect,
+      0.5,
+      100
+    );
 
-    const viewRotation = args.rotateCamera ? this.frame / 1000 : 0;
+    const viewRotation = params.rotateCamera ? this.frame / 1000 : 0;
 
     mat4.lookAt(
       viewMatrix,

src/sample/cornell/common.wgsl

@@ -107,8 +107,8 @@ fn rand() -> f32 {
                  11279  * 2539 * 23,
                  7919   * 631  * 5 * 3);
 
-  rnd = (rnd * C) ^ (rnd.yzx >> vec3(2u));
-  return f32(rnd.x) / f32(0xffffffff);
+  rnd = (rnd * C) ^ (rnd.yzx >> vec3(4u));
+  return f32(rnd.x ^ rnd.y) / f32(0xffffffff);
 }
 
 // Returns a random point within a unit sphere centered at (0,0,0).

src/sample/cornell/main.ts

@@ -84,6 +84,7 @@ const init: SampleInit = async ({ canvas, pageState, gui }) => {
 
     common.update({
       rotateCamera: params.rotateCamera,
+      aspect: canvas.width / canvas.height,
     });
     radiosity.run(commandEncoder);
 

src/sample/cornell/rasterizer.ts

@@ -112,7 +112,7 @@ export default class Rasterizer {
       fragment: {
         module: mod,
         entryPoint: 'fs_main',
-        targets: [{ format: 'rgba16float' }],
+        targets: [{ format: framebuffer.format }],
       },
       primitive: {
         topology: 'triangle-list',

src/sample/cornell/raytracer.ts

@@ -49,7 +49,7 @@ export default class Raytracer {
           visibility: GPUShaderStage.COMPUTE,
           storageTexture: {
             access: 'write-only',
-            format: 'rgba16float',
+            format: framebuffer.format,
             viewDimension: '2d',
           },
         },

src/sample/cornell/raytracer.wgsl

@@ -10,27 +10,53 @@
 override WorkgroupSizeX : u32;
 override WorkgroupSizeY : u32;
 
+const NumReflectionRays = 5;
+
 @compute @workgroup_size(WorkgroupSizeX, WorkgroupSizeY)
 fn main(@builtin(global_invocation_id) invocation_id : vec3u) {
-  // Calculate the fragment's NDC coordinates for the intersection of the near
-  // clip plane and far clip plane
-  let uv = vec2f(invocation_id.xy) / vec2f(textureDimensions(framebuffer).xy);
-  let ndcXY = (uv - 0.5) * vec2(2, -2);
-
-  // Transform the coordinates back into world space
-  var near = common_uniforms.inv_mvp * vec4f(ndcXY, 0.0, 1);
-  var far = common_uniforms.inv_mvp * vec4f(ndcXY, 1, 1);
-  near /= near.w;
-  far /= far.w;
-
-  // Create a ray that starts at the near clip plane, heading in the fragment's
-  // z-direction, and raytrace to find the nearest quad that the ray intersects.
-  let ray = Ray(near.xyz, normalize(far.xyz - near.xyz));
-  let hit = raytrace(ray);
-  let quad = quads[hit.quad];
+  if (all(invocation_id.xy < textureDimensions(framebuffer))) {
+    init_rand(invocation_id);
+
+    // Calculate the fragment's NDC coordinates for the intersection of the near
+    // clip plane and far clip plane
+    let uv = vec2f(invocation_id.xy) / vec2f(textureDimensions(framebuffer).xy);
+    let ndcXY = (uv - 0.5) * vec2(2, -2);
+
+    // Transform the coordinates back into world space
+    var near = common_uniforms.inv_mvp * vec4f(ndcXY, 0.0, 1);
+    var far = common_uniforms.inv_mvp * vec4f(ndcXY, 1, 1);
+    near /= near.w;
+    far /= far.w;
+
+    // Create a ray that starts at the near clip plane, heading in the fragment's
+    // z-direction, and raytrace to find the nearest quad that the ray intersects.
+    let ray = Ray(near.xyz, normalize(far.xyz - near.xyz));
+    let hit = raytrace(ray);
+
+    let hit_color = sample_hit(hit);
+    var normal = quads[hit.quad].plane.xyz;
 
+    // Fire a few rays off the surface to collect some reflections
+    let bounce = reflect(ray.dir, normal);
+    var reflection : vec3f;
+    for (var i = 0; i < NumReflectionRays; i++) {
+      let reflection_dir = normalize(bounce + rand_unit_sphere()*0.1);
+      let reflection_ray = Ray(hit.pos + bounce * 1e-5, reflection_dir);
+      let reflection_hit = raytrace(reflection_ray);
+      reflection += sample_hit(reflection_hit);
+    }
+    let color = mix(reflection / NumReflectionRays, hit_color, 0.95);
+
+    textureStore(framebuffer, invocation_id.xy, vec4(color, 1));
+  }
+}
+
+
+// Returns the sampled hit quad's lightmap at 'hit.uv', and adds the quad's
+// emissive value.
+fn sample_hit(hit : HitInfo) -> vec3f {
+  let quad = quads[hit.quad];
   // Sample the quad's lightmap, and add emissive.
-  let color = textureSampleLevel(lightmap, smpl, hit.uv, hit.quad, 0).rgb +
-              quad.emissive * quad.color;
-  textureStore(framebuffer, invocation_id.xy, vec4(color, 1));
+  return textureSampleLevel(lightmap, smpl, hit.uv, hit.quad, 0).rgb +
+         quad.emissive * quad.color;
 }