Optimize viewer shaders and camera controls; sample GPU profiles asynchronously

Amp-Thread-ID: https://ampcode.com/threads/T-01a0b382-1e79-712b-abd1-d29201c9bb5e
Co-authored-by: Heaust Azure <heaust.azure@gmail.com>

core/Cargo.toml

@@ -9,6 +9,7 @@ crate-type = ["cdylib"]
 path = "lib.rs"
 
 [dependencies]
+futures-channel = "0.3"
 js-sys = "0.3"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"

core/index.js

@@ -1,3 +1,5 @@
+import { SharedProfile } from "./profile.js";
+
 const views = Object.freeze({ f32: Float32Array, u32: Uint32Array, i32: Int32Array });
 
 export class SharedRows {
@@ -49,6 +51,8 @@ export class YawnCore {
   #arrays = new Map();
   #pending = new Map();
   #profileListeners = new Set();
+  #profile = new SharedProfile();
+  #profileTimer;
   #next = 1;
   /** Optional WebGPU features the device has (WebGPU spelling); populated once `ready` resolves. */
   features = new Set();
@@ -64,7 +68,7 @@ export class YawnCore {
     this.#worker.addEventListener("messageerror", () => this.#fail("WORKER_ERROR"));
     this.#worker.start?.();
     const offscreen = canvas.transferControlToOffscreen?.() ?? canvas;
-    this.ready = this.#request("init", { canvas: offscreen, arenaBytes }, [offscreen]).then(async result => {
+    this.ready = this.#request("init", { canvas: offscreen, arenaBytes, profileBuffer: this.#profile.buffer }, [offscreen]).then(async result => {
       this.#buffer = result.buffer;
       this.features = new Set(result.features);
       for (const descriptor of result.rows) this.#arrays.set(
@@ -195,7 +199,17 @@ export class YawnCore {
   onProfile(listener) {
     if (typeof listener !== "function") throw new TypeError("PROFILE_LISTENER");
     this.#profileListeners.add(listener);
-    return () => this.#profileListeners.delete(listener);
+    if (!this.#profileTimer) this.#profileTimer = setInterval(() => {
+      const stats = this.#profile.read();
+      if (stats) for (const listener of this.#profileListeners) listener(stats);
+    }, 100);
+    return () => {
+      this.#profileListeners.delete(listener);
+      if (!this.#profileListeners.size) {
+        clearInterval(this.#profileTimer);
+        this.#profileTimer = undefined;
+      }
+    };
   }
 
   array(name) {
@@ -213,10 +227,6 @@ export class YawnCore {
   }
 
   #message(message) {
-    if (message?.type === "profile") {
-      for (const listener of this.#profileListeners) listener(message.stats);
-      return;
-    }
     const pending = this.#pending.get(message?.request);
     if (!pending) return;
     this.#pending.delete(message.request);
@@ -231,6 +241,7 @@ export class YawnCore {
 
   dispose() {
     this.#fail("DISPOSED");
+    clearInterval(this.#profileTimer);
     this.#profileListeners.clear();
     this.#worker.terminate();
   }

core/lib.rs

@@ -302,7 +302,8 @@ impl Core {
         supported
     }
 
-    pub fn take_profile(&self) -> Option<String> {
-        self.render.take_profile()
+    /// Installs the worker-local sink that publishes completed GPU timings into shared memory.
+    pub fn on_profile(&self, callback: js_sys::Function) {
+        self.render.on_profile(callback);
     }
 }

core/package.json

@@ -4,5 +4,13 @@
   "description": "Rust/WASM render data and render graph core",
   "type": "module",
   "exports": "./index.js",
-  "files": ["index.js", "worker.js", "pkg"]
+  "scripts": {
+    "test": "node --test profile.test.js"
+  },
+  "files": [
+    "index.js",
+    "worker.js",
+    "profile.js",
+    "pkg"
+  ]
 }

core/profile.js

@@ -0,0 +1,42 @@
+// Single-writer/single-reader latest-sample mailbox. Header: state, UTF-8 byte length.
+// State 0 = empty, 1 = owned by a reader/writer, 2 = complete sample.
+// Neither side waits: an unread sample can be replaced, and a busy publication is skipped.
+export class SharedProfile {
+  constructor(buffer = new SharedArrayBuffer(1024 * 1024)) {
+    this.buffer = buffer;
+    this.header = new Int32Array(buffer, 0, 2);
+    this.bytes = new Uint8Array(buffer, 8);
+    this.encoder = new TextEncoder();
+    this.decoder = new TextDecoder();
+  }
+
+  write(json) {
+    const bytes = this.encoder.encode(json);
+    if (bytes.length > this.bytes.length) return false;
+    const state = Atomics.load(this.header, 0);
+    if (
+      state === 1 ||
+      Atomics.compareExchange(this.header, 0, state, 1) !== state
+    )
+      return false;
+    this.bytes.set(bytes);
+    this.header[1] = bytes.length;
+    Atomics.store(this.header, 0, 2);
+    return true;
+  }
+
+  read() {
+    if (Atomics.compareExchange(this.header, 0, 2, 1) !== 2) return;
+    let bytes;
+    try {
+      bytes = this.bytes.slice(0, this.header[1]);
+    } finally {
+      Atomics.store(this.header, 0, 0);
+    }
+    return JSON.parse(this.decoder.decode(bytes));
+  }
+
+  clear() {
+    Atomics.compareExchange(this.header, 0, 2, 0);
+  }
+}

core/profile.test.js

@@ -0,0 +1,122 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { Worker } from "node:worker_threads";
+import { SharedProfile } from "./profile.js";
+import { YawnCore } from "./index.js";
+
+test("shared profiles replace unread samples and exclude overlapping reads/writes", () => {
+  const writer = new SharedProfile(new SharedArrayBuffer(512));
+  const reader = new SharedProfile(writer.buffer);
+  assert.equal(reader.read(), undefined);
+  writer.write(JSON.stringify({ frame: 1, name: "GPU → SAB" }));
+  writer.write(JSON.stringify({ frame: 2, name: "light + opaque" }));
+  assert.deepEqual(reader.read(), { frame: 2, name: "light + opaque" });
+  assert.equal(reader.read(), undefined);
+  Atomics.store(writer.header, 0, 1);
+  assert.equal(writer.write('{"frame":3}'), false);
+  assert.equal(reader.read(), undefined);
+  Atomics.store(writer.header, 0, 0);
+  assert.equal(writer.write(" ".repeat(513)), false);
+  assert.equal(writer.write('{"frame":4}'), true);
+  writer.clear();
+  assert.equal(reader.read(), undefined);
+});
+
+test("concurrent shared profile publication never mixes samples", async () => {
+  const reader = new SharedProfile();
+  const done = new Int32Array(new SharedArrayBuffer(4));
+  const worker = new Worker(
+    `
+    const {workerData} = require('node:worker_threads');
+    (async () => {
+      const {SharedProfile} = await import(workerData.module);
+      const writer = new SharedProfile(workerData.buffer);
+      for (let frame = 1; frame <= 2000; frame++) {
+        const json = JSON.stringify({frame, values: Array(128).fill(frame * 3 + 7), name: 'GPU → SAB'});
+        while (!writer.write(json)) await new Promise(r => setImmediate(r));
+      }
+      Atomics.store(new Int32Array(workerData.done), 0, 1);
+    })();
+  `,
+    {
+      eval: true,
+      workerData: {
+        module: new URL("./profile.js", import.meta.url).href,
+        buffer: reader.buffer,
+        done: done.buffer,
+      },
+    },
+  );
+  const exit = new Promise((resolve, reject) => {
+    worker.once("error", reject);
+    worker.once("exit", (code) =>
+      code ? reject(new Error(`Worker exit ${code}`)) : resolve(),
+    );
+  });
+  try {
+    let last = 0;
+    while (!Atomics.load(done, 0) || Atomics.load(reader.header, 0) === 2) {
+      const sample = reader.read();
+      if (sample) {
+        assert.ok(sample.frame > last);
+        assert.deepEqual(sample.values, Array(128).fill(sample.frame * 3 + 7));
+        assert.equal(sample.name, "GPU → SAB");
+        last = sample.frame;
+      }
+      await new Promise((r) => setImmediate(r));
+    }
+    await exit;
+    assert.equal(last, 2000);
+  } finally {
+    await worker.terminate();
+  }
+});
+
+test("onProfile reads SAB without result messages and stops on unsubscribe/dispose", async (t) => {
+  t.mock.timers.enable({ apis: ["setInterval"] });
+  const messages = [];
+  let receive;
+  let writer;
+  const worker = {
+    addEventListener(type, listener) {
+      if (type === "message") receive = listener;
+    },
+    postMessage(message) {
+      messages.push(message.type);
+      writer = new SharedProfile(message.profileBuffer);
+      queueMicrotask(() =>
+        receive({
+          data: {
+            request: message.request,
+            result: {
+              buffer: new SharedArrayBuffer(64),
+              features: [],
+              rows: [],
+            },
+          },
+        }),
+      );
+    },
+    terminate() {},
+  };
+  const core = new YawnCore({}, { workerFactory: () => worker });
+  await core.ready;
+  const samples = [];
+  const stop = core.onProfile((sample) => samples.push(sample));
+  writer.write('{"frame":9,"milliseconds":1.25}');
+  t.mock.timers.tick(99);
+  assert.equal(samples.length, 0);
+  t.mock.timers.tick(1);
+  assert.deepEqual(samples, [{ frame: 9, milliseconds: 1.25 }]);
+  t.mock.timers.tick(100);
+  assert.equal(samples.length, 1);
+  assert.deepEqual(messages, ["init"]);
+  stop();
+  writer.write('{"frame":10}');
+  t.mock.timers.tick(100);
+  assert.equal(samples.length, 1);
+  core.onProfile((sample) => samples.push(sample));
+  core.dispose();
+  t.mock.timers.tick(100);
+  assert.equal(samples.length, 1);
+});

core/render_graph/gpu_resource.rs

@@ -1,8 +1,8 @@
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::num::NonZeroU64;
-use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::Arc;
 
+use futures_channel::oneshot;
 use serde::de::DeserializeOwned;
 use serde_json::Value;
 
@@ -31,7 +31,7 @@ pub struct GpuProfiler {
 
 pub struct ProfileMap {
     readback: Arc<wgpu::Buffer>,
-    state: Arc<AtomicU8>,
+    completion: oneshot::Receiver<Result<(), wgpu::BufferAsyncError>>,
     labels: Vec<String>,
     timestamp_period: f32,
 }
@@ -417,16 +417,15 @@ impl GpuResources {
 
     pub fn map_profile(&self, timestamp_period: f32) -> Option<ProfileMap> {
         let profiler = self.profiler.as_ref()?;
-        let state = Arc::new(AtomicU8::new(0));
-        let callback = state.clone();
+        let (complete, completion) = oneshot::channel();
         profiler
             .readback
             .map_async(wgpu::MapMode::Read, .., move |result| {
-                callback.store(if result.is_ok() { 1 } else { 2 }, Ordering::Release);
+                let _ = complete.send(result);
             });
         Some(ProfileMap {
             readback: profiler.readback.clone(),
-            state,
+            completion,
             labels: self
                 .passes
                 .iter()
@@ -438,11 +437,8 @@ impl GpuResources {
 }
 
 impl ProfileMap {
-    pub fn state(&self) -> u8 {
-        self.state.load(Ordering::Acquire)
-    }
-
-    pub fn read(self) -> Vec<(String, f64)> {
+    pub async fn read(self) -> Option<Vec<(String, f64)>> {
+        self.completion.await.ok()?.ok()?;
         let bytes = self.readback.get_mapped_range(..);
         let values = bytes
             .chunks_exact(8)
@@ -463,7 +459,7 @@ impl ProfileMap {
             .collect();
         drop(bytes);
         self.readback.unmap();
-        profile
+        Some(profile)
     }
 }
 

core/renderer/render.rs

@@ -10,26 +10,13 @@ use crate::render_data::RenderData;
 use crate::store::{Loadout, Store};
 
 #[wasm_bindgen(inline_js = r#"
-const channel = new MessageChannel();
-const ready = [];
-channel.port1.onmessage = () => ready.shift()?.();
-
 export function nextFrame() {
     return new Promise(resolve => requestAnimationFrame(resolve));
 }
-
-export function yieldTask() {
-    return new Promise(resolve => {
-        ready.push(resolve);
-        channel.port2.postMessage(0);
-    });
-}
 "#)]
 extern "C" {
     #[wasm_bindgen(js_name = nextFrame)]
     async fn next_frame();
-    #[wasm_bindgen(js_name = yieldTask)]
-    async fn yield_task();
 }
 
 pub struct RenderLoop {
@@ -42,7 +29,8 @@ pub struct RenderLoop {
     profiling: Cell<bool>,
     profile_pending: Cell<bool>,
     profile_after: Cell<f64>,
-    profile: RefCell<Option<String>>,
+    profile_callback: RefCell<Option<js_sys::Function>>,
+    profile_generation: Cell<u32>,
 }
 
 impl RenderLoop {
@@ -57,7 +45,8 @@ impl RenderLoop {
             profiling: Cell::new(false),
             profile_pending: Cell::new(false),
             profile_after: Cell::new(0.0),
-            profile: RefCell::new(None),
+            profile_callback: RefCell::new(None),
+            profile_generation: Cell::new(0),
         }
     }
 
@@ -81,13 +70,12 @@ impl RenderLoop {
     pub fn set_profiling(&self, enabled: bool) {
         self.profiling.set(enabled);
         self.profile_after.set(0.0);
-        if !enabled {
-            self.profile.borrow_mut().take();
-        }
+        self.profile_generation
+            .set(self.profile_generation.get().wrapping_add(1));
     }
 
-    pub fn take_profile(&self) -> Option<String> {
-        self.profile.borrow_mut().take()
+    pub fn on_profile(&self, callback: js_sys::Function) {
+        *self.profile_callback.borrow_mut() = Some(callback);
     }
 
     pub fn start(
@@ -152,15 +140,12 @@ impl RenderLoop {
                 };
                 if let Some((profile, adapter, width, height)) = submission {
                     control.profile_pending.set(true);
-                    control.profile_after.set(started + 250.0);
+                    control.profile_after.set(started + 100.0);
+                    let generation = control.profile_generation.get();
                     let control = control.clone();
                     wasm_bindgen_futures::spawn_local(async move {
-                        while profile.state() == 0 {
-                            yield_task().await;
-                        }
-                        if profile.state() == 1 {
-                            let passes = profile
-                                .read()
+                        if let Some(passes) = profile.read().await {
+                            let passes = passes
                                 .into_iter()
                                 .map(|(name, milliseconds)| {
                                     serde_json::json!({
@@ -173,18 +158,22 @@ impl RenderLoop {
                                 .iter()
                                 .filter_map(|pass| pass["milliseconds"].as_f64())
                                 .sum::<f64>();
-                            if control.profiling.get() {
-                                *control.profile.borrow_mut() = Some(
-                                    serde_json::json!({
-                                        "frame": frame,
-                                        "milliseconds": milliseconds,
-                                        "readbackMilliseconds": js_sys::Date::now() - started,
-                                        "adapter": adapter,
-                                        "canvas": { "width": width, "height": height },
-                                        "passes": passes,
-                                    })
-                                    .to_string(),
-                                );
+                            if control.profiling.get()
+                                && control.profile_generation.get() == generation
+                            {
+                                let sample = serde_json::json!({
+                                    "frame": frame,
+                                    "milliseconds": milliseconds,
+                                    "readbackMilliseconds": js_sys::Date::now() - started,
+                                    "adapter": adapter,
+                                    "canvas": { "width": width, "height": height },
+                                    "passes": passes,
+                                })
+                                .to_string();
+                                if let Some(callback) = control.profile_callback.borrow().as_ref() {
+                                    let _ = callback
+                                        .call1(&wasm_bindgen::JsValue::NULL, &sample.into());
+                                }
                             }
                         }
                         control.profile_pending.set(false);

core/worker.js

@@ -1,7 +1,8 @@
 import initWasm, { Core } from "./pkg/yawn_core.js";
+import { SharedProfile } from "./profile.js";
 
 let core;
-let profileTimer;
+let profile;
 // While a relocation is staged, every request except its commit waits here.
 let awaitingCommit = false;
 const queue = [];
@@ -15,6 +16,8 @@ async function handle(message) {
       if (core || !(message.canvas instanceof OffscreenCanvas)) fail("INIT");
       const wasm = await initWasm();
       core = new Core(message.arenaBytes);
+      profile = new SharedProfile(message.profileBuffer);
+      core.on_profile(json => profile.write(json));
       await core.initialize(message.canvas);
       const buffer = wasm.memory.buffer;
       if (!(buffer instanceof SharedArrayBuffer)) fail("WASM_MEMORY_NOT_SHARED");
@@ -70,11 +73,7 @@ async function handle(message) {
           break;
         case "set-profiler":
           result = core.set_profiler(Boolean(message.enabled));
-          clearInterval(profileTimer);
-          profileTimer = result && message.enabled ? setInterval(() => {
-            const stats = core.take_profile();
-            if (stats) postMessage({ type: "profile", stats: JSON.parse(stats) });
-          }, 250) : undefined;
+          profile.clear();
           break;
         default:
           fail("MESSAGE");