feat(host): vendor PyroWave + minimal Granite subset as crates/pyrowave-sys

Phase 0 of design/pyrowave-codec-plan.md — the opt-in wired-LAN ultra-low- latency codec. Vendored at upstream 509e4f88 (API 0.4.0, Granite 44362775, volk + vulkan-headers pins in PUNKTFUNK-VENDOR.txt), pruned to the 6.6 MB the standalone no-renderer build needs; scripts/vendor-pyrowave.sh reproduces the tree (a pin bump is protocol-affecting, plan §4.2). build.rs drives the wrapper CMakeLists (static archives incl. a static C-API lib upstream only ships shared) + bindgen over pyrowave.h; Linux and Windows only, empty stub elsewhere (Apple gets a native Metal port, §4.7). Offline-safe by construction: no network, no system lib, vendored Vulkan headers — same model as the opus dep (flatpak builder has no network). Phase-0 validation on .21 (RTX 5070 Ti, driver 610.43.03): - upstream pyrowave-c-test + interop test (incl. dmabuf/DRM-modifier Vulkan<->Vulkan) pass, from the pristine AND the pruned tree - GPU kernel times at ~1.6 bpp noise: encode/decode 0.090/0.042 ms @800p, 0.146/0.067 @1080p, 0.226/0.103 @1440p, 0.477/0.201 @4K — order of magnitude under NVENC's 1-2 ms retrieve, CBR lands within ~100 B of target - cargo test -p pyrowave-sys green (static link + API-version pin check) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-15 00:35:10 +02:00
parent 1b73361372
commit 4c3b11445c
396 changed files with 140058 additions and 0 deletions
@@ -0,0 +1,316 @@
+#version 450
+// Copyright (c) 2025 Hans-Kristian Arntzen
+// SPDX-License-Identifier: MIT
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_arithmetic : require
+#extension GL_KHR_shader_subgroup_ballot : require
+#extension GL_KHR_shader_subgroup_shuffle : require
+#extension GL_KHR_shader_subgroup_vote : require
+#extension GL_KHR_shader_subgroup_shuffle_relative : require
+#extension GL_KHR_shader_subgroup_clustered : require
+#extension GL_EXT_shader_8bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+
+#include "dwt_quant_scale.h"
+#include "constants.h"
+
+layout(local_size_x = 128) in;
+layout(constant_id = 1) const bool SkipQuantScale = false;
+
+layout(set = 0, binding = 0) uniform sampler2DArray uTexture;
+
+struct QuantStats
+{
+    float16_t square_error;
+    uint16_t payload_cost;
+};
+
+struct BlockMeta
+{
+    uint code_word;
+    uint offset;
+};
+
+// Fit into 64 bytes.
+struct BlockStats
+{
+    uint num_planes;
+    QuantStats errors[15];
+};
+
+layout(set = 0, binding = 1) writeonly buffer SSBOMeta
+{
+    BlockMeta meta[];
+} block_meta;
+
+layout(set = 0, binding = 2) writeonly buffer SSBOBlockStats
+{
+    BlockStats stats[];
+} block_stats;
+
+layout(set = 0, binding = 3) buffer Payloads
+{
+    layout(offset = 0) uint counter;
+    layout(offset = 8) uint8_t data[];
+} payload_data;
+
+#include "dwt_swizzle.h"
+
+layout(push_constant) uniform Registers
+{
+    ivec2 resolution;
+    ivec2 resolution_8x8_blocks;
+    vec2 inv_resolution;
+    float input_layer;
+    float quant_resolution;
+    int block_offset;
+    int block_stride;
+    float rdo_distortion_scale;
+} registers;
+
+float max4(vec4 v)
+{
+    vec2 v2 = max(v.xy, v.zw);
+    return max(v2.x, v2.y);
+}
+
+int max4(ivec4 v)
+{
+    ivec2 v2 = max(v.xy, v.zw);
+    return max(v2.x, v2.y);
+}
+
+int scan_clustered8(int v)
+{
+    for (uint i = 1; i < 8; i *= 2)
+    {
+        int up = subgroupShuffleUp(v, i);
+        v += (gl_SubgroupInvocationID & 7u) >= i ? up : 0;
+    }
+
+    return v;
+}
+
+void compute_quant_scale(float max_wave_texels, out uint quant_code, out float quant_scale)
+{
+    if (SkipQuantScale || max_wave_texels < 1.0)
+    {
+        quant_code = ENCODE_QUANT_IDENTITY;
+        quant_scale = 1.0;
+    }
+    else
+    {
+        int e;
+        frexp(max_wave_texels - 0.25, e);
+        float target_max = float(1 << e) - 0.25;
+        float inv_scale = max_wave_texels / target_max;
+        quant_code = encode_quant_scale(inv_scale);
+        quant_scale = 1.0 / decode_quant_scale(quant_code);
+    }
+}
+
+float compute_square_error(mat2x4 v, int q, out uint num_significant_values)
+{
+    v = mat2x4(abs(v[0]), abs(v[1]));
+    mat2x4 iv = mat2x4(floor(ldexp(v[0], ivec4(-q))), trunc(ldexp(v[1], ivec4(-q))));
+    num_significant_values = 0;
+    for (int j = 0; j < 2; j++)
+        for (int i = 0; i < 4; i++)
+            if (iv[j][i] != 0.0)
+                num_significant_values++;
+    iv[0] += mix(vec4(0.0), vec4(0.5), notEqual(iv[0], vec4(0.0)));
+    iv[1] += mix(vec4(0.0), vec4(0.5), notEqual(iv[1], vec4(0.0)));
+    iv = mat2x4(trunc(ldexp(iv[0], ivec4(q))), trunc(ldexp(iv[1], ivec4(q))));
+    mat2x4 err = v - iv;
+    num_significant_values = subgroupClusteredAdd(num_significant_values, 8);
+    return (dot(err[0], err[0]) + dot(err[1], err[1])) * registers.rdo_distortion_scale;
+}
+
+struct QuantResult
+{
+    float square_error;
+    int encode_cost_early;
+    int block4x2_shifted;
+    int encode_cost_late_bits;
+    int quality_planes;
+};
+
+QuantResult compute_quant_stats(mat2x4 v, int q, int msb, int block4x2_max, float inv_quant_squared)
+{
+    block4x2_max >>= q;
+
+    uint wave8_num_significants;
+    QuantResult result;
+
+    result.square_error = compute_square_error(v, q, wave8_num_significants) * inv_quant_squared;
+    result.block4x2_shifted = block4x2_max;
+
+    result.encode_cost_early = block4x2_max > 0 ? 1 : 0;
+    msb -= q;
+
+    result.quality_planes = 0;
+
+    if (msb >= 3)
+    {
+        result.quality_planes = msb - 2;
+        // Must encode the sign plane if we have quality planes.
+        result.encode_cost_early = result.quality_planes + 1;
+        result.block4x2_shifted >>= result.quality_planes;
+    }
+
+    result.encode_cost_early += findMSB(result.block4x2_shifted) + 1;
+    result.encode_cost_late_bits = 8 * subgroupClusteredAdd(max(result.encode_cost_early - 1, 0), 8) + int(wave8_num_significants);
+    return result;
+}
+
+float square(float v)
+{
+    return v * v;
+}
+
+void encode_payload(ivec2 block_index_8x8, mat2x4 texels)
+{
+    precise float max_subblock_texel = max(max4(abs(texels[0])), max4(abs(texels[1])));
+    precise float max_wave_texels = subgroupClusteredMax(max_subblock_texel, 8);
+    float quant_scale;
+    uint quant_code;
+    compute_quant_scale(max_wave_texels, quant_code, quant_scale);
+    texels *= quant_scale;
+    max_wave_texels *= quant_scale;
+    max_subblock_texel *= quant_scale;
+
+    float overall_quant_scale = registers.quant_resolution * quant_scale;
+    float inv_quant = 1.0 / overall_quant_scale;
+    float inv_quant_squared = inv_quant * inv_quant;
+    ivec4 abs_quant_texels0 = abs(ivec4(texels[0]));
+    ivec4 abs_quant_texels1 = abs(ivec4(texels[1]));
+    int max_absolute_value = int(max_wave_texels);
+    int block4x2_max = int(max_subblock_texel);
+
+    uint block_index = registers.block_offset + block_index_8x8.y * registers.block_stride + block_index_8x8.x;
+
+    // The entire block quantizes to zero.
+    if (max_absolute_value == 0)
+    {
+        if ((gl_SubgroupInvocationID & 7) == 0)
+        {
+            block_meta.meta[block_index] = BlockMeta(0, 0);
+            block_stats.stats[block_index].num_planes = 0;
+            block_stats.stats[block_index].errors[0] = QuantStats(float16_t(0.0), uint16_t(0));
+        }
+        return;
+    }
+
+    int msb = findMSB(max_absolute_value);
+
+    QuantResult result = compute_quant_stats(texels, 0, msb, block4x2_max, inv_quant_squared);
+    int scan = scan_clustered8(result.encode_cost_early);
+
+    uint global_offset = 0;
+
+    // For feedback, and allocation of payload.
+    if ((gl_SubgroupInvocationID & 7u) == 7u)
+        global_offset = atomicAdd(payload_data.counter, scan);
+    global_offset = subgroupShuffle(global_offset, gl_SubgroupInvocationID | 7u);
+
+    scan -= result.encode_cost_early;
+
+    // First, encode the code word.
+    int quality_planes = result.quality_planes;
+    uint code_word = quality_planes << Q_PLANES_OFFSET;
+    code_word = bitfieldInsert(code_word, quant_code, QUANT_SCALE_OFFSET, QUANT_SCALE_BITS);
+    uint plane_code = findMSB(result.block4x2_shifted) + 1;
+
+    uint merged_plane_code = plane_code << ((gl_SubgroupInvocationID & 7u) * 2u);
+    merged_plane_code |= subgroupShuffleXor(merged_plane_code, 1u);
+    merged_plane_code |= subgroupShuffleXor(merged_plane_code, 2u);
+    merged_plane_code |= subgroupShuffleXor(merged_plane_code, 4u);
+    code_word |= merged_plane_code;
+
+    if ((gl_SubgroupInvocationID & 7u) == 0u)
+    {
+        block_meta.meta[block_index] = BlockMeta(code_word, global_offset);
+        block_stats.stats[block_index].num_planes = msb + 1;
+        block_stats.stats[block_index].errors[0] = QuantStats(
+            float16_t(0.0), // We don't care about distortion from 0 quant since we've already made that decision.
+            uint16_t(result.encode_cost_late_bits));
+    }
+
+    for (int q = 1; q <= msb; q++)
+    {
+        QuantResult quant_result = compute_quant_stats(texels, q, msb, block4x2_max, inv_quant_squared);
+        float square_error = subgroupClusteredAdd(quant_result.square_error, 8);
+
+        if ((gl_SubgroupInvocationID & 7u) == 0)
+        {
+            block_stats.stats[block_index].errors[q] = QuantStats(
+                float16_t(min(square_error, 60000.0)),
+                uint16_t(quant_result.encode_cost_late_bits));
+        }
+    }
+
+    // Record distortion for throwing away everything.
+    float square_error = subgroupClusteredAdd((dot(texels[0], texels[0]) + dot(texels[1], texels[1])) * inv_quant_squared, 8);
+    if ((gl_SubgroupInvocationID & 7u) == 0)
+        block_stats.stats[block_index].errors[msb + 1] = QuantStats(float16_t(min(60000.0, square_error)), uint16_t(0));
+
+    uint byte_offset = scan + global_offset;
+    bool need_sign = result.block4x2_shifted != 0 || quality_planes != 0;
+
+    // Don't pack the sign plane until final pass, since we don't know how we quantize yet.
+    if (need_sign)
+    {
+        uvec4 s0 = uvec4(lessThan(texels[0], vec4(0.0))) << uvec4(0, 1, 2, 3);
+        uvec4 s1 = uvec4(lessThan(texels[1], vec4(0.0))) << uvec4(4, 5, 6, 7);
+        uint s = s0.x | s0.y | s0.z | s0.w | s1.x | s1.y | s1.z | s1.w;
+        payload_data.data[byte_offset++] = uint8_t(s);
+
+        int plane_iterations = quality_planes + int(plane_code);
+        int q = plane_iterations - 1;
+        do
+        {
+            s0 = uvec4(
+                bitfieldExtract(uint(abs_quant_texels0.x), q, 1),
+                bitfieldExtract(uint(abs_quant_texels0.y), q, 1),
+                bitfieldExtract(uint(abs_quant_texels0.z), q, 1),
+                bitfieldExtract(uint(abs_quant_texels0.w), q, 1));
+            s1 = uvec4(
+                bitfieldExtract(uint(abs_quant_texels1.x), q, 1),
+                bitfieldExtract(uint(abs_quant_texels1.y), q, 1),
+                bitfieldExtract(uint(abs_quant_texels1.z), q, 1),
+                bitfieldExtract(uint(abs_quant_texels1.w), q, 1));
+            s0 <<= uvec4(0, 1, 2, 3);
+            s1 <<= uvec4(4, 5, 6, 7);
+            s = s0.x | s0.y | s0.z | s0.w | s1.x | s1.y | s1.z | s1.w;
+            payload_data.data[byte_offset++] = uint8_t(s);
+            q--;
+        } while (q >= 0);
+    }
+}
+
+void main()
+{
+    uint local_index = gl_SubgroupID * gl_SubgroupSize + gl_SubgroupInvocationID;
+    uint block_local_index = bitfieldExtract(local_index, 0, 3);
+    uint block_x = bitfieldExtract(local_index, 3, 2);
+    uint block_y = bitfieldExtract(local_index, 5, 2);
+
+    // Each thread individually encodes 8 values.
+    ivec2 local_coord = unswizzle8x8(block_local_index << 3);
+
+    ivec2 coord = ivec2(gl_WorkGroupID.xy) * 32;
+    coord += 8 * ivec2(block_x, block_y);
+    coord += local_coord;
+
+    ivec2 block_index = 4 * ivec2(gl_WorkGroupID.xy) + ivec2(block_x, block_y);
+
+    vec3 uv = vec3(vec2(coord) * registers.inv_resolution, registers.input_layer);
+    vec4 texels0 = textureGatherOffset(uTexture, uv, ivec2(1, 1)).wxzy;
+    vec4 texels1 = textureGatherOffset(uTexture, uv, ivec2(3, 1)).wxzy;
+    precise vec4 scaled_texels0 = texels0 * registers.quant_resolution;
+    precise vec4 scaled_texels1 = texels1 * registers.quant_resolution;
+    bool in_bounds = all(lessThan(block_index, registers.resolution_8x8_blocks));
+    if (in_bounds)
+        encode_payload(block_index, mat2x4(scaled_texels0, scaled_texels1));
+}