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,379 @@
+#version 450
+// Copyright (c) 2025 Hans-Kristian Arntzen
+// SPDX-License-Identifier: MIT
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_ballot : require
+#extension GL_KHR_shader_subgroup_arithmetic : require
+#extension GL_KHR_shader_subgroup_shuffle_relative : require
+#extension GL_KHR_shader_subgroup_shuffle : require
+#extension GL_KHR_shader_subgroup_clustered : require
+#extension GL_KHR_shader_subgroup_vote : require
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+#extension GL_EXT_shader_8bit_storage : require
+
+#include "constants.h"
+
+layout(local_size_x = 64) in;
+
+struct BlockMeta
+{
+    uint code_word;
+    uint offset;
+};
+
+struct BitstreamPacket
+{
+    uint offset;
+    uint num_words;
+};
+
+layout(set = 0, binding = 0) writeonly buffer BitstreamPayload
+{
+    uint data[];
+} bitstream_data;
+
+layout(set = 0, binding = 0) writeonly buffer BitstreamPayload16Bit
+{
+    uint16_t data[];
+} bitstream_data_16b;
+
+layout(set = 0, binding = 0) writeonly buffer BitstreamPayload8Bit
+{
+    uint8_t data[];
+} bitstream_data_8b;
+
+layout(set = 0, binding = 1) writeonly buffer BitstreamMeta
+{
+    BitstreamPacket packets[];
+} bitstream_meta;
+
+layout(set = 0, binding = 2) readonly buffer SSBOMeta
+{
+    BlockMeta meta[];
+} block_meta;
+
+layout(set = 0, binding = 3) buffer Payloads
+{
+    layout(offset = 4) uint bitstream_payload_counter;
+    layout(offset = 8) uint8_t data[];
+} payload_data;
+
+struct QuantStats
+{
+    float16_t square_error;
+    uint16_t payload_cost;
+};
+
+struct BlockStats
+{
+    uint num_planes;
+    QuantStats errors[15];
+};
+
+layout(set = 0, binding = 4) readonly buffer SSBOBlockStats
+{
+    BlockStats stats[];
+} block_stats;
+
+layout(set = 0, binding = 5) readonly buffer RateControlQuant
+{
+    int data[];
+} quant_data;
+
+layout(push_constant) uniform Registers
+{
+    ivec2 resolution;
+    ivec2 resolution_32x32_blocks;
+    ivec2 resolution_8x8_blocks;
+    uint quant_resolution_code;
+    uint sequence_code;
+    int block_offset_32x32;
+    int block_stride_32x32;
+    int block_offset_8x8;
+    int block_stride_8x8;
+} registers;
+
+uint compute_required_8x8_size(uint control_word)
+{
+    int q_bits = int(bitfieldExtract(control_word, Q_PLANES_OFFSET, Q_PLANES_BITS));
+    uint lsbs = control_word & 0x5555u;
+    uint msbs = control_word & 0xaaaau;
+    uint msbs_shift = msbs >> 1;
+    msbs |= msbs_shift;
+    return bitCount(lsbs) + bitCount(msbs) + q_bits * 8;
+}
+
+uint quantize_code_word(uint control_word, int quant)
+{
+    if (quant != 0 && control_word != 0)
+    {
+        int q_bits = int(bitfieldExtract(control_word, Q_PLANES_OFFSET, Q_PLANES_BITS));
+        int sub_quant = min(q_bits, quant);
+        q_bits -= sub_quant;
+        quant -= sub_quant;
+
+        if (quant != 0)
+        {
+            quant = min(quant, 3);
+
+            uint plane0 = control_word & 0x5555u;
+            uint plane1 = (control_word & 0xaaaau) >> 1;
+            uint plane2 = plane0 & plane1;
+
+            do
+            {
+                plane0 = plane1;
+                plane1 = plane2;
+                plane2 = 0;
+                quant--;
+            } while (quant != 0);
+
+            plane0 &= ~plane1;
+
+            uint new_control_word = plane0 | (plane1 << 1);
+            control_word = bitfieldInsert(control_word, new_control_word, 0, 16);
+        }
+
+        control_word = bitfieldInsert(control_word, uint(q_bits), Q_PLANES_OFFSET, Q_PLANES_BITS);
+    }
+
+    return control_word;
+}
+
+uint copy_bytes(inout uint output_offset, uint input_offset, uint count)
+{
+    uint significant_mask = 0;
+
+    do
+    {
+        uint in_data = uint(payload_data.data[input_offset]);
+        // If we observe any 1 in the non-sign planes, it's not deadzone quantized.
+        significant_mask |= in_data;
+        bitstream_data_8b.data[output_offset++] = uint8_t(in_data);
+        count--;
+        input_offset++;
+    } while (count > 0);
+
+    return significant_mask;
+}
+
+uint modify_quant_code(uint code, int quant)
+{
+    int e = int(bitfieldExtract(code, 3, 5));
+    e = max(e - quant, 0);
+    code = bitfieldInsert(code, e, 3, 5);
+    return code;
+}
+
+uint inclusive_add_clustered16(uint v)
+{
+    for (uint i = 1; i < 16; i *= 2)
+    {
+        uint up = subgroupShuffleUp(v, i);
+        v += (gl_SubgroupInvocationID & 15) >= i ? up : 0;
+    }
+
+    return v;
+}
+
+shared uint shared_sign_bank[4][1024 / 32];
+uint pending_sign_write = 0;
+uint pending_sign_mask = 0;
+
+void append_sign_plane(uint bank, inout uint local_sign_offset, uint sign_mask, uint significant_mask)
+{
+    // Clock out one bit a time. This seems kinda slow.
+    while (significant_mask != 0)
+    {
+        int bit = findLSB(significant_mask);
+        significant_mask &= significant_mask - 1;
+        int out_bit = int(local_sign_offset & 31u);
+        pending_sign_write = bitfieldInsert(pending_sign_write, bitfieldExtract(sign_mask, bit, 1), out_bit, 1);
+        pending_sign_mask = bitfieldInsert(pending_sign_mask, 1, out_bit, 1);
+
+        if (out_bit == 31)
+        {
+            if (pending_sign_mask == ~0u)
+            {
+                shared_sign_bank[bank][local_sign_offset / 32] = pending_sign_write;
+            }
+            else
+            {
+                atomicAnd(shared_sign_bank[bank][local_sign_offset / 32], ~pending_sign_mask);
+                atomicOr(shared_sign_bank[bank][local_sign_offset / 32], pending_sign_write & pending_sign_mask);
+            }
+
+            pending_sign_mask = 0;
+        }
+
+        local_sign_offset++;
+    }
+}
+
+void flush_sign_plane(uint bank, uint local_sign_offset)
+{
+    if (pending_sign_mask != 0)
+    {
+        atomicAnd(shared_sign_bank[bank][local_sign_offset / 32], ~pending_sign_mask);
+        atomicOr(shared_sign_bank[bank][local_sign_offset / 32], pending_sign_write & pending_sign_mask);
+        pending_sign_mask = 0;
+    }
+}
+
+void main()
+{
+    uint index = gl_SubgroupInvocationID + gl_SubgroupSize * gl_SubgroupID;
+    uint linear_block_32x32_index = index >> 4;
+    ivec2 block32x32_index = 2 * ivec2(gl_WorkGroupID.xy);
+    block32x32_index.x += int(bitfieldExtract(index, 4, 1));
+    block32x32_index.y += int(bitfieldExtract(index, 5, 1));
+    ivec2 local_block_index = ivec2(bitfieldExtract(index, 0, 2), bitfieldExtract(index, 2, 2));
+    ivec2 block8x8_index = 4 * block32x32_index + local_block_index;
+
+    BlockMeta meta;
+    int quant;
+
+    bool in_range_8x8 = all(lessThan(block8x8_index, registers.resolution_8x8_blocks));
+    bool in_range_32x32 = all(lessThan(block32x32_index, registers.resolution_32x32_blocks));
+    uint num_bits_for_q = 0;
+
+    if (in_range_32x32)
+    {
+        int block_index = registers.block_offset_32x32 +
+            registers.block_stride_32x32 * block32x32_index.y +
+            block32x32_index.x;
+        quant = quant_data.data[block_index];
+    }
+    else
+    {
+        quant = 0;
+    }
+
+    if (in_range_8x8)
+    {
+        int block_index = registers.block_offset_8x8 +
+            registers.block_stride_8x8 * block8x8_index.y +
+            block8x8_index.x;
+        meta = block_meta.meta[block_index];
+        uint num_planes = block_stats.stats[block_index].num_planes;
+        num_bits_for_q = uint(block_stats.stats[block_index].errors[min(num_planes, quant)].payload_cost);
+    }
+    else
+    {
+        meta = BlockMeta(0, 0);
+    }
+
+    uint code_word = quantize_code_word(meta.code_word, quant);
+    bool active_code_word = (code_word & 0xffffu) != 0;
+
+    uvec4 code_word_ballot = subgroupBallot(active_code_word);
+    uint local_ballot = gl_SubgroupSize >= 64 && linear_block_32x32_index >= 2 ? code_word_ballot.y : code_word_ballot.x;
+    local_ballot = bitfieldExtract(local_ballot, int(16u * (linear_block_32x32_index & 1u)), 16);
+
+    uint required_plane_bytes = compute_required_8x8_size(code_word);
+    uint required_sign_bits = num_bits_for_q - required_plane_bytes * 8;
+
+    uint required_bits_with_meta = num_bits_for_q;
+    if (required_bits_with_meta != 0)
+        required_bits_with_meta += 24;
+
+    const uint HeaderSize = 2;
+    bool writes_header =
+        all(lessThan(block32x32_index, registers.resolution_32x32_blocks)) && (index & 15u) == 15u;
+
+    uint payload_total_bits = subgroupClusteredAdd(required_bits_with_meta, 16);
+    uint payload_total_words = (payload_total_bits + 31) / 32;
+    if (payload_total_words != 0)
+        payload_total_words += HeaderSize;
+
+    uint global_payload_offset = 0;
+    if (writes_header && payload_total_words != 0)
+        global_payload_offset = atomicAdd(payload_data.bitstream_payload_counter, payload_total_words);
+    global_payload_offset = subgroupShuffle(global_payload_offset, gl_SubgroupInvocationID | 15u);
+
+    if (writes_header)
+    {
+        uint block_index = registers.block_offset_32x32 +
+            block32x32_index.y * registers.block_stride_32x32 + block32x32_index.x;
+
+        if (payload_total_words != 0)
+        {
+            bitstream_data.data[global_payload_offset + 0] =
+                local_ballot | (payload_total_words << 16) | (registers.sequence_code << 28);
+            bitstream_data.data[global_payload_offset + 1] =
+                modify_quant_code(registers.quant_resolution_code, quant) | (block_index << 8);
+        }
+
+        bitstream_meta.packets[block_index] = BitstreamPacket(global_payload_offset, payload_total_words);
+    }
+
+    uint total_subblocks = bitCount(local_ballot);
+
+    uint total_sign_bits = inclusive_add_clustered16(required_sign_bits);
+    uint local_planes_offset = inclusive_add_clustered16(required_plane_bytes) - required_plane_bytes;
+    uint local_sign_offset = total_sign_bits - required_sign_bits;
+    uint global_planes_offset = 4 * global_payload_offset + 3 * total_subblocks + 4 * HeaderSize;
+    uint global_sign_offset = global_planes_offset + subgroupClusteredAdd(required_plane_bytes, 16);
+    global_planes_offset += local_planes_offset;
+
+    uint total_sign_bytes = (subgroupShuffle(total_sign_bits, gl_SubgroupInvocationID | 15u) + 7) / 8;
+
+    // Followed by N code words which map to the local ballot of active 16x16 regions.
+    if (active_code_word)
+    {
+        uint block_header_offset = bitCount(bitfieldExtract(
+            local_ballot, 0, local_block_index.y * 4 + local_block_index.x));
+
+        uint in_q_bits = bitfieldExtract(meta.code_word, Q_PLANES_OFFSET, Q_PLANES_BITS);
+        uint out_q_bits = bitfieldExtract(code_word, Q_PLANES_OFFSET, Q_PLANES_BITS);
+        uint input_offset = meta.offset;
+        uint output_offset = global_planes_offset;
+
+        for (int bit_offset = 0; bit_offset < 16; bit_offset += 2)
+        {
+            uint out_planes = bitfieldExtract(code_word, bit_offset, 2) + out_q_bits;
+            uint in_planes = bitfieldExtract(meta.code_word, bit_offset, 2) + in_q_bits;
+            if (in_planes != 0)
+                in_planes++;
+
+            uint sign_plane = uint(payload_data.data[input_offset]);
+
+            if (out_planes != 0)
+            {
+                uint significant_mask = copy_bytes(output_offset, input_offset + 1, out_planes);
+                append_sign_plane(linear_block_32x32_index, local_sign_offset, sign_plane, significant_mask);
+            }
+
+            input_offset += in_planes;
+        }
+
+        flush_sign_plane(linear_block_32x32_index, local_sign_offset);
+
+        bitstream_data_16b.data[2 * global_payload_offset + block_header_offset + 2 * HeaderSize] =
+            uint16_t(code_word);
+        bitstream_data_8b.data[4 * global_payload_offset + 2 * total_subblocks + block_header_offset + 4 * HeaderSize] =
+            uint8_t(code_word >> 16);
+    }
+
+    subgroupBarrier();
+
+    // Copy out all sign planes for any given group.
+    for (uint i = index & 15u; i < total_sign_bytes / 4; i += 16)
+    {
+        uint sign_word = shared_sign_bank[linear_block_32x32_index][i];
+        uint offset_8b = global_sign_offset + 4 * i;
+        bitstream_data_8b.data[offset_8b + 0] = uint8_t(sign_word >> 0);
+        bitstream_data_8b.data[offset_8b + 1] = uint8_t(sign_word >> 8);
+        bitstream_data_8b.data[offset_8b + 2] = uint8_t(sign_word >> 16);
+        bitstream_data_8b.data[offset_8b + 3] = uint8_t(sign_word >> 24);
+    }
+
+    // Copy out any stragglers.
+    for (uint i = (total_sign_bytes & ~3u) + (index & 15u); i < total_sign_bytes; i += 16)
+    {
+        uint sign_word = shared_sign_bank[linear_block_32x32_index][i / 4];
+        uint offset_8b = global_sign_offset + i;
+        bitstream_data_8b.data[offset_8b] = uint8_t(sign_word >> (8 * (i & 3u)));
+    }
+}