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>
This commit is contained in:
@@ -0,0 +1,391 @@
|
||||
#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_vote : require
|
||||
#extension GL_KHR_shader_subgroup_arithmetic : require
|
||||
#extension GL_KHR_shader_subgroup_shuffle_relative : require
|
||||
#extension GL_EXT_samplerless_texture_functions : require
|
||||
|
||||
#if STORAGE_MODE == 0
|
||||
#extension GL_EXT_shader_8bit_storage : require
|
||||
#extension GL_EXT_shader_16bit_storage : require
|
||||
#endif
|
||||
|
||||
layout(local_size_x = 128) in;
|
||||
|
||||
layout(set = 0, binding = 0) writeonly uniform image2DArray uDequantImg;
|
||||
|
||||
layout(set = 0, binding = 1) readonly buffer PayloadOffsets
|
||||
{
|
||||
uint data[];
|
||||
} payload_offsets;
|
||||
|
||||
#if STORAGE_MODE == 0
|
||||
layout(set = 0, binding = 2) readonly buffer Payloads
|
||||
{
|
||||
uint data[];
|
||||
} payload_data_u32;
|
||||
|
||||
layout(set = 0, binding = 2) readonly buffer Payloads16
|
||||
{
|
||||
uint16_t data[];
|
||||
} payload_data_u16;
|
||||
|
||||
layout(set = 0, binding = 2) readonly buffer Payloads8
|
||||
{
|
||||
uint8_t data[];
|
||||
} payload_data_u8;
|
||||
#elif STORAGE_MODE == 1
|
||||
layout(set = 0, binding = 2) uniform usamplerBuffer PayloadU32;
|
||||
layout(set = 0, binding = 3) uniform mediump usamplerBuffer PayloadU16;
|
||||
layout(set = 0, binding = 4) uniform mediump usamplerBuffer PayloadU8;
|
||||
#else
|
||||
layout(set = 0, binding = 2) uniform utexture2D PayloadU32;
|
||||
layout(set = 0, binding = 3) uniform mediump utexture2D PayloadU16;
|
||||
layout(set = 0, binding = 4) uniform mediump utexture2D PayloadU8;
|
||||
#endif
|
||||
|
||||
#include "dwt_swizzle.h"
|
||||
#include "dwt_quant_scale.h"
|
||||
#include "constants.h"
|
||||
|
||||
layout(push_constant) uniform Registers
|
||||
{
|
||||
ivec2 resolution;
|
||||
int output_layer;
|
||||
int block_offset_32x32;
|
||||
int block_stride_32x32;
|
||||
} registers;
|
||||
|
||||
#if STORAGE_MODE == 1
|
||||
uint read_payload_u8(int coord)
|
||||
{
|
||||
return texelFetch(PayloadU8, coord).x;
|
||||
}
|
||||
|
||||
uint read_payload_u16(int coord)
|
||||
{
|
||||
return texelFetch(PayloadU16, coord).x;
|
||||
}
|
||||
|
||||
uint read_payload_u32(int coord)
|
||||
{
|
||||
return texelFetch(PayloadU32, coord).x;
|
||||
}
|
||||
#elif STORAGE_MODE == 2
|
||||
uint read_payload_u8(uint coord)
|
||||
{
|
||||
uint x = bitfieldExtract(coord, 0, 12);
|
||||
uint y = bitfieldExtract(coord, 12, 20);
|
||||
return texelFetch(PayloadU8, ivec2(x, y), 0).x;
|
||||
}
|
||||
|
||||
uint read_payload_u16(uint coord)
|
||||
{
|
||||
uint x = bitfieldExtract(coord, 0, 11);
|
||||
uint y = bitfieldExtract(coord, 11, 21);
|
||||
return texelFetch(PayloadU16, ivec2(x, y), 0).x;
|
||||
}
|
||||
|
||||
uint read_payload_u32(uint coord)
|
||||
{
|
||||
uint x = bitfieldExtract(coord, 0, 10);
|
||||
uint y = bitfieldExtract(coord, 10, 22);
|
||||
return texelFetch(PayloadU32, ivec2(x, y), 0).x;
|
||||
}
|
||||
#endif
|
||||
|
||||
mat2x4 decode_payload(uint code_word, uint q_bits, uint offset, uint block_index)
|
||||
{
|
||||
bool empty_block = code_word == 0;
|
||||
if (empty_block)
|
||||
return mat2x4(vec4(0.0), vec4(0.0));
|
||||
|
||||
int bit_offset = 2 * int(block_index);
|
||||
|
||||
// First, we need to compute the offset that our 4x2 block starts on.
|
||||
uint lsbs = code_word & 0x5555u;
|
||||
uint msbs = code_word & 0xaaaau;
|
||||
uint msbs_shift = msbs >> 1;
|
||||
msbs |= msbs_shift;
|
||||
|
||||
uint byte_offset =
|
||||
bitCount(bitfieldExtract(lsbs, 0, bit_offset)) +
|
||||
bitCount(bitfieldExtract(msbs, 0, bit_offset)) +
|
||||
q_bits * block_index + offset;
|
||||
|
||||
#if STORAGE_MODE == 0
|
||||
// Eagerly load the data to keep latency down.
|
||||
// Also forces the descriptor to be loaded early.
|
||||
uint payload = uint(payload_data_u8.data[byte_offset]);
|
||||
#else
|
||||
uint payload = read_payload_u8(int(byte_offset));
|
||||
#endif
|
||||
|
||||
uint local_control_word = bitfieldExtract(code_word, bit_offset, 2);
|
||||
int decoded_abs[8] = int[8](0, 0, 0, 0, 0, 0, 0, 0);
|
||||
int plane_iterations = int(q_bits + local_control_word);
|
||||
|
||||
for (int q = plane_iterations - 1; q >= 0; q--)
|
||||
{
|
||||
for (int b = 0; b < 8; b++)
|
||||
{
|
||||
int decoded = int(bitfieldExtract(payload, b, 1));
|
||||
decoded_abs[b] = bitfieldInsert(decoded_abs[b], decoded, q, 1);
|
||||
}
|
||||
byte_offset++;
|
||||
#if STORAGE_MODE == 0
|
||||
payload = uint(payload_data_u8.data[byte_offset]);
|
||||
#else
|
||||
payload = read_payload_u8(int(byte_offset));
|
||||
#endif
|
||||
}
|
||||
|
||||
mat2x4 m;
|
||||
|
||||
for (int i = 0; i < 4; i++)
|
||||
{
|
||||
for (int j = 0; j < 2; j++)
|
||||
{
|
||||
float v = float(decoded_abs[i * 2 + j]);
|
||||
if (v != 0.0)
|
||||
v += 0.5;
|
||||
m[j][i] = v;
|
||||
}
|
||||
}
|
||||
|
||||
return m;
|
||||
}
|
||||
|
||||
shared uint shared_sign_offset;
|
||||
shared uint shared_plane_byte_offsets[16];
|
||||
shared uint shared_sign_scan[128 / 4];
|
||||
|
||||
const int MaxScaleExp = 4;
|
||||
|
||||
float decode_quant(uint quant_code)
|
||||
{
|
||||
// Custom FP formulation for numbers in (0, 16) range.
|
||||
int e = MaxScaleExp - int(quant_code >> 3);
|
||||
int m = int(quant_code) & 0x7;
|
||||
float inv_quant = (1.0 / (8.0 * 1024.0 * 1024.0)) * float((8 + m) * (1 << (20 + e)));
|
||||
return inv_quant;
|
||||
}
|
||||
|
||||
uint scan_subgroups(uint v)
|
||||
{
|
||||
for (uint i = 1; i < gl_NumSubgroups; i *= 2)
|
||||
{
|
||||
uint up = subgroupShuffleUp(v, i);
|
||||
v += gl_SubgroupInvocationID >= i ? up : 0;
|
||||
}
|
||||
|
||||
return v;
|
||||
}
|
||||
|
||||
void scan_subgroups_fallback(uint local_index)
|
||||
{
|
||||
barrier();
|
||||
|
||||
// Slow LDS fallback for devices with wave size smaller than 16.
|
||||
bool active_lane = local_index < gl_NumSubgroups;
|
||||
|
||||
uint v = 0;
|
||||
if (active_lane)
|
||||
v = shared_sign_scan[local_index];
|
||||
|
||||
for (uint i = 1; i < gl_NumSubgroups; i *= 2)
|
||||
{
|
||||
uint up = 0;
|
||||
bool do_work = local_index >= i && active_lane;
|
||||
if (do_work)
|
||||
up = shared_sign_scan[local_index - i];
|
||||
|
||||
// Resolve write-after-read hazard.
|
||||
barrier();
|
||||
|
||||
if (do_work)
|
||||
{
|
||||
v += up;
|
||||
shared_sign_scan[local_index] = v;
|
||||
}
|
||||
|
||||
barrier();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void main()
|
||||
{
|
||||
uint local_index = gl_SubgroupID * gl_SubgroupSize + gl_SubgroupInvocationID;
|
||||
|
||||
int block_index_32x32 = int(registers.block_offset_32x32 +
|
||||
gl_WorkGroupID.y * registers.block_stride_32x32 +
|
||||
gl_WorkGroupID.x);
|
||||
|
||||
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);
|
||||
uint linear_block = block_y * 4 + block_x;
|
||||
|
||||
// Each thread individually decodes 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;
|
||||
|
||||
uint offset_u32 = payload_offsets.data[block_index_32x32];
|
||||
|
||||
if (offset_u32 == ~0u)
|
||||
{
|
||||
for (int j = 0; j < 2; j++)
|
||||
for (int i = 0; i < 4; i++)
|
||||
imageStore(uDequantImg, ivec3(coord + ivec2(i, j), registers.output_layer), vec4(0.0));
|
||||
return;
|
||||
}
|
||||
|
||||
#if STORAGE_MODE == 0
|
||||
uint ballot = payload_data_u32.data[offset_u32] & 0xffff;
|
||||
uint q_code = payload_data_u32.data[offset_u32 + 1] & 0xff;
|
||||
#else
|
||||
uint ballot = read_payload_u16(2 * int(offset_u32));
|
||||
uint q_code = read_payload_u8(4 * int(offset_u32) + 4);
|
||||
#endif
|
||||
|
||||
if (local_index < 16)
|
||||
{
|
||||
uint control_word = 0;
|
||||
uint q_bits = 0;
|
||||
|
||||
if (bitfieldExtract(ballot, int(local_index), 1) != 0)
|
||||
{
|
||||
uint local_code_offset = bitCount(bitfieldExtract(ballot, 0, int(local_index)));
|
||||
#if STORAGE_MODE == 0
|
||||
control_word = uint(payload_data_u16.data[offset_u32 * 2 + 4 + local_code_offset]);
|
||||
q_bits = uint(payload_data_u8.data[offset_u32 * 4 + 8 + bitCount(ballot) * 2 + local_code_offset]) & 0xfu;
|
||||
#else
|
||||
control_word = read_payload_u16(int(offset_u32 * 2 + 4 + local_code_offset));
|
||||
q_bits = read_payload_u8(int(offset_u32 * 4 + 8 + bitCount(ballot) * 2 + local_code_offset)) & 0xfu;
|
||||
#endif
|
||||
}
|
||||
|
||||
uint lsbs = control_word & 0x5555u;
|
||||
uint msbs = control_word & 0xaaaau;
|
||||
uint msbs_shift = msbs >> 1;
|
||||
msbs |= msbs_shift;
|
||||
uint byte_cost = bitCount(lsbs) + bitCount(msbs) + q_bits * 8;
|
||||
|
||||
uint byte_scan = offset_u32 * 4 + 8 + 3 * bitCount(ballot) + subgroupInclusiveAdd(byte_cost);
|
||||
if (local_index == 15)
|
||||
shared_sign_offset = 8 * byte_scan;
|
||||
shared_plane_byte_offsets[local_index] = byte_scan - byte_cost;
|
||||
}
|
||||
|
||||
barrier();
|
||||
|
||||
mat2x4 v;
|
||||
int significant_count;
|
||||
|
||||
if (bitfieldExtract(ballot, int(linear_block), 1) != 0)
|
||||
{
|
||||
uint local_code_offset = bitCount(bitfieldExtract(ballot, 0, int(linear_block)));
|
||||
|
||||
#if STORAGE_MODE == 0
|
||||
uint control_word = uint(payload_data_u16.data[offset_u32 * 2 + 4 + local_code_offset]);
|
||||
uint control_word2 = uint(payload_data_u8.data[offset_u32 * 4 + 8 + bitCount(ballot) * 2 + local_code_offset]);
|
||||
#else
|
||||
uint control_word = read_payload_u16(int(offset_u32 * 2 + 4 + local_code_offset));
|
||||
uint control_word2 = read_payload_u8(int(offset_u32 * 4 + 8 + bitCount(ballot) * 2 + local_code_offset));
|
||||
#endif
|
||||
|
||||
v = decode_payload(control_word, control_word2 & 0xfu,
|
||||
shared_plane_byte_offsets[linear_block], block_local_index);
|
||||
|
||||
significant_count = 0;
|
||||
for (int j = 0; j < 2; j++)
|
||||
for (int i = 0; i < 4; i++)
|
||||
significant_count += int(v[j][i] != 0.0);
|
||||
|
||||
float q = decode_quant(q_code);
|
||||
float inv_scale = q * decode_quant_scale(bitfieldExtract(control_word2, QUANT_SCALE_OFFSET - 16, QUANT_SCALE_BITS));
|
||||
|
||||
v *= inv_scale;
|
||||
}
|
||||
else
|
||||
{
|
||||
v = mat2x4(vec4(0.0), vec4(0.0));
|
||||
significant_count = 0;
|
||||
}
|
||||
|
||||
// Figure out how many significant coefficients we have.
|
||||
int significant_scan = subgroupInclusiveAdd(significant_count);
|
||||
if (gl_SubgroupInvocationID == gl_SubgroupSize - 1)
|
||||
shared_sign_scan[gl_SubgroupID] = significant_scan;
|
||||
|
||||
if (gl_NumSubgroups <= 8)
|
||||
{
|
||||
barrier();
|
||||
if (gl_SubgroupSize <= 32)
|
||||
{
|
||||
// Should be more robust since not all compilers properly understand the shuffle up pattern.
|
||||
// AMD is known to understand it well.
|
||||
if (local_index < gl_NumSubgroups)
|
||||
shared_sign_scan[local_index] = subgroupInclusiveAdd(shared_sign_scan[local_index]);
|
||||
}
|
||||
else
|
||||
{
|
||||
if (local_index < gl_NumSubgroups)
|
||||
shared_sign_scan[local_index] = scan_subgroups(shared_sign_scan[local_index]);
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
else
|
||||
{
|
||||
scan_subgroups_fallback(local_index);
|
||||
}
|
||||
|
||||
// Compute where we need to start reading sign bits from.
|
||||
uint sign_offset = shared_sign_offset + significant_scan - significant_count;
|
||||
if (gl_SubgroupID != 0)
|
||||
sign_offset += shared_sign_scan[gl_SubgroupID - 1];
|
||||
|
||||
// Read out all sign bits we could possibly access per thread.
|
||||
// On AMD at least, this 64-bit load should be vectorizable.
|
||||
#if STORAGE_MODE == 0
|
||||
uint sign_word = payload_data_u32.data[sign_offset / 32 + 0];
|
||||
uint sign_word_upper = payload_data_u32.data[sign_offset / 32 + 1];
|
||||
#else
|
||||
uint sign_word = read_payload_u32(int(sign_offset / 32 + 0));
|
||||
uint sign_word_upper = read_payload_u32(int(sign_offset / 32 + 1));
|
||||
#endif
|
||||
|
||||
uint masked_sign_offset = sign_offset & 31u;
|
||||
if (masked_sign_offset != 0)
|
||||
{
|
||||
sign_word >>= masked_sign_offset;
|
||||
sign_word |= sign_word_upper << (32 - masked_sign_offset);
|
||||
}
|
||||
|
||||
int sign_counter = 0;
|
||||
|
||||
// Clock out the sign bits as needed.
|
||||
for (int i = 0; i < 4; i++)
|
||||
{
|
||||
for (int j = 0; j < 2; j++)
|
||||
{
|
||||
if (v[j][i] != 0.0)
|
||||
{
|
||||
v[j][i] *= 1.0 - 2.0 * float(bitfieldExtract(sign_word, sign_counter, 1));
|
||||
sign_counter++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Write output.
|
||||
for (int j = 0; j < 2; j++)
|
||||
for (int i = 0; i < 4; i++)
|
||||
imageStore(uDequantImg, ivec3(coord + ivec2(i, j), registers.output_layer), vec4(v[j][i]));
|
||||
}
|
||||
Reference in New Issue
Block a user