@woosh/meep-engine 3.11.3 → 3.13.0

package/src/shade/renderer/postprocess/nss/shader_nss_postprocess.js

@@ -1,292 +1,291 @@
-import { CodeChunk } from "../../shader/compiler/CodeChunk.js";
-import { ComputeShader } from "../../shader/ComputeShader.js";
-import { ShaderResourceSetDescriptor } from "../../shader/resource/ShaderResourceSetDescriptor.js";
-import { SAMPLER_LINEAR_UPSCALE } from "../../texture/SAMPLER_LINEAR_UPSCALE.js";
-import { NSS_WORKGROUP_SIZE } from "./NSS_CONSTANTS.js";
-import { NSS_SETTINGS_STRUCT } from "./NSS_SETTINGS_STRUCT.js";
-
-/**
- * NSS postprocess pass. Runs at output-resolution (one thread per output
- * pixel). Reads:
- *   - the per-render-pixel KPN columns (z=2..5 of the network output 3D
- *     texture, in REVERSE order — slot 5 is column 0). Each column has
- *     4 channels = the 4 row weights for that column.
- *   - the per-render-pixel ThetaAlpha (z=1): theta in .x (history blend
- *     weight, [0, 1]), alpha in .y (accumulation rate, scaled to
- *     [0.05, 0.4])
- *   - the jittered scene colour at render-res (used for the spatial taps)
- *   - the warped history at output-res
- *   - the nearest-depth offset from preprocess (used to dilate velocity)
- *
- * Per-output-pixel logic mirrors ARM's reference (UE5.5 NSSPostprocess.usf):
- *   1. Solve which render-res pixel feeds each of the 4x4 KPN taps.
- *   2. Apply the KPN to those taps, normalising by the weight sum.
- *   3. Warp the history (bilinear) using velocity dilated by the nearest
- *      depth offset.
- *   4. Blend the KPN colour with warped history using theta.
- *   5. In tonemapped (Karis) space, accumulate the center tap with alpha.
- *   6. Inverse tonemap, divide by EXPOSURE, write.
- *
- * @see ARM reference: UE5.5/Shaders/Private/NSSPostprocess.usf and
- *      NSSCommon.ush in github.com/arm/neural-graphics-for-unreal
- */
-
-const resources = new ShaderResourceSetDescriptor();
-resources.createGroup()
-    .addSampler("sLinear", "filtering", SAMPLER_LINEAR_UPSCALE)
-    // ── KPN coefficients + ThetaAlpha + FB (network output, 3D, depth=6) ──
-    .addTexture("tNetworkOut", "float", "3d")
-    // ── render-res inputs ──
-    .addTexture("tColour", "float")          // jittered scene colour
-    .addTexture("tVelocity", "float")        // pixel-space motion vectors
-    .addTexture("tNearestOffset", "float")   // packed (-1, 0, 1)² offset
-    // ── output-res input ──
-    .addTexture("tHistory", "float")         // previous frame upscaled colour
-    // ── settings ──
-    .addUniform("settings", NSS_SETTINGS_STRUCT)
-    // ── outputs ──
-    .addStorageTexture("tOutput", "rgba16float", "2d", "write-only")
-;
-
-
-const body = CodeChunk.from(
-    //language=WGSL
-    `
-const NSS_EXPOSURE: f32 = 7.38905609893; // e^2, matches ARM's NSSCommon.ush
-
-// Karis tonemap (max-of-channels) — matches ARM's NSSCommon.ush.
-fn nss_tonemap(rgb: vec3<f32>) -> vec3<f32> {
-    let m = max(max(rgb.r, rgb.g), rgb.b);
-    return rgb / (1.0 + m);
-}
-
-// Karis inverse — clamp to dodge the inf at m == 1.0.
-fn nss_tonemap_inv(rgb: vec3<f32>) -> vec3<f32> {
-    let c = clamp(rgb, vec3<f32>(0.0), vec3<f32>(0.99999));
-    let m = max(max(c.r, c.g), c.b);
-    return c / (1.0 - m);
-}
-
-fn nss_warped_history(prev_uv: vec2<f32>) -> vec3<f32> {
-    let on_screen = all(prev_uv >= vec2<f32>(0.0)) && all(prev_uv <= vec2<f32>(1.0));
-    if (!on_screen) {
-        return vec3<f32>(0.0);
-    }
-    // ARM uses Catmull-Rom by default for the history warp; bilinear is the
-    // simpler approximation we keep here. (Catmull-Rom is in their
-    // CatmullRomWithDering helper — TODO: port if quality demands it.)
-    return max(vec3<f32>(0.0), textureSampleLevel(tHistory, sLinear, prev_uv, 0.0).rgb);
-}
-
-// Bilinear-sample a single z-slice of the network output 3D texture at a
-// render-res pixel position. Mirrors ARM's ManualBilinear4 — without this
-// every 2x2 output block in 2x mode would share the same per-render-pixel
-// theta/alpha/KPN values, leaving sharp blocky transitions. We do manual
-// 4-tap bilinear (rather than textureSampleLevel with the linear sampler)
-// so adjacent z layers — which carry semantically different head outputs —
-// don't get filtered together.
-fn nss_sample_layer_bilinear(render_pos: vec2<f32>, render_res_max: vec2<i32>, layer: i32) -> vec4<f32> {
-    let pos = render_pos - 0.5;
-    let base = floor(pos);
-    let frac = pos - base;
-    let bi = vec2<i32>(base);
-    let lo = vec2<i32>(0, 0);
-    let s00 = textureLoad(tNetworkOut, vec3<i32>(clamp(bi,                  lo, render_res_max), layer), 0);
-    let s10 = textureLoad(tNetworkOut, vec3<i32>(clamp(bi + vec2<i32>(1,0), lo, render_res_max), layer), 0);
-    let s01 = textureLoad(tNetworkOut, vec3<i32>(clamp(bi + vec2<i32>(0,1), lo, render_res_max), layer), 0);
-    let s11 = textureLoad(tNetworkOut, vec3<i32>(clamp(bi + vec2<i32>(1,1), lo, render_res_max), layer), 0);
-    return mix(mix(s00, s10, frac.x), mix(s01, s11, frac.x), frac.y);
-}
-
-// SolveInputLocation: figure out which render-res pixel maps to a given
-// output-res pixel position, given upscale ratio and jitter. Direct port
-// of ARM's helper (NSSPostprocess.usf:80). Returns -1 when no render-res
-// sample lands at this output position.
-fn nss_solve_input_location(out_idx: i32, upscale: f32, jitter: f32) -> i32 {
-    let lower = f32(out_idx) / upscale - (-jitter + 0.5);
-    let upper = f32(out_idx + 1) / upscale - (-jitter + 0.5);
-    let candidate = i32(floor(upper));
-    if (lower < f32(candidate)) {
-        return candidate;
-    }
-    return -1;
-}
-
-@compute @workgroup_size(${NSS_WORKGROUP_SIZE}, ${NSS_WORKGROUP_SIZE}, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let out_res_u = settings.output_resolution;
-    let render_res_u = settings.render_resolution;
-
-    if (any(gid.xy >= out_res_u)) {
-        return;
-    }
-
-    let out_st = vec2<i32>(gid.xy);
-    let render_res = vec2<i32>(render_res_u);
-
-    let out_res_f = vec2<f32>(out_res_u);
-    let render_res_f = vec2<f32>(render_res_u);
-    let upscale = out_res_f / render_res_f;
-
-    // Render-res pixel that hosts this output pixel's nearest network sample.
-    // (Used for the per-render-res velocity / nearest-offset taps; the
-    // network's per-pixel theta/alpha/KPN outputs are bilinear-sampled at
-    // the OUTPUT pixel's render-res position instead — see below.)
-    let render_centre = (vec2<f32>(out_st) + 0.5) / upscale;
-    let coef_st = vec2<i32>(floor(render_centre));
-    let render_res_max = render_res - vec2<i32>(1, 1);
-    let coef_st_clamped = clamp(coef_st, vec2<i32>(0, 0), render_res_max);
-
-    // ── KPN columns ──
-    // ARM's quantized model maps OutputIndices.KPNFilterCol0 = 5,
-    // KPNFilterCol1 = 4, KPNFilterCol2 = 3, KPNFilterCol3 = 2 (NSS.cpp:243).
-    // Each column has 4 channels = the 4 row weights for that column.
-    // Bilinear at render_centre (matches ARM's ManualBilinear4) — without
-    // this, every 2x2 output block in 2x mode shares one set of weights and
-    // we get blocky/jittery transitions across cell boundaries.
-    let col0 = nss_sample_layer_bilinear(render_centre, render_res_max, 5);
-    let col1 = nss_sample_layer_bilinear(render_centre, render_res_max, 4);
-    let col2 = nss_sample_layer_bilinear(render_centre, render_res_max, 3);
-    let col3 = nss_sample_layer_bilinear(render_centre, render_res_max, 2);
-    // Force tiny positive — ARM does this to dodge 0/0 in the
-    // weight-normalisation divide (NSSPostprocess.usf:118).
-    let kpn_cols = array<vec4<f32>, 4>(
-        max(col0, vec4<f32>(1e-7)),
-        max(col1, vec4<f32>(1e-7)),
-        max(col2, vec4<f32>(1e-7)),
-        max(col3, vec4<f32>(1e-7)),
-    );
-
-    // ── Apply KPN ──
-    // Outer loop is column index (matches ARM's "for x in 0..3" - we
-    // load the 4 cols above and pick by index here). Inner loop is row
-    // index. Each (col, row) tap reads the render-res sample at
-    // OutputPixelPos + (col-1, row-1) and accumulates colour * weight,
-    // with the centre tap (col=1, row=1) also captured for the
-    // alpha-blend stage.
-    var filtered_rgb = vec3<f32>(0.0);
-    var weight_sum = 0.0;
-    var center_tap_rgb = vec3<f32>(0.0);
-    var center_tap_present = 0.0;
-    for (var x = 0; x < 4; x = x + 1) {
-        let col = kpn_cols[x];
-        for (var y = 0; y < 4; y = y + 1) {
-            let weight = col[y];
-            let tap_out = out_st + vec2<i32>(x - 1, y - 1);
-
-            let in_x = nss_solve_input_location(tap_out.x, upscale.x, settings.jitter.x * settings.jitter_sign.x);
-            let in_y = nss_solve_input_location(tap_out.y, upscale.y, settings.jitter.y * settings.jitter_sign.y);
-            if (in_x < 0 || in_y < 0) {
-                continue;
-            }
-
-            let sample_st = clamp(vec2<i32>(in_x, in_y), vec2<i32>(0, 0), render_res - vec2<i32>(1, 1));
-            let sample_rgb = max(vec3<f32>(0.0), textureLoad(tColour, sample_st, 0).rgb) * NSS_EXPOSURE;
-
-            filtered_rgb += sample_rgb * weight;
-            weight_sum += weight;
-
-            if (x == 1 && y == 1) {
-                center_tap_rgb = sample_rgb;
-                center_tap_present = 1.0;
-            }
-        }
-    }
-    let kpn_rgb = filtered_rgb / max(weight_sum, 1e-5);
-
-    // ── Warp history ──
-    // Decode the nearest-depth offset to {-1, 0, 1}², use it to dilate
-    // the velocity tap to the depth-nearest sample (same trick as
-    // preprocess).
-    let nearest_packed = textureLoad(tNearestOffset, coef_st_clamped, 0).rg;
-    let nearest_offset = vec2<i32>(round(nearest_packed * 2.0 - 1.0));
-
-    var velocity = textureLoad(tVelocity, coef_st_clamped + nearest_offset, 0).rg;
-    // Motion zero-clamp (ARM NSSPostprocess.usf:215). Sub-pixel velocity on
-    // near-static frames drifts prev_uv by fractions of a pixel each frame
-    // and bilinear-resamples the history at slightly different sub-pixel
-    // positions — visible as flicker even when nothing in the scene is
-    // moving. ARM's threshold is 0.1 input pixels.
-    if (length(velocity) <= 0.1) {
-        velocity = vec2<f32>(0.0);
-    }
-    let velocity_out = velocity / (render_res_f / out_res_f); // render-res pixels → output-res pixels
-
-    let prev_out_st = vec2<f32>(out_st) + 0.5 - velocity_out;
-    let prev_uv = prev_out_st / out_res_f;
-
-    let history_rgb = nss_warped_history(prev_uv) * NSS_EXPOSURE;
-    let on_screen_f = select(0.0, 1.0, all(prev_uv >= vec2<f32>(0.0)) && all(prev_uv <= vec2<f32>(1.0)));
-
-    // ── ThetaAlpha (head slot 1) ──
-    // ARM's quantized model maps OutputIndices.ThetaAlpha = 1
-    // (NSS.cpp:242). Channel x = theta (history blend weight, [0, 1] from
-    // sigmoid LUT), channel y = alpha (raw [0, 1], remapped to
-    // [0.05, 0.4] per NSSPostprocess.usf:248).
-    // Bilinear at render_centre — same reasoning as the KPN cols above.
-    let ta = nss_sample_layer_bilinear(render_centre, render_res_max, 1);
-    let theta_network = saturate(ta.x);
-    let theta_raw = select(theta_network, saturate(settings.theta_override), settings.theta_override >= 0.0);
-    let theta = theta_raw * saturate(settings.history_validity) * on_screen_f;
-    let alpha = (0.35 * saturate(ta.y) + 0.05) * settings.alpha_blend_scale;
-
-    // ── Stage 1: blend KPN with warped history ──
-    var blended = mix(kpn_rgb, history_rgb, theta);
-
-    // ── Stage 2: accumulate centre tap in tonemapped space ──
-    blended = mix(
-        nss_tonemap(blended),
-        nss_tonemap(center_tap_rgb),
-        alpha * center_tap_present,
-    );
-
-    var final_rgb = max(vec3<f32>(0.0), nss_tonemap_inv(blended)) / NSS_EXPOSURE;
-
-    // Debug visualisation. Useful for SEEING per-pixel network output
-    // stability: a stable scene should produce a stable, spatially-varying
-    // grayscale. Uniform gray = network has lost spatial discrimination
-    // (only the bias survives). Flickering = network output itself is the
-    // noise source.
-    //
-    //   1 = theta (slot 1.x sigmoid)
-    //   2 = alpha (slot 1.y sigmoid)
-    //   3 = KPN col 0 channel 0 (the (col=0, row=0) kernel weight at slot 5.x)
-    //   4 = KPN total weight sum (all 16 weights, should be ~1 for a healthy
-    //       normalised kernel)
-    //   5 = FB channel 0 (slot 0.x feedback)
-    let dv = settings.debug_view;
-    if (dv == 1u) {
-        final_rgb = vec3<f32>(theta_network);
-    } else if (dv == 2u) {
-        final_rgb = vec3<f32>(saturate(ta.y));
-    } else if (dv == 3u) {
-        final_rgb = vec3<f32>(kpn_cols[0].x);
-    } else if (dv == 4u) {
-        // Sum the 4 channels of each col, then sum across cols. Visualises
-        // how much "weight" the network is allocating in total. ARM's
-        // post-norm KPN should sum to ~1; if uniformly mid-grey here too,
-        // every head is bias-dominated.
-        let sum_c = (kpn_cols[0] + kpn_cols[1] + kpn_cols[2] + kpn_cols[3]);
-        let total = sum_c.x + sum_c.y + sum_c.z + sum_c.w;
-        final_rgb = vec3<f32>(total * 0.25); // /4 to keep it visually in [0, 1]-ish
-    } else if (dv == 5u) {
-        let fb = nss_sample_layer_bilinear(render_centre, render_res_max, 0);
-        final_rgb = vec3<f32>(fb.x);
-    }
-
-    // Alpha channel of the output stores the effective history weight, so
-    // downstream passes (e.g. sharpening, post-filtering) can dampen
-    // their effect where temporal accumulation is doing the heavy lifting.
-    textureStore(tOutput, out_st, vec4<f32>(final_rgb, theta));
-}
-    `, [
-        NSS_SETTINGS_STRUCT.declaration_chunk,
-    ]
-);
-
-export const shader_nss_postprocess = ComputeShader.from({
-    label: "NSS / postprocess",
-    body,
-    resources,
-    group_size: [NSS_WORKGROUP_SIZE, NSS_WORKGROUP_SIZE, 1],
-});
+import { CodeChunk } from "../../shader/compiler/CodeChunk.js";
+import { ComputeShader } from "../../shader/ComputeShader.js";
+import { ShaderResourceSetDescriptor } from "../../shader/resource/ShaderResourceSetDescriptor.js";
+import { SAMPLER_LINEAR_UPSCALE } from "../../texture/SAMPLER_LINEAR_UPSCALE.js";
+import { NSS_WORKGROUP_SIZE } from "./NSS_CONSTANTS.js";
+import { NSS_SETTINGS_STRUCT } from "./NSS_SETTINGS_STRUCT.js";
+
+/**
+ * NSS postprocess pass. Runs at output-resolution (one thread per output
+ * pixel). Reads:
+ *   - the per-render-pixel KPN columns (z=2..5 of the network output 3D
+ *     texture, in REVERSE order — slot 5 is column 0). Each column has
+ *     4 channels = the 4 row weights for that column.
+ *   - the per-render-pixel ThetaAlpha (z=1): theta in .x (history blend
+ *     weight, [0, 1]), alpha in .y (accumulation rate, scaled to
+ *     [0.05, 0.4])
+ *   - the jittered scene colour at render-res (used for the spatial taps)
+ *   - the warped history at output-res
+ *   - the nearest-depth offset from preprocess (used to dilate velocity)
+ *
+ * Per-output-pixel logic mirrors ARM's reference (UE5.5 NSSPostprocess.usf):
+ *   1. Solve which render-res pixel feeds each of the 4x4 KPN taps.
+ *   2. Apply the KPN to those taps, normalising by the weight sum.
+ *   3. Warp the history (bilinear) using velocity dilated by the nearest
+ *      depth offset.
+ *   4. Blend the KPN colour with warped history using theta.
+ *   5. In tonemapped (Karis) space, accumulate the center tap with alpha.
+ *   6. Inverse tonemap, divide by EXPOSURE, write.
+ *
+ * @see ARM reference: UE5.5/Shaders/Private/NSSPostprocess.usf and
+ *      NSSCommon.ush in github.com/arm/neural-graphics-for-unreal
+ */
+
+const resources = new ShaderResourceSetDescriptor();
+resources.createGroup()
+    .addSampler("sLinear", "filtering", SAMPLER_LINEAR_UPSCALE)
+    // ── KPN coefficients + ThetaAlpha + FB (network output, 3D, depth=6) ──
+    .addTexture("tNetworkOut", "float", "3d")
+    // ── render-res inputs ──
+    .addTexture("tColour", "float")          // jittered scene colour
+    .addTexture("tVelocity", "float")        // pixel-space motion vectors
+    .addTexture("tNearestOffset", "float")   // packed (-1, 0, 1)² offset
+    // ── output-res input ──
+    .addTexture("tHistory", "float")         // previous frame upscaled colour
+    // ── settings ──
+    .addUniform("settings", NSS_SETTINGS_STRUCT)
+    // ── outputs ──
+    .addStorageTexture("tOutput", "rgba16float", "2d", "write-only")
+;
+
+
+const body = CodeChunk.from(
+    //language=WGSL
+    `
+const NSS_EXPOSURE: f32 = 7.38905609893; // e^2, matches ARM's NSSCommon.ush
+
+// Karis tonemap (max-of-channels) — matches ARM's NSSCommon.ush.
+fn nss_tonemap(rgb: vec3<f32>) -> vec3<f32> {
+    let m = max(max(rgb.r, rgb.g), rgb.b);
+    return rgb / (1.0 + m);
+}
+
+// Karis inverse — clamp to dodge the inf at m == 1.0.
+fn nss_tonemap_inv(rgb: vec3<f32>) -> vec3<f32> {
+    let c = clamp(rgb, vec3<f32>(0.0), vec3<f32>(0.99999));
+    let m = max(max(c.r, c.g), c.b);
+    return c / (1.0 - m);
+}
+
+fn nss_warped_history(prev_uv: vec2<f32>) -> vec3<f32> {
+    let on_screen = all(prev_uv >= vec2<f32>(0.0)) && all(prev_uv <= vec2<f32>(1.0));
+    if (!on_screen) {
+        return vec3<f32>(0.0);
+    }
+    // ARM uses Catmull-Rom by default for the history warp; bilinear is the
+    // simpler approximation we keep here. (Catmull-Rom is in their
+    // CatmullRomWithDering helper — TODO: port if quality demands it.)
+    return max(vec3<f32>(0.0), textureSampleLevel(tHistory, sLinear, prev_uv, 0.0).rgb);
+}
+
+// Bilinear-sample a single z-slice of the network output 3D texture at a
+// render-res pixel position. Mirrors ARM's ManualBilinear4 — without this
+// every 2x2 output block in 2x mode would share the same per-render-pixel
+// theta/alpha/KPN values, leaving sharp blocky transitions. We do manual
+// 4-tap bilinear (rather than textureSampleLevel with the linear sampler)
+// so adjacent z layers — which carry semantically different head outputs —
+// don't get filtered together.
+fn nss_sample_layer_bilinear(render_pos: vec2<f32>, render_res_max: vec2<i32>, layer: i32) -> vec4<f32> {
+    let pos = render_pos - 0.5;
+    let base = floor(pos);
+    let frac = pos - base;
+    let bi = vec2<i32>(base);
+    let lo = vec2<i32>(0, 0);
+    let s00 = textureLoad(tNetworkOut, vec3<i32>(clamp(bi,                  lo, render_res_max), layer), 0);
+    let s10 = textureLoad(tNetworkOut, vec3<i32>(clamp(bi + vec2<i32>(1,0), lo, render_res_max), layer), 0);
+    let s01 = textureLoad(tNetworkOut, vec3<i32>(clamp(bi + vec2<i32>(0,1), lo, render_res_max), layer), 0);
+    let s11 = textureLoad(tNetworkOut, vec3<i32>(clamp(bi + vec2<i32>(1,1), lo, render_res_max), layer), 0);
+    return mix(mix(s00, s10, frac.x), mix(s01, s11, frac.x), frac.y);
+}
+
+// SolveInputLocation: figure out which render-res pixel maps to a given
+// output-res pixel position, given upscale ratio and jitter. Direct port
+// of ARM's helper (NSSPostprocess.usf:80). Returns -1 when no render-res
+// sample lands at this output position.
+fn nss_solve_input_location(out_idx: i32, upscale: f32, jitter: f32) -> i32 {
+    let lower = f32(out_idx) / upscale - (-jitter + 0.5);
+    let upper = f32(out_idx + 1) / upscale - (-jitter + 0.5);
+    let candidate = i32(floor(upper));
+    if (lower < f32(candidate)) {
+        return candidate;
+    }
+    return -1;
+}
+
+@compute @workgroup_size(${NSS_WORKGROUP_SIZE}, ${NSS_WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let out_res_u = settings.output_resolution;
+    let render_res_u = settings.render_resolution;
+
+    if (any(gid.xy >= out_res_u)) {
+        return;
+    }
+
+    let out_st = vec2<i32>(gid.xy);
+    let render_res = vec2<i32>(render_res_u);
+
+    let out_res_f = vec2<f32>(out_res_u);
+    let render_res_f = vec2<f32>(render_res_u);
+    let upscale = out_res_f / render_res_f;
+
+    // Render-res pixel that hosts this output pixel's nearest network sample.
+    // (Used for the per-render-res velocity / nearest-offset taps; the
+    // network's per-pixel theta/alpha/KPN outputs are bilinear-sampled at
+    // the OUTPUT pixel's render-res position instead — see below.)
+    let render_centre = (vec2<f32>(out_st) + 0.5) / upscale;
+    let coef_st = vec2<i32>(floor(render_centre));
+    let render_res_max = render_res - vec2<i32>(1, 1);
+    let coef_st_clamped = clamp(coef_st, vec2<i32>(0, 0), render_res_max);
+
+    // ── KPN columns ──
+    // ARM's quantized model maps OutputIndices.KPNFilterCol0 = 5,
+    // KPNFilterCol1 = 4, KPNFilterCol2 = 3, KPNFilterCol3 = 2 (NSS.cpp:243).
+    // Each column has 4 channels = the 4 row weights for that column.
+    // Bilinear at render_centre (matches ARM's ManualBilinear4) — without
+    // this, every 2x2 output block in 2x mode shares one set of weights and
+    // we get blocky/jittery transitions across cell boundaries.
+    let col0 = nss_sample_layer_bilinear(render_centre, render_res_max, 5);
+    let col1 = nss_sample_layer_bilinear(render_centre, render_res_max, 4);
+    let col2 = nss_sample_layer_bilinear(render_centre, render_res_max, 3);
+    let col3 = nss_sample_layer_bilinear(render_centre, render_res_max, 2);
+    // Force tiny positive — ARM does this to dodge 0/0 in the
+    // weight-normalisation divide (NSSPostprocess.usf:118).
+    let kpn_cols = array<vec4<f32>, 4>(
+        max(col0, vec4<f32>(1e-7)),
+        max(col1, vec4<f32>(1e-7)),
+        max(col2, vec4<f32>(1e-7)),
+        max(col3, vec4<f32>(1e-7)),
+    );
+
+    // ── Apply KPN ──
+    // Outer loop is column index (matches ARM's "for x in 0..3" - we
+    // load the 4 cols above and pick by index here). Inner loop is row
+    // index. Each (col, row) tap reads the render-res sample at
+    // OutputPixelPos + (col-1, row-1) and accumulates colour * weight,
+    // with the centre tap (col=1, row=1) also captured for the
+    // alpha-blend stage.
+    var filtered_rgb = vec3<f32>(0.0);
+    var weight_sum = 0.0;
+    var center_tap_rgb = vec3<f32>(0.0);
+    var center_tap_present = 0.0;
+    for (var x = 0; x < 4; x = x + 1) {
+        let col = kpn_cols[x];
+        for (var y = 0; y < 4; y = y + 1) {
+            let weight = col[y];
+            let tap_out = out_st + vec2<i32>(x - 1, y - 1);
+
+            let in_x = nss_solve_input_location(tap_out.x, upscale.x, settings.jitter.x * settings.jitter_sign.x);
+            let in_y = nss_solve_input_location(tap_out.y, upscale.y, settings.jitter.y * settings.jitter_sign.y);
+            if (in_x < 0 || in_y < 0) {
+                continue;
+            }
+
+            let sample_st = clamp(vec2<i32>(in_x, in_y), vec2<i32>(0, 0), render_res - vec2<i32>(1, 1));
+            let sample_rgb = max(vec3<f32>(0.0), textureLoad(tColour, sample_st, 0).rgb) * NSS_EXPOSURE;
+
+            filtered_rgb += sample_rgb * weight;
+            weight_sum += weight;
+
+            if (x == 1 && y == 1) {
+                center_tap_rgb = sample_rgb;
+                center_tap_present = 1.0;
+            }
+        }
+    }
+    let kpn_rgb = filtered_rgb / max(weight_sum, 1e-5);
+
+    // ── Warp history ──
+    // Decode the nearest-depth offset to {-1, 0, 1}², use it to dilate
+    // the velocity tap to the depth-nearest sample (same trick as
+    // preprocess).
+    let nearest_packed = textureLoad(tNearestOffset, coef_st_clamped, 0).rg;
+    let nearest_offset = vec2<i32>(round(nearest_packed * 2.0 - 1.0));
+
+    var velocity = textureLoad(tVelocity, coef_st_clamped + nearest_offset, 0).rg;
+    // Motion zero-clamp (ARM NSSPostprocess.usf:215). Sub-pixel velocity on
+    // near-static frames drifts prev_uv by fractions of a pixel each frame
+    // and bilinear-resamples the history at slightly different sub-pixel
+    // positions — visible as flicker even when nothing in the scene is
+    // moving. ARM's threshold is 0.1 input pixels.
+    if (length(velocity) <= 0.1) {
+        velocity = vec2<f32>(0.0);
+    }
+    let velocity_out = velocity / (render_res_f / out_res_f); // render-res pixels → output-res pixels
+
+    let prev_out_st = vec2<f32>(out_st) + 0.5 - velocity_out;
+    let prev_uv = prev_out_st / out_res_f;
+
+    let history_rgb = nss_warped_history(prev_uv) * NSS_EXPOSURE;
+    let on_screen_f = select(0.0, 1.0, all(prev_uv >= vec2<f32>(0.0)) && all(prev_uv <= vec2<f32>(1.0)));
+
+    // ── ThetaAlpha (head slot 1) ──
+    // ARM's quantized model maps OutputIndices.ThetaAlpha = 1
+    // (NSS.cpp:242). Channel x = theta (history blend weight, [0, 1] from
+    // sigmoid LUT), channel y = alpha (raw [0, 1], remapped to
+    // [0.05, 0.4] per NSSPostprocess.usf:248).
+    // Bilinear at render_centre — same reasoning as the KPN cols above.
+    let ta = nss_sample_layer_bilinear(render_centre, render_res_max, 1);
+    let theta_network = saturate(ta.x);
+    let theta_raw = select(theta_network, saturate(settings.theta_override), settings.theta_override >= 0.0);
+    let theta = theta_raw * saturate(settings.history_validity) * on_screen_f;
+    let alpha = (0.35 * saturate(ta.y) + 0.05) * settings.alpha_blend_scale;
+
+    // ── Stage 1: blend KPN with warped history ──
+    var blended = mix(kpn_rgb, history_rgb, theta);
+
+    // ── Stage 2: accumulate centre tap in tonemapped space ──
+    blended = mix(
+        nss_tonemap(blended),
+        nss_tonemap(center_tap_rgb),
+        alpha * center_tap_present,
+    );
+
+    var final_rgb = max(vec3<f32>(0.0), nss_tonemap_inv(blended)) / NSS_EXPOSURE;
+
+    // Debug visualisation. Useful for SEEING per-pixel network output
+    // stability: a stable scene should produce a stable, spatially-varying
+    // grayscale. Uniform gray = network has lost spatial discrimination
+    // (only the bias survives). Flickering = network output itself is the
+    // noise source.
+    //
+    //   1 = theta (slot 1.x sigmoid)
+    //   2 = alpha (slot 1.y sigmoid)
+    //   3 = KPN col 0 channel 0 (the (col=0, row=0) kernel weight at slot 5.x)
+    //   4 = KPN total weight sum (all 16 weights, should be ~1 for a healthy
+    //       normalised kernel)
+    //   5 = FB channel 0 (slot 0.x feedback)
+    let dv = settings.debug_view;
+    if (dv == 1u) {
+        final_rgb = vec3<f32>(theta_network);
+    } else if (dv == 2u) {
+        final_rgb = vec3<f32>(saturate(ta.y));
+    } else if (dv == 3u) {
+        final_rgb = vec3<f32>(kpn_cols[0].x);
+    } else if (dv == 4u) {
+        // Sum the 4 channels of each col, then sum across cols. Visualises
+        // how much "weight" the network is allocating in total. ARM's
+        // post-norm KPN should sum to ~1; if uniformly mid-grey here too,
+        // every head is bias-dominated.
+        let sum_c = (kpn_cols[0] + kpn_cols[1] + kpn_cols[2] + kpn_cols[3]);
+        let total = sum_c.x + sum_c.y + sum_c.z + sum_c.w;
+        final_rgb = vec3<f32>(total * 0.25); // /4 to keep it visually in [0, 1]-ish
+    } else if (dv == 5u) {
+        let fb = nss_sample_layer_bilinear(render_centre, render_res_max, 0);
+        final_rgb = vec3<f32>(fb.x);
+    }
+
+    // Alpha channel of the output stores the effective history weight, so
+    // downstream passes (e.g. sharpening, post-filtering) can dampen
+    // their effect where temporal accumulation is doing the heavy lifting.
+    textureStore(tOutput, out_st, vec4<f32>(final_rgb, theta));
+}
+    `, [
+        NSS_SETTINGS_STRUCT.declaration_chunk,
+    ]
+);
+
+export const shader_nss_postprocess = ComputeShader.from({
+    label: "NSS / postprocess",
+    body,
+    resources,
+});

package/src/shade/renderer/postprocess/nss/shader_nss_preprocess.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"shader_nss_preprocess.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/postprocess/nss/shader_nss_preprocess.js"],"names":[],"mappings":"AAiOA,kDAKG;8BArO2B,+BAA+B"}
+{"version":3,"file":"shader_nss_preprocess.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/postprocess/nss/shader_nss_preprocess.js"],"names":[],"mappings":"AAiOA,kDAIG;8BApO2B,+BAA+B"}