@tinypoly/gpu-renderer 0.0.2 → 0.0.4

package/src/sceneAssets.ts

@@ -1,5 +1,6 @@
 import { environmentFromImage, type PackedEnvironment } from "./environment.js";
 import { compileShaderPrograms, type ShaderPrograms } from "./shaderProgram.js";
+import { bakeSkyOnDevice, compileSkyProgram } from "./skyBake.js";
 import { CpuComputePipeline, CpuEnvironment } from "./upstream.js";
 import type { SerializedMaterial, SerializedScene, SerializedTexture, UniformValue } from "./upstream.js";
 
@@ -13,6 +14,8 @@ export interface SceneAssets {
   scene: SerializedScene;
   programs: ShaderPrograms;
   sky: ProceduralSkyEnvironment | null;
+  /** The baked sky, kept by the renderer so a scene with the same sky does not bake it again. */
+  skyBake: { key: string; baked: Promise<ProceduralSkyEnvironment | null> } | null;
 }
 
 export interface ProceduralSkyEnvironment {
@@ -83,12 +86,10 @@ function bakeComputations(source: SerializedScene): SerializedScene {
   };
 }
 
-/** Bakes a procedural sky with the CPU renderer's own interpreter, at the size its environment uses. */
-function bakeSky(scene: SerializedScene): ProceduralSkyEnvironment | null {
-  const sky = scene.sky;
-
-  if (!sky) return null;
+type ProceduralSky = NonNullable<SerializedScene["sky"]>;
 
+/** Bakes a procedural sky with the CPU renderer's own interpreter: the fallback, when the device cannot run it. */
+function bakeSkyOnCpu(sky: ProceduralSky): Float32Array | null {
   const generator = CpuEnvironment.fromProcedural(sky, SKY_WIDTH);
   let step = generator.next();
 
@@ -97,22 +98,66 @@ function bakeSky(scene: SerializedScene): ProceduralSkyEnvironment | null {
   // The snapshot is the environment's public face; its texture is the baked panorama.
   const texture = step.value.snapshot().texture;
 
-  if (!texture || !(texture.data instanceof Float32Array)) return null;
+  return texture && texture.data instanceof Float32Array ? texture.data : null;
+}
+
+/**
+ * Bakes a procedural sky at the size its environment uses. On a device the sky's shader is translated and run there,
+ * over every texel at once; the CPU renderer's interpreter, which runs it texel by texel, is what is left when there is
+ * no device or the shader does not translate.
+ */
+async function bakeSky(sky: ProceduralSky, device?: GPUDevice): Promise<ProceduralSkyEnvironment | null> {
+  let radiance: Float32Array | null = null;
+
+  if (device) {
+    try {
+      radiance = await bakeSkyOnDevice(device, await compileSkyProgram(sky.fragmentShader), sky.uniforms, SKY_WIDTH);
+    } catch(error) {
+      console.warn("GPU renderer: the procedural sky is baked on the CPU, its shader did not run on the device.", error);
+    }
+  }
+
+  radiance ??= bakeSkyOnCpu(sky);
 
-  const environment = environmentFromImage(texture.data, texture.width, texture.height, "procedural-sky");
+  if (!radiance) return null;
+
+  const environment = environmentFromImage(radiance, SKY_WIDTH, SKY_WIDTH / 2, "procedural-sky");
 
   if (!environment) return null;
 
   return { environment, intensity: sky.intensity, background: sky.background, uniforms: sky.uniforms };
 }
 
-export async function prepareSceneAssets(source: SerializedScene): Promise<SceneAssets> {
+/** What decides a baked sky: the same shader with the same uniforms bakes the same panorama. */
+export function skyKey(sky: ProceduralSky | undefined): string {
+  return sky ? JSON.stringify([sky.fragmentShader, sky.uniforms]) : "";
+}
+
+export interface SceneAssetOptions {
+  /** Where the sky is baked; without one, on the CPU. */
+  device?: GPUDevice;
+  /** A sky baked for an earlier scene, reused when its key matches this scene's. */
+  sky?: { key: string; baked: Promise<ProceduralSkyEnvironment | null> } | null;
+}
+
+export async function prepareSceneAssets(
+  source: SerializedScene, options: SceneAssetOptions = {}): Promise<SceneAssets> {
   const scene = source.meshes.some(mesh => mesh.materials.some(material => material?.compute))
     ? bakeComputations(source)
     : source;
 
-  const sky = bakeSky(scene);
-  const programs = await compileShaderPrograms(scene, scene.sky?.foregroundShader);
+  const key = skyKey(scene.sky);
+  const reused = options.sky && options.sky.key === key ? options.sky.baked : null;
+  const baked = reused ?? (scene.sky ? bakeSky(scene.sky, options.device) : Promise.resolve(null));
+  const [sky, programs] = await Promise.all([baked, compileShaderPrograms(scene, scene.sky?.foregroundShader)]);
+  // The sky the scene uses, with its baked panorama, for the next scene to reuse.
+  const skyBake = scene.sky ? { key, baked } : null;
+
+  // The panorama comes from the shader and its uniforms; how strongly it lights and whether it paints the background
+  // are this scene's own.
+  const current = sky && scene.sky
+    ? { ...sky, intensity: scene.sky.intensity, background: scene.sky.background, uniforms: scene.sky.uniforms }
+    : null;
 
-  return { scene, programs, sky };
+  return { scene, programs, sky: current, skyBake };
 }

package/src/skyBake.ts

@@ -0,0 +1,113 @@
+import { compileProgram, type ShaderProgram } from "./shaderProgram.js";
+import { createShaderModule } from "./shaderModule.js";
+import type { UniformValue } from "./upstream.js";
+import { BVH_WGSL } from "./wgsl/bvh.wgsl.js";
+import { COMMON_WGSL } from "./wgsl/common.wgsl.js";
+import { RANDOM_WGSL } from "./wgsl/random.wgsl.js";
+import { SAMPLING_WGSL } from "./wgsl/sampling.wgsl.js";
+import { SHADER_SUPPORT_WGSL } from "./wgsl/shaderSupport.wgsl.js";
+import { TEXTURE_WGSL } from "./wgsl/texture.wgsl.js";
+
+const WORKGROUP = 8;
+
+/**
+ * The bake's own kernel: the kernel's support code up to the translated shaders, the sky's program, and an entry that
+ * runs it once per texel of the panorama, along the direction the CPU renderer's bake gives that texel. The output
+ * goes to the binding the path tracer accumulates into, which is the one read-write buffer the support code declares.
+ */
+export function skyBakeSource(program: ShaderProgram, width: number, height: number): string {
+  return `${COMMON_WGSL}${RANDOM_WGSL}${BVH_WGSL}${SAMPLING_WGSL}${TEXTURE_WGSL}${SHADER_SUPPORT_WGSL}
+${program.wgsl}
+
+@compute @workgroup_size(${WORKGROUP}, ${WORKGROUP})
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+  if (id.x >= ${width}u || id.y >= ${height}u) {
+    return;
+  }
+
+  let u = (f32(id.x) + 0.5) / ${width}.0;
+  let v = (f32(id.y) + 0.5) / ${height}.0;
+  let phi = (u - 0.5) * 2.0 * PI;
+  let theta = v * PI;
+  var input: ShaderInput;
+  input.data = 0u;
+  input.direction = vec3<f32>(sin(theta) * cos(phi), cos(theta), sin(theta) * sin(phi));
+  input.fragCoord = vec4<f32>(f32(id.x) + 0.5, f32(id.y) + 0.5, 0.0, 1.0);
+  input.diffuseColor = vec4<f32>(0.0, 0.0, 0.0, 1.0);
+
+  accumulation[id.y * ${width}u + id.x] = p${program.id}_run(input).color;
+}
+`;
+}
+
+/** Translates a procedural sky's shader for the bake; the sky reads its direction from `vDirection`, as the CPU's does. */
+export function compileSkyProgram(fragmentShader: string): Promise<ShaderProgram> {
+  return compileProgram(0, "foreground", { fragment: fragmentShader, defines: {} }, true);
+}
+
+/**
+ * Runs a procedural sky's shader on the device over an equirectangular panorama and reads the radiance back, as RGBA
+ * floats, row by row from the top: the same image the CPU renderer's interpreter bakes, in a fraction of the time. A
+ * texel that comes out not finite reads black, where the CPU bake would have refused the whole sky.
+ */
+export async function bakeSkyOnDevice(
+  device: GPUDevice, program: ShaderProgram, uniforms: Record<string, UniformValue>, width: number,
+): Promise<Float32Array> {
+  const height = width / 2;
+  const texels = width * height;
+  const words = program.packUniforms(uniforms, new Map());
+  const module = createShaderModule(device, "sky bake", skyBakeSource(program, width, height));
+
+  const layout = device.createBindGroupLayout({
+    label: "sky bake",
+    entries: [
+      { binding: 7, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+      { binding: 8, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
+    ],
+  });
+
+  const pipeline = await device.createComputePipelineAsync({
+    label: "sky bake",
+    layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
+    compute: { module, entryPoint: "main" },
+  });
+
+  const output = device.createBuffer({ size: texels * 16, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC });
+
+  const data = device.createBuffer({
+    size: Math.max(16, words.byteLength), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+  });
+
+  const staging = device.createBuffer({ size: texels * 16, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
+
+  try {
+    device.queue.writeBuffer(data, 0, words);
+
+    const encoder = device.createCommandEncoder({ label: "sky bake" });
+    const pass = encoder.beginComputePass({ label: "sky bake" });
+
+    pass.setPipeline(pipeline);
+    pass.setBindGroup(0, device.createBindGroup({
+      layout,
+      entries: [{ binding: 7, resource: { buffer: output } }, { binding: 8, resource: { buffer: data } }],
+    }));
+    pass.dispatchWorkgroups(Math.ceil(width / WORKGROUP), Math.ceil(height / WORKGROUP));
+    pass.end();
+    encoder.copyBufferToBuffer(output, 0, staging, 0, texels * 16);
+    device.queue.submit([encoder.finish()]);
+
+    await staging.mapAsync(GPUMapMode.READ);
+
+    const radiance = new Float32Array(staging.getMappedRange().slice(0));
+
+    staging.unmap();
+
+    for (let i = 0; i < radiance.length; i++) if (!Number.isFinite(radiance[i])) radiance[i] = 0;
+
+    return radiance;
+  } finally {
+    output.destroy();
+    data.destroy();
+    staging.destroy();
+  }
+}