@driftengine/core 3.63.0 → 4.1.0

package/src/render/shaders/network.wgsl.ts

@@ -0,0 +1,238 @@
+/**
+ * The one network evaluator the device has, in single or half precision.
+ *
+ * **Every network the engine evaluates on a device goes through this text.** A DriftTexture's
+ * decode interpreter reads its networks from uniform tables; reconstruction's refinement tier and
+ * capture's priors read theirs from storage buffers, in half precision where the device has it.
+ * Those are different bindings and the same arithmetic, so the arithmetic is written once and the
+ * weights are reached through a function the including module defines — the arrangement
+ * `shadow.wgsl.ts` uses for its map. A second evaluator would be a second set of numerical
+ * conventions to keep in step, which `@driftengine/texture`'s `inference.ts` names as the thing this
+ * engine refuses.
+ *
+ * **The reference is `evalNetwork`, and in half precision `evalNetworkHalf`**, operation for
+ * operation: the weighted inputs summed in index order, then the bias, then the rectifier on hidden
+ * layers and nothing on the output. The weights are laid out as `NNET` stores them — each layer's
+ * matrix row by row by output, then its biases.
+ *
+ * **Half precision needs `enable f16;` at the top of the including module** and a device created
+ * with `shader-f16`. It is chosen per network by its consumer, and `activationBound` is what that
+ * consumer asks first: past 65,504 the format has no finite value.
+ *
+ * `scripts/gpu-parity.mjs` runs both instantiations against the references, and
+ * `scripts/wgsl-handwritten.test.mjs` has `naga` validate them without a device.
+ */
+
+export interface NetworkWgslOptions {
+  /** The prefix every function this defines carries, so a module can include more than one. */
+  readonly name: string;
+  /** The scalar the network computes in. `f16` needs `enable f16;` in the including module. */
+  readonly scalar: 'f32' | 'f16';
+  /** The widest layer, which is the length of the function-scope arrays the layers pass through. */
+  readonly width: number;
+  /** The most hidden layers a network may have here. */
+  readonly hidden: number;
+}
+
+/**
+ * The evaluator, as text.
+ *
+ * Defines `fn <name>Eval(first, inputs, hidden, widths, outputs, input) -> array<S, width>`:
+ * `first` is where the network's weights start, `widths` its hidden widths, and `input` its first
+ * `inputs` values. The including module defines `fn <name>Weight(index: u32) -> S`.
+ *
+ * **Every count is clamped to what the arrays hold.** A loader refuses a network wider or deeper
+ * than this was instantiated for; the clamps are what keep a malformed one from indexing past a
+ * function-scope array, which the language would otherwise clamp silently and differently.
+ */
+export function networkWgsl(options: NetworkWgslOptions): string {
+  const { name, scalar: s, width: w, hidden: h } = options;
+  return /* wgsl */ `
+fn ${name}Eval(
+  first: u32,
+  inputs: u32,
+  hidden: u32,
+  widths: array<u32, ${h}>,
+  outputs: u32,
+  input: array<${s}, ${w}>,
+) -> array<${s}, ${w}> {
+  var current = input;
+  var next: array<${s}, ${w}>;
+  var at = first;
+  var previous = min(inputs, ${w}u);
+  let layers = min(hidden, ${h}u);
+  for (var layer = 0u; layer < layers; layer = layer + 1u) {
+    let width = min(widths[layer], ${w}u);
+    for (var o = 0u; o < width; o = o + 1u) {
+      var total = ${s}(0.0);
+      for (var i = 0u; i < previous; i = i + 1u) {
+        total = total + current[i] * ${name}Weight(at + o * previous + i);
+      }
+      /* Rectified linear on hidden layers and nothing on the output, as inference.ts insists. */
+      next[o] = max(total + ${name}Weight(at + width * previous + o), ${s}(0.0));
+    }
+    at = at + previous * width + width;
+    previous = width;
+    current = next;
+  }
+  var result: array<${s}, ${w}>;
+  let count = min(outputs, ${w}u);
+  for (var o = 0u; o < count; o = o + 1u) {
+    var total = ${s}(0.0);
+    for (var i = 0u; i < previous; i = i + 1u) {
+      total = total + current[i] * ${name}Weight(at + o * previous + i);
+    }
+    result[o] = total + ${name}Weight(at + count * previous + o);
+  }
+  return result;
+}
+`;
+}
+
+/** A network's shape, which the fixed form bakes in. The same shape `evalNetwork` takes. */
+export interface FixedNetworkShape {
+  readonly inputs: number;
+  readonly hidden: readonly number[];
+  readonly outputs: number;
+}
+
+export interface FixedNetworkWgslOptions {
+  /** The prefix every function this defines carries, as for the variable form. */
+  readonly name: string;
+  readonly scalar: 'f32' | 'f16';
+  readonly shape: FixedNetworkShape;
+  /** Where the network's weights start in what `<name>Weight` reads. */
+  readonly first?: number;
+}
+
+/**
+ * The same evaluator with its shape fixed when the pipeline is built, for a consumer whose network
+ * is known then — reconstruction's refinement tier, capture's priors.
+ *
+ * **Nine to fourteen times faster, measured**, 2026-09-19 on an RX 9070 XT, once per pixel of a
+ * 1280×720 frame: a network of 11-8-3 in 0.031 ms against the variable form's 0.286, and 11-16-16-3
+ * in 0.076 against 1.09 (`tools/recon-train/budget.ts`). The variable form indexes function-scope
+ * arrays with counts it reads at run time, which a compiler cannot keep in registers; with every
+ * count a constant the loops unroll and the arrays are registers. The DriftTexture interpreter keeps
+ * the variable form, because its networks change per material and its pipeline does not.
+ *
+ * **The arithmetic is the variable form's, in the same order** — the weighted inputs summed in index
+ * order, then the bias, then the rectifier on hidden layers and nothing on the output, over weights
+ * laid out as `NNET` stores them — so `evalNetwork` and `evalNetworkHalf` are its references too,
+ * and `gpu-parity.mjs` holds it to them the same way.
+ *
+ * Defines `fn <name>Eval(input: array<S, inputs>) -> array<S, outputs>`. The including module
+ * defines `fn <name>Weight(index: u32) -> S`, as for the variable form.
+ */
+export function networkFixedWgsl(options: FixedNetworkWgslOptions): string {
+  const { name, scalar: s, shape } = options;
+  const layers = [shape.inputs, ...shape.hidden, shape.outputs];
+  let at = options.first ?? 0;
+  let body = '';
+  let previous = 'input';
+  for (let layer = 1; layer < layers.length; layer += 1) {
+    const width = layers[layer] as number;
+    const fanIn = layers[layer - 1] as number;
+    const output = layer === layers.length - 1;
+    const target = output ? 'result' : `layer${layer}`;
+    const bias = at + width * fanIn;
+    const value = output
+      ? `total + ${name}Weight(${bias}u + o)`
+      : `max(total + ${name}Weight(${bias}u + o), ${s}(0.0))`;
+    body += `
+  var ${target}: array<${s}, ${width}>;
+  for (var o = 0u; o < ${width}u; o = o + 1u) {
+    var total = ${s}(0.0);
+    for (var i = 0u; i < ${fanIn}u; i = i + 1u) {
+      total = total + ${previous}[i] * ${name}Weight(${at}u + o * ${fanIn}u + i);
+    }
+    ${target}[o] = ${value};
+  }`;
+    at = bias + width;
+    previous = target;
+  }
+  return /* wgsl */ `
+fn ${name}Eval(input: array<${s}, ${shape.inputs}>) -> array<${s}, ${shape.outputs}> {${body}
+  return result;
+}
+`;
+}
+
+/**
+ * An entry point that evaluates one fixed-shape network over many inputs, so a harness can hold
+ * the fixed form to `evalNetwork` and `evalNetworkHalf` as it holds the variable one. Inputs and
+ * outputs cross as `f32`, as the variable form's parity entry point has them.
+ */
+export function networkFixedParityWgsl(scalar: 'f32' | 'f16', shape: FixedNetworkShape): string {
+  return /* wgsl */ `${scalar === 'f16' ? 'enable f16;\n' : ''}
+@group(0) @binding(0) var<storage, read> weights: array<${scalar}>;
+@group(0) @binding(1) var<storage, read> inputs: array<f32>;
+@group(0) @binding(2) var<storage, read_write> outputs: array<f32>;
+
+fn parityWeight(index: u32) -> ${scalar} {
+  return weights[index];
+}
+${networkFixedWgsl({ name: 'parity', scalar, shape })}
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+  let count = arrayLength(&inputs) / ${shape.inputs}u;
+  if (id.x >= count) { return; }
+  var input: array<${scalar}, ${shape.inputs}>;
+  for (var i = 0u; i < ${shape.inputs}u; i = i + 1u) {
+    input[i] = ${scalar}(inputs[id.x * ${shape.inputs}u + i]);
+  }
+  let result = parityEval(input);
+  for (var o = 0u; o < ${shape.outputs}u; o = o + 1u) {
+    outputs[id.x * ${shape.outputs}u + o] = f32(result[o]);
+  }
+}
+`;
+}
+
+/** The widest network the parity entry points evaluate. */
+export const NETWORK_PARITY_WIDTH = 16;
+/** The deepest. */
+export const NETWORK_PARITY_HIDDEN = 4;
+/** Unsigned words per case: first, inputs, hidden count, outputs, then the hidden widths. */
+export const NETWORK_CASE_WORDS = 4 + NETWORK_PARITY_HIDDEN;
+
+/**
+ * An entry point that evaluates one network per invocation from storage buffers, so a harness can
+ * compare the device with `evalNetwork` or `evalNetworkHalf`.
+ *
+ * Inputs and outputs cross as `f32` whatever the network computes in: a half-precision network's
+ * inputs are converted on the way in, as `evalNetworkHalf` rounds them, and its outputs widen on
+ * the way out, which loses nothing.
+ */
+export function networkParityWgsl(scalar: 'f32' | 'f16'): string {
+  const w = NETWORK_PARITY_WIDTH;
+  return /* wgsl */ `${scalar === 'f16' ? 'enable f16;\n' : ''}
+@group(0) @binding(0) var<storage, read> cases: array<u32>;
+@group(0) @binding(1) var<storage, read> weights: array<${scalar}>;
+@group(0) @binding(2) var<storage, read> inputs: array<f32>;
+@group(0) @binding(3) var<storage, read_write> outputs: array<f32>;
+
+fn parityWeight(index: u32) -> ${scalar} {
+  return weights[index];
+}
+${networkWgsl({ name: 'parity', scalar, width: w, hidden: NETWORK_PARITY_HIDDEN })}
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+  let count = arrayLength(&cases) / ${NETWORK_CASE_WORDS}u;
+  if (id.x >= count) { return; }
+  let at = id.x * ${NETWORK_CASE_WORDS}u;
+  var widths: array<u32, ${NETWORK_PARITY_HIDDEN}>;
+  for (var layer = 0u; layer < ${NETWORK_PARITY_HIDDEN}u; layer = layer + 1u) {
+    widths[layer] = cases[at + 4u + layer];
+  }
+  var input: array<${scalar}, ${w}>;
+  for (var i = 0u; i < ${w}u; i = i + 1u) {
+    input[i] = ${scalar}(inputs[id.x * ${w}u + i]);
+  }
+  let result = parityEval(cases[at], cases[at + 1u], cases[at + 2u], widths, cases[at + 3u], input);
+  for (var o = 0u; o < ${w}u; o = o + 1u) {
+    outputs[id.x * ${w}u + o] = f32(result[o]);
+  }
+}
+`;
+}

package/src/render/shaders/recon/motion.wgsl.ts

@@ -0,0 +1,101 @@
+/**
+ * Where a drawn surface was last frame, written into the reconstruction's motion target.
+ *
+ * **A pass of its own rather than a second output on every material.** The alternative is a second
+ * colour attachment on every generated fragment stage, which `ARCHITECTURE.md` §1 prices at 196,910
+ * gzipped bytes for one flag — and it would be paid by every consumer, including the ones who never
+ * reconstruct anything. This draws only the meshes a caller said had moved, with one small shader
+ * over the position buffer alone.
+ *
+ * **What a pixel carries**: the offset from where it is now to where its surface was, in the uv the
+ * resolve measures in; the view depth it had then, which is what the disocclusion compares against
+ * last frame's depth; and a flag. `recon/resolve.ts`'s `historyWeight` is the reader, and a texel
+ * this pass did not write is a zero flag, which means *derive the camera's motion from the depth* —
+ * so a scene that says nothing about its movers is exactly the scene this pass never ran for.
+ *
+ * **The rasterisation matrix is handed over whole, and that is deliberate.** The scene is drawn by
+ * shaders `naga` generated, every one of which ends its vertex stage with `gl_Position.y = -y`, and
+ * `CLIP_CORRECTION` carries a negation to cancel it — §3 rows 56 and 57. This shader is hand
+ * written and has no such negation, so rather than reason about halves of that pair per draw, the
+ * renderer multiplies the two together once a frame and passes the product. The vertex stage then
+ * has one job: put the position exactly where the scene put it, so the depth test compares like
+ * with like.
+ *
+ * **The motion itself is measured from the *unjittered* matrices**, which is the whole of §3 row
+ * 121: the history stands for a picture of the scene as it would have been drawn with no jitter, so
+ * a motion vector that carried the jitter difference would move every sample by a fraction of a
+ * pixel in a pattern that repeats with the jitter period — shimmer, diagnosed as anything but its
+ * cause.
+ */
+
+/** Floats in the per-frame block: three matrices. */
+export const MOTION_FRAME_FLOATS = 48;
+
+/**
+ * Bytes a draw's own block takes.
+ *
+ * Two matrices is 128, and a dynamic uniform offset must be a multiple of 256 — the same alignment
+ * §3 row 50's neighbours are about. So each draw takes a whole 256 and half of it is padding, which
+ * is the cheaper half of the trade against a bind group per draw inside the frame loop.
+ */
+export const MOTION_DRAW_STRIDE = 256;
+
+/** Floats a draw actually writes inside that stride. */
+export const MOTION_DRAW_FLOATS = 32;
+
+/** The motion target's format. Half float holds a uv to a hundredth of a pixel and a view depth. */
+export const MOTION_TARGET_FORMAT: GPUTextureFormat = 'rgba16float';
+
+export const RECON_MOTION_WGSL = /* wgsl */ `
+struct MotionFrame {
+  /* The scene's own clip transform, with the generated shaders' y negation already folded in. */
+  raster: mat4x4<f32>,
+  /* This frame's and last frame's view-projections, unjittered, in the matrices' own convention. */
+  viewProj: mat4x4<f32>,
+  previousViewProj: mat4x4<f32>,
+}
+
+struct MotionDraw {
+  model: mat4x4<f32>,
+  previousModel: mat4x4<f32>,
+}
+
+@group(0) @binding(0) var<uniform> frame: MotionFrame;
+@group(0) @binding(1) var<uniform> draw: MotionDraw;
+
+struct Varyings {
+  @builtin(position) clip: vec4<f32>,
+  /* Clip positions rather than uv: an interpolator divides by w for us, so dividing these by their
+     own interpolated w in the fragment stage is the perspective-correct answer at every pixel. */
+  @location(0) now: vec4<f32>,
+  @location(1) was: vec4<f32>,
+}
+
+@vertex
+fn motionVert(@location(0) position: vec3<f32>) -> Varyings {
+  let world = draw.model * vec4<f32>(position, 1.0);
+  let wasWorld = draw.previousModel * vec4<f32>(position, 1.0);
+  var out: Varyings;
+  out.clip = frame.raster * world;
+  out.now = frame.viewProj * world;
+  out.was = frame.previousViewProj * wasWorld;
+  return out;
+}
+
+/* Normalised device coordinates to the uv the resolve measures in: v increases upward, as the
+   references' images do, and the accessor that reads this target is what flips the texture's rows. */
+fn motionUv(clip: vec4<f32>) -> vec2<f32> {
+  return clip.xy / clip.w * 0.5 + vec2<f32>(0.5);
+}
+
+@fragment
+fn motionFrag(in: Varyings) -> @location(0) vec4<f32> {
+  /*
+   * A surface behind last frame's eye has a w that is not positive, and there is no uv for it. The
+   * flag stays set: the disocclusion refuses a depth that is not positive, so the resolve drops the
+   * history rather than reading it from wherever a divide by a negative w would have pointed.
+   */
+  let offset = select(vec2<f32>(0.0), motionUv(in.was) - motionUv(in.now), in.was.w > 0.0);
+  return vec4<f32>(offset, in.was.w, 1.0);
+}
+`;