@woosh/meep-engine 3.11.3 → 3.13.0

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

@@ -1,226 +1,225 @@
-import { f32, u32 } from "../../shader/type/WEBGPU_WGSL_PRIMITIVE_TYPES.js";
-import { WebGPUArray } from "../../shader/type/WebGPUArray.js";
-import { CodeChunk } from "../../shader/compiler/CodeChunk.js";
-import { ComputeShader } from "../../shader/ComputeShader.js";
-import { ShaderResourceSetDescriptor } from "../../shader/resource/ShaderResourceSetDescriptor.js";
-import { NSS_WORKGROUP_SIZE } from "./NSS_CONSTANTS.js";
-import { NSS_LAYER_CONFIG_STRUCT } from "./NSS_LAYER_CONFIG_STRUCT.js";
-import { NSS_SETTINGS_STRUCT } from "./NSS_SETTINGS_STRUCT.js";
-
-/**
- * One 3x3 (or 1x1) convolution layer of the NSS network — int8 weights,
- * f32 biases, per-channel TOSA RESCALE applied before activation.
- *
- * Pipeline per output channel:
- *   acc = sum(input * unpack4x8snorm(weight))      (over ic, ky, kx)
- *   acc += bias                                     (when has_bias)
- *   acc *= rescale                                   (per-channel, from ARM's RESCALE multipliers/shifts)
- *   acc = saturate to [-1, 1]                        (mirrors int8 saturation)
- *   acc = activation(acc)                            (ReLU clamps to [0, 1]; identity passes through)
- *   textureStore(...)
- *
- * The TOSA RESCALE half is what keeps intermediate activations bounded —
- * without it the encoder's 64-channel layers blow up by 4-5 orders of
- * magnitude and downstream layers degenerate.
- *
- * Weights live in a shared packed-int8 storage buffer (`array<u32>`,
- * 4 int8 weights per u32 along the in_channels dimension). Biases and
- * rescales live in separate f32 storage buffers. `output_layer_offset`
- * lets multiple parallel output heads share one destination texture.
- *
- * Edge handling: clamp-to-edge.
- */
-const resources = new ShaderResourceSetDescriptor();
-resources.createGroup()
-    .addTexture("tInput", "float", "3d")
-    .addStorageTexture("tOutput", "rgba16float", "3d", "write-only")
-    .addStorageBuffer("weights", WebGPUArray.from(u32))
-    .addStorageBuffer("biases", WebGPUArray.from(f32))
-    .addStorageBuffer("rescales", WebGPUArray.from(f32))
-    // Per-head transfer LUTs (TOSA TABLE op). 256 int8 entries packed
-    // 4-per-u32. Only output heads (layers with `has_lut == 1`) read
-    // from this; body layers ignore the binding and use the inline
-    // ReLU activation_kind path.
-    .addStorageBuffer("luts", WebGPUArray.from(u32))
-    .addUniform("layer_config", NSS_LAYER_CONFIG_STRUCT)
-    // Shared per-frame settings — currently only network_acc_scale is read,
-    // but binding the whole struct keeps it consistent with the other passes.
-    .addUniform("settings", NSS_SETTINGS_STRUCT)
-;
-
-
-const body = CodeChunk.from(
-    //language=WGSL
-    `
-fn nss_activation(x: vec4<f32>, kind: u32) -> vec4<f32> {
-    if (kind == 1u) {
-        // ReLU on top of the rescaled+saturated input. The input is
-        // already in [-1, 1], so this gives [0, 1].
-        return max(vec4<f32>(0.0), x);
-    } else if (kind == 2u) {
-        // Leaky ReLU, alpha = 0.1
-        return select(x, 0.1 * x, x < vec4<f32>(0.0));
-    }
-    // identity — used by the output heads
-    return x;
-}
-
-fn nss_load_input_pixel(coord: vec2<i32>, group_z: u32) -> vec4<f32> {
-    let dim = textureDimensions(tInput);
-    let clamped = clamp(coord, vec2<i32>(0, 0), vec2<i32>(dim.xy) - vec2<i32>(1, 1));
-    return textureLoad(tInput, vec3<i32>(clamped, i32(group_z)), 0);
-}
-
-// Look up one entry in a layer's per-head LUT (TOSA TABLE op).
-//
-// value_fp is the un-rescaled fp accumulator scaled by the per-channel
-// rescale_factor (= 127 * mult / 2^shift) - i.e. the int8 RESCALE
-// result without the output_zp bias yet applied. output_zp is the
-// per-head signed int8 bias from TOSA RESCALE; adding it before
-// clamp+shift puts the LUT INDEX in the trained transition zone (each
-// head's LUT has almost all of its dynamic range concentrated in a
-// window centred on uint8 = output_zp + 128 - without applying it,
-// the index lands in the flat tail and the network output collapses
-// to a constant).
-//
-// LUT indexing: int8 input shifted to [0, 255]; LUT is 256 int8s
-// packed 4-per-u32. ARM uses this to encode the trained activation
-// function for the heads (sigmoid for ThetaAlpha, tanh-like S-curves
-// for the KPN columns and FB).
-//
-// Output dequantisation matches ARM's Dequantize: clamp the int8 to
-// [-127, 127] (folding -128 into -127 to dodge an off-by-one when the
-// LUT table happens to store it), then map to fp via (x + 127) / 255
-// - giving fp in [0, 1] (NOT [-1, 1] as a naive /127 would suggest).
-fn nss_apply_lut(value_fp: f32, lut_off_u32: u32, output_zp: i32) -> f32 {
-    let int8_unbiased = i32(round(value_fp * 127.0));
-    let int8_in = clamp(int8_unbiased + output_zp, -128, 127);
-    let idx = u32(int8_in + 128); // shift to [0, 255]
-    let u32_idx = idx / 4u;
-    let lane = idx % 4u;
-    // unpack4xI8 returns vec4<i32> with sign-extended int8 values.
-    let four = unpack4xI8(luts[lut_off_u32 + u32_idx]);
-    let int8_out = max(four[lane], -127);
-    return f32(int8_out + 127) / 255.0;
-}
-
-@compute @workgroup_size(${NSS_WORKGROUP_SIZE}, ${NSS_WORKGROUP_SIZE}, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let out_dim = textureDimensions(tOutput);
-    let layer_out_groups = layer_config.out_channels / 4u;
-
-    if (any(gid.xy >= out_dim.xy) || gid.z >= layer_out_groups) {
-        return;
-    }
-
-    let in_channels = layer_config.in_channels;
-    let k = layer_config.kernel_size;
-    let weights_off = layer_config.weights_offset;
-    let bias_off = layer_config.bias_offset;
-    let rescale_off = layer_config.rescale_offset;
-    let lut_off = layer_config.lut_offset;
-    let activation = layer_config.activation_kind;
-    let has_bias = layer_config.has_bias;
-    let has_lut = layer_config.has_lut;
-    let out_layer_off = layer_config.output_layer_offset;
-    let output_zp = layer_config.output_zp;
-
-    let in_groups = in_channels / 4u;
-
-    let pixel = vec2<i32>(gid.xy);
-    let out_group = gid.z;
-    let out_base_channel = out_group * 4u;
-
-    // Bias seeds the accumulator (already dequantized at conversion time).
-    var acc = vec4<f32>(0.0);
-    if (has_bias != 0u) {
-        let b0 = biases[bias_off + out_base_channel + 0u];
-        let b1 = biases[bias_off + out_base_channel + 1u];
-        let b2 = biases[bias_off + out_base_channel + 2u];
-        let b3 = biases[bias_off + out_base_channel + 3u];
-        acc = vec4<f32>(b0, b1, b2, b3);
-    }
-
-    let half_k = i32(k) / 2;
-
-    // Weights: (out_channels, ky, kx, in_channels) row-major, packed int8
-    // along the innermost dim.
-    let row_stride_u32 = in_channels / 4u;
-    let kxy_stride_u32 = k * row_stride_u32;
-    let oc_stride_u32 = k * kxy_stride_u32;
-
-    for (var ky = 0u; ky < k; ky = ky + 1u) {
-        for (var kx = 0u; kx < k; kx = kx + 1u) {
-            let tap = pixel + vec2<i32>(i32(kx) - half_k, i32(ky) - half_k);
-
-            for (var ig = 0u; ig < in_groups; ig = ig + 1u) {
-                let v = nss_load_input_pixel(tap, ig);
-
-                // 4 output lanes: 1 u32 = 4 int8 weights for this lane's
-                // (ky, kx, ic_group of 4 ic lanes).
-                let base_off = weights_off
-                    + ky * kxy_stride_u32
-                    + kx * row_stride_u32
-                    + ig;
-
-                let w0 = unpack4x8snorm(weights[base_off + (out_base_channel + 0u) * oc_stride_u32]);
-                let w1 = unpack4x8snorm(weights[base_off + (out_base_channel + 1u) * oc_stride_u32]);
-                let w2 = unpack4x8snorm(weights[base_off + (out_base_channel + 2u) * oc_stride_u32]);
-                let w3 = unpack4x8snorm(weights[base_off + (out_base_channel + 3u) * oc_stride_u32]);
-
-                acc.x += dot(v, w0);
-                acc.y += dot(v, w1);
-                acc.z += dot(v, w2);
-                acc.w += dot(v, w3);
-            }
-        }
-    }
-
-    // Per-channel TOSA RESCALE (= 127 * mult / 2^shift, baked at
-    // conversion time). Then saturate to [-1, 1] (the int8 saturation
-    // that ARM's RESCALE applies implicitly).
-    //
-    // network_acc_scale (default 1.0) is a global multiplier on the
-    // accumulator — debug knob to compensate for the fp/int scale mismatch
-    // between our pipeline and ARM's training distribution. With it at 1
-    // the head LUTs return their centre values (≈0.498) on most pixels
-    // because the conv accumulators land too close to 0 to push the LUT
-    // out of its transition zone.
-    let r0 = rescales[rescale_off + out_base_channel + 0u];
-    let r1 = rescales[rescale_off + out_base_channel + 1u];
-    let r2 = rescales[rescale_off + out_base_channel + 2u];
-    let r3 = rescales[rescale_off + out_base_channel + 3u];
-
-    let rescaled = clamp(
-        acc * vec4<f32>(r0, r1, r2, r3) * settings.network_acc_scale,
-        vec4<f32>(-1.0),
-        vec4<f32>( 1.0),
-    );
-
-    var out_v: vec4<f32>;
-    if (has_lut != 0u) {
-        // Output heads: per-head transfer LUT replaces the activation.
-        out_v = vec4<f32>(
-            nss_apply_lut(rescaled.x, lut_off, output_zp),
-            nss_apply_lut(rescaled.y, lut_off, output_zp),
-            nss_apply_lut(rescaled.z, lut_off, output_zp),
-            nss_apply_lut(rescaled.w, lut_off, output_zp),
-        );
-    } else {
-        out_v = nss_activation(rescaled, activation);
-    }
-
-    textureStore(tOutput, vec3<i32>(pixel, i32(out_layer_off + out_group)), out_v);
-}
-    `, [
-        NSS_LAYER_CONFIG_STRUCT.declaration_chunk,
-        NSS_SETTINGS_STRUCT.declaration_chunk,
-    ]
-);
-
-export const shader_nss_conv_layer = ComputeShader.from({
-    label: "NSS / conv layer",
-    body,
-    resources,
-    group_size: [NSS_WORKGROUP_SIZE, NSS_WORKGROUP_SIZE, 1],
-});
+import { f32, u32 } from "../../shader/type/WEBGPU_WGSL_PRIMITIVE_TYPES.js";
+import { WebGPUArray } from "../../shader/type/WebGPUArray.js";
+import { CodeChunk } from "../../shader/compiler/CodeChunk.js";
+import { ComputeShader } from "../../shader/ComputeShader.js";
+import { ShaderResourceSetDescriptor } from "../../shader/resource/ShaderResourceSetDescriptor.js";
+import { NSS_WORKGROUP_SIZE } from "./NSS_CONSTANTS.js";
+import { NSS_LAYER_CONFIG_STRUCT } from "./NSS_LAYER_CONFIG_STRUCT.js";
+import { NSS_SETTINGS_STRUCT } from "./NSS_SETTINGS_STRUCT.js";
+
+/**
+ * One 3x3 (or 1x1) convolution layer of the NSS network — int8 weights,
+ * f32 biases, per-channel TOSA RESCALE applied before activation.
+ *
+ * Pipeline per output channel:
+ *   acc = sum(input * unpack4x8snorm(weight))      (over ic, ky, kx)
+ *   acc += bias                                     (when has_bias)
+ *   acc *= rescale                                   (per-channel, from ARM's RESCALE multipliers/shifts)
+ *   acc = saturate to [-1, 1]                        (mirrors int8 saturation)
+ *   acc = activation(acc)                            (ReLU clamps to [0, 1]; identity passes through)
+ *   textureStore(...)
+ *
+ * The TOSA RESCALE half is what keeps intermediate activations bounded —
+ * without it the encoder's 64-channel layers blow up by 4-5 orders of
+ * magnitude and downstream layers degenerate.
+ *
+ * Weights live in a shared packed-int8 storage buffer (`array<u32>`,
+ * 4 int8 weights per u32 along the in_channels dimension). Biases and
+ * rescales live in separate f32 storage buffers. `output_layer_offset`
+ * lets multiple parallel output heads share one destination texture.
+ *
+ * Edge handling: clamp-to-edge.
+ */
+const resources = new ShaderResourceSetDescriptor();
+resources.createGroup()
+    .addTexture("tInput", "float", "3d")
+    .addStorageTexture("tOutput", "rgba16float", "3d", "write-only")
+    .addStorageBuffer("weights", WebGPUArray.from(u32))
+    .addStorageBuffer("biases", WebGPUArray.from(f32))
+    .addStorageBuffer("rescales", WebGPUArray.from(f32))
+    // Per-head transfer LUTs (TOSA TABLE op). 256 int8 entries packed
+    // 4-per-u32. Only output heads (layers with `has_lut == 1`) read
+    // from this; body layers ignore the binding and use the inline
+    // ReLU activation_kind path.
+    .addStorageBuffer("luts", WebGPUArray.from(u32))
+    .addUniform("layer_config", NSS_LAYER_CONFIG_STRUCT)
+    // Shared per-frame settings — currently only network_acc_scale is read,
+    // but binding the whole struct keeps it consistent with the other passes.
+    .addUniform("settings", NSS_SETTINGS_STRUCT)
+;
+
+
+const body = CodeChunk.from(
+    //language=WGSL
+    `
+fn nss_activation(x: vec4<f32>, kind: u32) -> vec4<f32> {
+    if (kind == 1u) {
+        // ReLU on top of the rescaled+saturated input. The input is
+        // already in [-1, 1], so this gives [0, 1].
+        return max(vec4<f32>(0.0), x);
+    } else if (kind == 2u) {
+        // Leaky ReLU, alpha = 0.1
+        return select(x, 0.1 * x, x < vec4<f32>(0.0));
+    }
+    // identity — used by the output heads
+    return x;
+}
+
+fn nss_load_input_pixel(coord: vec2<i32>, group_z: u32) -> vec4<f32> {
+    let dim = textureDimensions(tInput);
+    let clamped = clamp(coord, vec2<i32>(0, 0), vec2<i32>(dim.xy) - vec2<i32>(1, 1));
+    return textureLoad(tInput, vec3<i32>(clamped, i32(group_z)), 0);
+}
+
+// Look up one entry in a layer's per-head LUT (TOSA TABLE op).
+//
+// value_fp is the un-rescaled fp accumulator scaled by the per-channel
+// rescale_factor (= 127 * mult / 2^shift) - i.e. the int8 RESCALE
+// result without the output_zp bias yet applied. output_zp is the
+// per-head signed int8 bias from TOSA RESCALE; adding it before
+// clamp+shift puts the LUT INDEX in the trained transition zone (each
+// head's LUT has almost all of its dynamic range concentrated in a
+// window centred on uint8 = output_zp + 128 - without applying it,
+// the index lands in the flat tail and the network output collapses
+// to a constant).
+//
+// LUT indexing: int8 input shifted to [0, 255]; LUT is 256 int8s
+// packed 4-per-u32. ARM uses this to encode the trained activation
+// function for the heads (sigmoid for ThetaAlpha, tanh-like S-curves
+// for the KPN columns and FB).
+//
+// Output dequantisation matches ARM's Dequantize: clamp the int8 to
+// [-127, 127] (folding -128 into -127 to dodge an off-by-one when the
+// LUT table happens to store it), then map to fp via (x + 127) / 255
+// - giving fp in [0, 1] (NOT [-1, 1] as a naive /127 would suggest).
+fn nss_apply_lut(value_fp: f32, lut_off_u32: u32, output_zp: i32) -> f32 {
+    let int8_unbiased = i32(round(value_fp * 127.0));
+    let int8_in = clamp(int8_unbiased + output_zp, -128, 127);
+    let idx = u32(int8_in + 128); // shift to [0, 255]
+    let u32_idx = idx / 4u;
+    let lane = idx % 4u;
+    // unpack4xI8 returns vec4<i32> with sign-extended int8 values.
+    let four = unpack4xI8(luts[lut_off_u32 + u32_idx]);
+    let int8_out = max(four[lane], -127);
+    return f32(int8_out + 127) / 255.0;
+}
+
+@compute @workgroup_size(${NSS_WORKGROUP_SIZE}, ${NSS_WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let out_dim = textureDimensions(tOutput);
+    let layer_out_groups = layer_config.out_channels / 4u;
+
+    if (any(gid.xy >= out_dim.xy) || gid.z >= layer_out_groups) {
+        return;
+    }
+
+    let in_channels = layer_config.in_channels;
+    let k = layer_config.kernel_size;
+    let weights_off = layer_config.weights_offset;
+    let bias_off = layer_config.bias_offset;
+    let rescale_off = layer_config.rescale_offset;
+    let lut_off = layer_config.lut_offset;
+    let activation = layer_config.activation_kind;
+    let has_bias = layer_config.has_bias;
+    let has_lut = layer_config.has_lut;
+    let out_layer_off = layer_config.output_layer_offset;
+    let output_zp = layer_config.output_zp;
+
+    let in_groups = in_channels / 4u;
+
+    let pixel = vec2<i32>(gid.xy);
+    let out_group = gid.z;
+    let out_base_channel = out_group * 4u;
+
+    // Bias seeds the accumulator (already dequantized at conversion time).
+    var acc = vec4<f32>(0.0);
+    if (has_bias != 0u) {
+        let b0 = biases[bias_off + out_base_channel + 0u];
+        let b1 = biases[bias_off + out_base_channel + 1u];
+        let b2 = biases[bias_off + out_base_channel + 2u];
+        let b3 = biases[bias_off + out_base_channel + 3u];
+        acc = vec4<f32>(b0, b1, b2, b3);
+    }
+
+    let half_k = i32(k) / 2;
+
+    // Weights: (out_channels, ky, kx, in_channels) row-major, packed int8
+    // along the innermost dim.
+    let row_stride_u32 = in_channels / 4u;
+    let kxy_stride_u32 = k * row_stride_u32;
+    let oc_stride_u32 = k * kxy_stride_u32;
+
+    for (var ky = 0u; ky < k; ky = ky + 1u) {
+        for (var kx = 0u; kx < k; kx = kx + 1u) {
+            let tap = pixel + vec2<i32>(i32(kx) - half_k, i32(ky) - half_k);
+
+            for (var ig = 0u; ig < in_groups; ig = ig + 1u) {
+                let v = nss_load_input_pixel(tap, ig);
+
+                // 4 output lanes: 1 u32 = 4 int8 weights for this lane's
+                // (ky, kx, ic_group of 4 ic lanes).
+                let base_off = weights_off
+                    + ky * kxy_stride_u32
+                    + kx * row_stride_u32
+                    + ig;
+
+                let w0 = unpack4x8snorm(weights[base_off + (out_base_channel + 0u) * oc_stride_u32]);
+                let w1 = unpack4x8snorm(weights[base_off + (out_base_channel + 1u) * oc_stride_u32]);
+                let w2 = unpack4x8snorm(weights[base_off + (out_base_channel + 2u) * oc_stride_u32]);
+                let w3 = unpack4x8snorm(weights[base_off + (out_base_channel + 3u) * oc_stride_u32]);
+
+                acc.x += dot(v, w0);
+                acc.y += dot(v, w1);
+                acc.z += dot(v, w2);
+                acc.w += dot(v, w3);
+            }
+        }
+    }
+
+    // Per-channel TOSA RESCALE (= 127 * mult / 2^shift, baked at
+    // conversion time). Then saturate to [-1, 1] (the int8 saturation
+    // that ARM's RESCALE applies implicitly).
+    //
+    // network_acc_scale (default 1.0) is a global multiplier on the
+    // accumulator — debug knob to compensate for the fp/int scale mismatch
+    // between our pipeline and ARM's training distribution. With it at 1
+    // the head LUTs return their centre values (≈0.498) on most pixels
+    // because the conv accumulators land too close to 0 to push the LUT
+    // out of its transition zone.
+    let r0 = rescales[rescale_off + out_base_channel + 0u];
+    let r1 = rescales[rescale_off + out_base_channel + 1u];
+    let r2 = rescales[rescale_off + out_base_channel + 2u];
+    let r3 = rescales[rescale_off + out_base_channel + 3u];
+
+    let rescaled = clamp(
+        acc * vec4<f32>(r0, r1, r2, r3) * settings.network_acc_scale,
+        vec4<f32>(-1.0),
+        vec4<f32>( 1.0),
+    );
+
+    var out_v: vec4<f32>;
+    if (has_lut != 0u) {
+        // Output heads: per-head transfer LUT replaces the activation.
+        out_v = vec4<f32>(
+            nss_apply_lut(rescaled.x, lut_off, output_zp),
+            nss_apply_lut(rescaled.y, lut_off, output_zp),
+            nss_apply_lut(rescaled.z, lut_off, output_zp),
+            nss_apply_lut(rescaled.w, lut_off, output_zp),
+        );
+    } else {
+        out_v = nss_activation(rescaled, activation);
+    }
+
+    textureStore(tOutput, vec3<i32>(pixel, i32(out_layer_off + out_group)), out_v);
+}
+    `, [
+        NSS_LAYER_CONFIG_STRUCT.declaration_chunk,
+        NSS_SETTINGS_STRUCT.declaration_chunk,
+    ]
+);
+
+export const shader_nss_conv_layer = ComputeShader.from({
+    label: "NSS / conv layer",
+    body,
+    resources,
+});

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

@@ -1 +1 @@
-{"version":3,"file":"shader_nss_extract_feedback.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/postprocess/nss/shader_nss_extract_feedback.js"],"names":[],"mappings":"AA4CA,wDAKG;8BAhD2B,+BAA+B"}
+{"version":3,"file":"shader_nss_extract_feedback.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/postprocess/nss/shader_nss_extract_feedback.js"],"names":[],"mappings":"AA4CA,wDAIG;8BA/C2B,+BAA+B"}

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

@@ -1,50 +1,49 @@
-import { CodeChunk } from "../../shader/compiler/CodeChunk.js";
-import { ComputeShader } from "../../shader/ComputeShader.js";
-import { ShaderResourceSetDescriptor } from "../../shader/resource/ShaderResourceSetDescriptor.js";
-import { NSS_WORKGROUP_SIZE } from "./NSS_CONSTANTS.js";
-
-/**
- * Copies the feedback channel-group (z = 0 — head slot 0 in ARM's
- * quantized model output ordering, see NSS_NETWORK_SCHEMA) of the
- * network's 3D output tensor into a standalone 2D rgba16f texture, so
- * the next frame's preprocess pass can sample it as a regular 2D
- * texture (with bilinear filtering for the warp).
- *
- * Trivial pass — included as a separate shader rather than folded into
- * postprocess because postprocess runs at output-res while feedback is
- * needed at render-res.
- */
-const resources = new ShaderResourceSetDescriptor();
-resources.createGroup()
-    .addTexture("tNetworkOut", "float", "3d")
-    .addStorageTexture("tFeedback", "rgba16float", "2d", "write-only")
-;
-
-const body = CodeChunk.from(
-    //language=WGSL
-    `
-@compute @workgroup_size(${NSS_WORKGROUP_SIZE}, ${NSS_WORKGROUP_SIZE}, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let dim = textureDimensions(tFeedback);
-
-    if (any(gid.xy >= dim)) {
-        return;
-    }
-
-    let st = vec2<i32>(gid.xy);
-
-    // Slot 0 of the network output is the feedback tensor (ARM's
-    // quantized-model OutputIndices.Feedback = 0).
-    let fb = textureLoad(tNetworkOut, vec3<i32>(st, 0), 0);
-
-    textureStore(tFeedback, st, fb);
-}
-    `
-);
-
-export const shader_nss_extract_feedback = ComputeShader.from({
-    label: "NSS / extract feedback",
-    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 { NSS_WORKGROUP_SIZE } from "./NSS_CONSTANTS.js";
+
+/**
+ * Copies the feedback channel-group (z = 0 — head slot 0 in ARM's
+ * quantized model output ordering, see NSS_NETWORK_SCHEMA) of the
+ * network's 3D output tensor into a standalone 2D rgba16f texture, so
+ * the next frame's preprocess pass can sample it as a regular 2D
+ * texture (with bilinear filtering for the warp).
+ *
+ * Trivial pass — included as a separate shader rather than folded into
+ * postprocess because postprocess runs at output-res while feedback is
+ * needed at render-res.
+ */
+const resources = new ShaderResourceSetDescriptor();
+resources.createGroup()
+    .addTexture("tNetworkOut", "float", "3d")
+    .addStorageTexture("tFeedback", "rgba16float", "2d", "write-only")
+;
+
+const body = CodeChunk.from(
+    //language=WGSL
+    `
+@compute @workgroup_size(${NSS_WORKGROUP_SIZE}, ${NSS_WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let dim = textureDimensions(tFeedback);
+
+    if (any(gid.xy >= dim)) {
+        return;
+    }
+
+    let st = vec2<i32>(gid.xy);
+
+    // Slot 0 of the network output is the feedback tensor (ARM's
+    // quantized-model OutputIndices.Feedback = 0).
+    let fb = textureLoad(tNetworkOut, vec3<i32>(st, 0), 0);
+
+    textureStore(tFeedback, st, fb);
+}
+    `
+);
+
+export const shader_nss_extract_feedback = ComputeShader.from({
+    label: "NSS / extract feedback",
+    body,
+    resources,
+});

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

@@ -1 +1 @@
-{"version":3,"file":"shader_nss_postprocess.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/postprocess/nss/shader_nss_postprocess.js"],"names":[],"mappings":"AA8RA,mDAKG;8BAlS2B,+BAA+B"}
+{"version":3,"file":"shader_nss_postprocess.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/postprocess/nss/shader_nss_postprocess.js"],"names":[],"mappings":"AA8RA,mDAIG;8BAjS2B,+BAA+B"}