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,221 @@
|
||||
#version 450
|
||||
#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_vote : require
|
||||
#extension GL_KHR_shader_subgroup_shuffle_relative : require
|
||||
#extension GL_KHR_shader_subgroup_shuffle : 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 = 64) in;
|
||||
|
||||
struct BlockMeta
|
||||
{
|
||||
uint code_word;
|
||||
uint offset;
|
||||
};
|
||||
|
||||
struct RDOperation
|
||||
{
|
||||
int quant;
|
||||
uint block_offset_saving;
|
||||
};
|
||||
|
||||
const int BLOCK_SPACE_SUBDIVISION = 16;
|
||||
|
||||
layout(set = 0, binding = 0) buffer Buckets
|
||||
{
|
||||
uint count;
|
||||
uint consumed_payload;
|
||||
layout(offset = 64) uint total_savings_per_bucket[128 * BLOCK_SPACE_SUBDIVISION];
|
||||
RDOperation rdo_operations[];
|
||||
} buckets;
|
||||
|
||||
struct QuantStats
|
||||
{
|
||||
float16_t square_error;
|
||||
uint16_t payload_cost;
|
||||
};
|
||||
|
||||
struct BlockStats
|
||||
{
|
||||
uint num_planes;
|
||||
QuantStats errors[15];
|
||||
};
|
||||
|
||||
layout(set = 0, binding = 1) readonly buffer SSBOBlockStats
|
||||
{
|
||||
BlockStats stats[];
|
||||
} block_stats;
|
||||
|
||||
layout(push_constant) uniform Registers
|
||||
{
|
||||
ivec2 resolution;
|
||||
ivec2 resolution_8x8_blocks;
|
||||
int block_offset_8x8;
|
||||
int block_stride_8x8;
|
||||
int block_offset_32x32;
|
||||
int block_stride_32x32;
|
||||
uint total_wg_count;
|
||||
uint num_blocks_aligned;
|
||||
uint block_index_shamt;
|
||||
} registers;
|
||||
|
||||
shared uint shared_rate_cost[16];
|
||||
shared float shared_distortion[16];
|
||||
shared uint shared_tmp[4];
|
||||
|
||||
// Perform operations that cause lower distortion first.
|
||||
uint distortion_to_bucket_index(float d, float cost, float d_base, float cost_base)
|
||||
{
|
||||
if (cost == cost_base)
|
||||
return 0;
|
||||
|
||||
// Compress a large range into 64 possible buckets.
|
||||
// Every band is ~1.5 dB.
|
||||
// Greedily chase least added (weighted) distortion per byte removed from code stream.
|
||||
float index = 60.0 + 2.0 * log2(max(d - d_base, 0.0) / (cost_base - cost));
|
||||
return uint(max(index + 0.5, 0.0));
|
||||
}
|
||||
|
||||
uint inclusive_max_clustered16(uint v)
|
||||
{
|
||||
// Ensures that we never end up with a value > 127.
|
||||
v = min(v, 128 - 16 + gl_SubgroupInvocationID);
|
||||
|
||||
for (uint i = 1; i < 16; i *= 2)
|
||||
{
|
||||
// Ensure monotonic progression for buckets.
|
||||
// Separate every quant level out by at least one bucket.
|
||||
uint up = subgroupShuffleUp(v, i) + i;
|
||||
v = max(v, gl_SubgroupInvocationID >= i ? up : 0);
|
||||
}
|
||||
|
||||
return v;
|
||||
}
|
||||
|
||||
void emit_rdo_operations()
|
||||
{
|
||||
float distortion;
|
||||
float cost;
|
||||
|
||||
if (gl_SubgroupInvocationID < 16)
|
||||
{
|
||||
cost = float(shared_rate_cost[gl_SubgroupInvocationID]);
|
||||
distortion = shared_distortion[gl_SubgroupInvocationID];
|
||||
}
|
||||
else
|
||||
{
|
||||
// Dummy values.
|
||||
cost = float(shared_rate_cost[gl_SubgroupInvocationID]);
|
||||
distortion = 1e30;
|
||||
}
|
||||
|
||||
uint bucket_index = distortion_to_bucket_index(distortion, cost, shared_distortion[0], float(shared_rate_cost[0]));
|
||||
if (gl_SubgroupInvocationID == 0)
|
||||
bucket_index = 0;
|
||||
|
||||
// Constraints:
|
||||
// bucket_index for Q1 must be less than bucket_index for Q2 if Q1 < Q2.
|
||||
// If a high quant target sees very favorable RD, lower bucket indices for lower Q values.
|
||||
uint inclusive_bucket_index = inclusive_max_clustered16(bucket_index);
|
||||
|
||||
if (gl_SubgroupInvocationID == 0)
|
||||
{
|
||||
uint unquantized_cost = shared_rate_cost[0];
|
||||
atomicAdd(buckets.consumed_payload, unquantized_cost);
|
||||
}
|
||||
else if (gl_SubgroupInvocationID < 16)
|
||||
{
|
||||
uint saving = shared_rate_cost[gl_SubgroupInvocationID - 1] - shared_rate_cost[gl_SubgroupInvocationID];
|
||||
|
||||
if (saving != 0)
|
||||
{
|
||||
ivec2 block32x32_index = ivec2(gl_WorkGroupID.xy);
|
||||
int block_index = registers.block_offset_32x32 +
|
||||
block32x32_index.y * registers.block_stride_32x32 + block32x32_index.x;
|
||||
uint subdivision = block_index >> registers.block_index_shamt;
|
||||
atomicAdd(buckets.total_savings_per_bucket[inclusive_bucket_index * BLOCK_SPACE_SUBDIVISION + subdivision], saving);
|
||||
buckets.rdo_operations[block_index + inclusive_bucket_index * registers.num_blocks_aligned] =
|
||||
RDOperation(int(gl_SubgroupInvocationID), block_index | (saving << 16));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void main()
|
||||
{
|
||||
// Each workgroup processes a 64x64 block and computes all possible rate wins for every potential quant rate.
|
||||
uint index = gl_SubgroupInvocationID + gl_SubgroupSize * gl_SubgroupID;
|
||||
ivec2 block32x32_index = ivec2(gl_WorkGroupID.xy);
|
||||
ivec2 local_block_index = ivec2(bitfieldExtract(index, 0, 2), bitfieldExtract(index, 2, 2));
|
||||
ivec2 block8x8_index = 4 * block32x32_index + local_block_index;
|
||||
|
||||
uint num_active_planes;
|
||||
|
||||
bool block8x8_in_range = all(lessThan(block8x8_index, registers.resolution_8x8_blocks));
|
||||
int block_index_8x8 = registers.block_offset_8x8 +
|
||||
registers.block_stride_8x8 * block8x8_index.y +
|
||||
block8x8_index.x;
|
||||
|
||||
if (block8x8_in_range)
|
||||
num_active_planes = block_stats.stats[block_index_8x8].num_planes;
|
||||
|
||||
uint bit_index = index >> 4;
|
||||
|
||||
for (uint i = bit_index; i < 16; i += 4)
|
||||
{
|
||||
float dist = 0.0;
|
||||
uint cost = 0;
|
||||
|
||||
if (block8x8_in_range)
|
||||
{
|
||||
QuantStats stats = block_stats.stats[block_index_8x8].errors[min(i, num_active_planes)];
|
||||
dist = float(stats.square_error);
|
||||
cost = uint(stats.payload_cost);
|
||||
}
|
||||
|
||||
// 16 bits to encode the control codes, 8 bits to encode Q bits + quant scale.
|
||||
// Cost is encoded in terms of bits. 8x8 blocks are decoded in isolation.
|
||||
if (cost != 0)
|
||||
cost += 24;
|
||||
|
||||
if (gl_SubgroupSize == 16)
|
||||
{
|
||||
cost = subgroupAdd(cost);
|
||||
dist = subgroupAdd(dist);
|
||||
}
|
||||
else
|
||||
{
|
||||
cost += subgroupShuffleXor(cost, 1);
|
||||
cost += subgroupShuffleXor(cost, 2);
|
||||
cost += subgroupShuffleXor(cost, 4);
|
||||
cost += subgroupShuffleXor(cost, 8);
|
||||
|
||||
dist += subgroupShuffleXor(dist, 1);
|
||||
dist += subgroupShuffleXor(dist, 2);
|
||||
dist += subgroupShuffleXor(dist, 4);
|
||||
dist += subgroupShuffleXor(dist, 8);
|
||||
}
|
||||
|
||||
if ((index & 15u) == 0u)
|
||||
{
|
||||
// Need to encode a header.
|
||||
// We can eliminate 32x32 blocks if everything decodes to 0.
|
||||
if (cost != 0)
|
||||
cost += 64;
|
||||
|
||||
// Each packet is aligned to 4 bytes for practical reasons.
|
||||
shared_rate_cost[i] = (cost + 31) >> 5;
|
||||
shared_distortion[i] = dist;
|
||||
}
|
||||
}
|
||||
|
||||
barrier();
|
||||
|
||||
if (gl_SubgroupID == 0)
|
||||
emit_rdo_operations();
|
||||
}
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
#version 450
|
||||
|
||||
layout(local_size_x = 512) in;
|
||||
|
||||
const int BLOCK_SPACE_SUBDIVISION = 16;
|
||||
|
||||
layout(set = 0, binding = 0) buffer Buckets
|
||||
{
|
||||
layout(offset = 64) uvec4 total_savings_per_bucket[128 * BLOCK_SPACE_SUBDIVISION / 4];
|
||||
} buckets;
|
||||
|
||||
shared uint shared_scan[512];
|
||||
|
||||
void main()
|
||||
{
|
||||
uvec4 v = buckets.total_savings_per_bucket[gl_LocalInvocationIndex];
|
||||
v.y += v.x;
|
||||
v.z += v.y;
|
||||
v.w += v.z;
|
||||
shared_scan[gl_LocalInvocationIndex] = v.w;
|
||||
|
||||
barrier();
|
||||
|
||||
for (uint step = 1u; step < gl_WorkGroupSize.x / 2u; step *= 2u)
|
||||
{
|
||||
barrier();
|
||||
|
||||
uint shuffled_up = 0;
|
||||
if (gl_LocalInvocationIndex >= step)
|
||||
shuffled_up = shared_scan[gl_LocalInvocationIndex - step];
|
||||
|
||||
barrier();
|
||||
|
||||
v += shuffled_up;
|
||||
shared_scan[gl_LocalInvocationIndex] = v.w;
|
||||
}
|
||||
|
||||
buckets.total_savings_per_bucket[gl_LocalInvocationIndex] = v;
|
||||
}
|
||||
@@ -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)));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
#ifndef CONSTANTS_H_
|
||||
#define CONSTANTS_H_
|
||||
|
||||
const int Q_PLANES_OFFSET = 16;
|
||||
const int Q_PLANES_BITS = 4;
|
||||
|
||||
const int QUANT_SCALE_OFFSET = 20;
|
||||
const int QUANT_SCALE_BITS = 4;
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,234 @@
|
||||
#version 450
|
||||
// Copyright (c) 2025 Hans-Kristian Arntzen
|
||||
// SPDX-License-Identifier: MIT
|
||||
#extension GL_KHR_shader_subgroup_basic : require
|
||||
#if FP16
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
|
||||
#endif
|
||||
|
||||
layout(local_size_x = 64) in;
|
||||
|
||||
layout(set = 0, binding = 0) uniform mediump sampler2D uTexture;
|
||||
layout(set = 0, binding = 1) writeonly uniform mediump image2DArray uOutput;
|
||||
|
||||
layout(constant_id = 0) const bool DCShift = false;
|
||||
|
||||
layout(push_constant) uniform Registers
|
||||
{
|
||||
ivec2 resolution;
|
||||
vec2 inv_resolution;
|
||||
ivec2 aligned_resolution;
|
||||
};
|
||||
|
||||
uint local_index;
|
||||
|
||||
#include "dwt_common.h"
|
||||
|
||||
vec2 generate_mirror_uv(ivec2 coord)
|
||||
{
|
||||
coord -= ivec2(lessThan(coord, ivec2(0)));
|
||||
coord += 1;
|
||||
ivec2 end_mirrored_clamp = (2 * aligned_resolution) - resolution;
|
||||
ivec2 past_wrapped_coord = coord + 2 * (resolution - aligned_resolution) + 1;
|
||||
coord = mix(min(coord, resolution), past_wrapped_coord, greaterThanEqual(coord, end_mirrored_clamp));
|
||||
|
||||
return vec2(coord) * inv_resolution;
|
||||
}
|
||||
|
||||
void load_image_with_apron()
|
||||
{
|
||||
ivec2 base_coord = ivec2(gl_WorkGroupID.xy) * ivec2(BLOCK_SIZE, BLOCK_SIZE) - APRON;
|
||||
ivec2 local_coord0 = 2 * unswizzle8x8(local_index);
|
||||
ivec2 coord0 = base_coord + local_coord0;
|
||||
|
||||
VEC4 texels0 = VEC4(textureGather(uTexture, generate_mirror_uv(coord0))).wzxy;
|
||||
VEC4 texels1 = VEC4(textureGather(uTexture, generate_mirror_uv(coord0 + ivec2(16, 0)))).wzxy;
|
||||
VEC4 texels2 = VEC4(textureGather(uTexture, generate_mirror_uv(coord0 + ivec2(0, 16)))).wzxy;
|
||||
VEC4 texels3 = VEC4(textureGather(uTexture, generate_mirror_uv(coord0 + ivec2(16, 16)))).wzxy;
|
||||
if (DCShift) { texels0 -= FLOAT(0.5); texels1 -= FLOAT(0.5); texels2 -= FLOAT(0.5); texels3 -= FLOAT(0.5); }
|
||||
|
||||
int local_coord0_y_half = local_coord0.y >> 1;
|
||||
|
||||
// Pack two lines together in one vec2. This allows packed FP16 math easily by processing two lines in parallel.
|
||||
store_shared(local_coord0_y_half + 0, local_coord0.x + 0, texels0.xz);
|
||||
store_shared(local_coord0_y_half + 0, local_coord0.x + 1, texels0.yw);
|
||||
store_shared(local_coord0_y_half + 0, local_coord0.x + 16, texels1.xz);
|
||||
store_shared(local_coord0_y_half + 0, local_coord0.x + 17, texels1.yw);
|
||||
store_shared(local_coord0_y_half + 8, local_coord0.x + 0, texels2.xz);
|
||||
store_shared(local_coord0_y_half + 8, local_coord0.x + 1, texels2.yw);
|
||||
store_shared(local_coord0_y_half + 8, local_coord0.x + 16, texels3.xz);
|
||||
store_shared(local_coord0_y_half + 8, local_coord0.x + 17, texels3.yw);
|
||||
|
||||
// Load the top-right apron
|
||||
{
|
||||
ivec2 local_coord = ivec2(BLOCK_SIZE + 2 * (local_index % 4u), 2 * (local_index / 4u));
|
||||
VEC4 texels = VEC4(textureGather(uTexture, generate_mirror_uv(base_coord + local_coord))).wzxy;
|
||||
if (DCShift) { texels -= FLOAT(0.5); }
|
||||
store_shared(local_coord.y >> 1, local_coord.x + 0, texels.xz);
|
||||
store_shared(local_coord.y >> 1, local_coord.x + 1, texels.yw);
|
||||
}
|
||||
|
||||
// Load the bottom-left apron
|
||||
{
|
||||
ivec2 local_coord = ivec2(2 * (local_index % 16u), BLOCK_SIZE + 2 * (local_index / 16u));
|
||||
VEC4 texels = VEC4(textureGather(uTexture, generate_mirror_uv(base_coord + local_coord))).wzxy;
|
||||
if (DCShift) { texels -= FLOAT(0.5); }
|
||||
store_shared(local_coord.y >> 1, local_coord.x + 0, texels.xz);
|
||||
store_shared(local_coord.y >> 1, local_coord.x + 1, texels.yw);
|
||||
}
|
||||
|
||||
if (local_index < 16)
|
||||
{
|
||||
// Load the bottom-right apron
|
||||
ivec2 local_coord = ivec2(BLOCK_SIZE + 2 * (local_index % 4u), BLOCK_SIZE + 2 * (local_index / 4u));
|
||||
VEC4 texels = VEC4(textureGather(uTexture, generate_mirror_uv(base_coord + local_coord))).wzxy;
|
||||
if (DCShift) { texels -= FLOAT(0.5); }
|
||||
store_shared(local_coord.y >> 1, local_coord.x + 0, texels.xz);
|
||||
store_shared(local_coord.y >> 1, local_coord.x + 1, texels.yw);
|
||||
}
|
||||
}
|
||||
|
||||
void forward_transform8x2()
|
||||
{
|
||||
const int SIZE = 8;
|
||||
const int PADDED_SIZE = SIZE + 2 * APRON;
|
||||
const int PADDED_SIZE_HALF = PADDED_SIZE / 2;
|
||||
VEC2 values[PADDED_SIZE];
|
||||
|
||||
ivec2 local_coord = ivec2(8 * (local_index % 4u), local_index / 4u);
|
||||
|
||||
for (int i = 0; i < PADDED_SIZE; i++)
|
||||
{
|
||||
VEC2 v = load_shared(local_coord.y, local_coord.x + i);
|
||||
values[i] = v;
|
||||
}
|
||||
|
||||
// CDF 9/7 lifting steps.
|
||||
// Arith go brrr.
|
||||
for (int i = 1; i < PADDED_SIZE - 1; i += 2)
|
||||
values[i] += ALPHA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 2; i < PADDED_SIZE - 2; i += 2)
|
||||
values[i] += BETA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 3; i < PADDED_SIZE - 3; i += 2)
|
||||
values[i] += GAMMA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 4; i < PADDED_SIZE - 4; i += 2)
|
||||
values[i] += DELTA * (values[i - 1] + values[i + 1]);
|
||||
|
||||
// Avoid WAR hazard.
|
||||
barrier();
|
||||
|
||||
for (int i = APRON_HALF; i < PADDED_SIZE_HALF - APRON_HALF; i++)
|
||||
{
|
||||
VEC2 a = values[2 * i + 0];
|
||||
VEC2 b = values[2 * i + 1];
|
||||
|
||||
// Filter kernel rescale.
|
||||
a *= inv_K;
|
||||
b *= K;
|
||||
|
||||
// Transpose the 2x2 block.
|
||||
VEC2 t0 = VEC2(a.x, b.x);
|
||||
VEC2 t1 = VEC2(a.y, b.y);
|
||||
|
||||
// Transpose write
|
||||
int y_coord = (local_coord.x >> 1) + (i - APRON_HALF);
|
||||
store_shared(y_coord, 2 * local_coord.y + 0, t0);
|
||||
store_shared(y_coord, 2 * local_coord.y + 1, t1);
|
||||
}
|
||||
}
|
||||
|
||||
void forward_transform4x2(bool active_lane, int y_offset)
|
||||
{
|
||||
const int SIZE = 4;
|
||||
const int PADDED_SIZE = SIZE + 2 * APRON;
|
||||
const int PADDED_SIZE_HALF = PADDED_SIZE / 2;
|
||||
VEC2 values[PADDED_SIZE];
|
||||
|
||||
ivec2 local_coord = ivec2(4 * (local_index % 8u), local_index / 8u + y_offset);
|
||||
|
||||
if (active_lane)
|
||||
{
|
||||
for (int i = 0; i < PADDED_SIZE; i++)
|
||||
{
|
||||
VEC2 v = load_shared(local_coord.y, local_coord.x + i);
|
||||
values[i] = v;
|
||||
}
|
||||
|
||||
// CDF 9/7 lifting steps.
|
||||
// Arith go brrr.
|
||||
for (int i = 1; i < PADDED_SIZE - 1; i += 2)
|
||||
values[i] += ALPHA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 2; i < PADDED_SIZE - 2; i += 2)
|
||||
values[i] += BETA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 3; i < PADDED_SIZE - 3; i += 2)
|
||||
values[i] += GAMMA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 4; i < PADDED_SIZE - 4; i += 2)
|
||||
values[i] += DELTA * (values[i - 1] + values[i + 1]);
|
||||
}
|
||||
|
||||
// Avoid WAR hazard.
|
||||
barrier();
|
||||
|
||||
if (active_lane)
|
||||
{
|
||||
for (int i = APRON_HALF; i < PADDED_SIZE_HALF - APRON_HALF; i++)
|
||||
{
|
||||
VEC2 a = values[2 * i + 0];
|
||||
VEC2 b = values[2 * i + 1];
|
||||
|
||||
// Filter kernel rescale.
|
||||
a *= inv_K;
|
||||
b *= K;
|
||||
|
||||
// Transpose the 2x2 block.
|
||||
VEC2 t0 = VEC2(a.x, b.x);
|
||||
VEC2 t1 = VEC2(a.y, b.y);
|
||||
|
||||
// Transpose write
|
||||
int y_coord = (local_coord.x >> 1) + (i - APRON_HALF);
|
||||
store_shared(y_coord, 2 * local_coord.y + 0, t0);
|
||||
store_shared(y_coord, 2 * local_coord.y + 1, t1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void main()
|
||||
{
|
||||
local_index = gl_SubgroupID * gl_SubgroupSize + gl_SubgroupInvocationID;
|
||||
|
||||
load_image_with_apron();
|
||||
|
||||
barrier();
|
||||
|
||||
// Horizontal transform.
|
||||
forward_transform8x2();
|
||||
|
||||
// Also need to transform the apron.
|
||||
forward_transform4x2(local_index < 32, BLOCK_SIZE_HALF);
|
||||
|
||||
barrier();
|
||||
|
||||
// Vertical transform.
|
||||
forward_transform8x2();
|
||||
|
||||
barrier();
|
||||
|
||||
ivec2 local_coord = unswizzle8x8(local_index);
|
||||
for (int y = local_coord.y; y < BLOCK_SIZE_HALF; y += 8)
|
||||
{
|
||||
for (int x = local_coord.x * 2; x < BLOCK_SIZE; x += 16)
|
||||
{
|
||||
VEC2 v0 = load_shared(y, x + 0);
|
||||
VEC2 v1 = load_shared(y, x + 1);
|
||||
|
||||
int img_x = x >> 1;
|
||||
int img_y = y;
|
||||
|
||||
ivec2 base_image_coord = ivec2(gl_WorkGroupID.xy) * (BLOCK_SIZE / 2) + ivec2(img_x, img_y);
|
||||
imageStore(uOutput, ivec3(base_image_coord, 0), v0.xxxx);
|
||||
imageStore(uOutput, ivec3(base_image_coord, 2), v0.yyyy);
|
||||
imageStore(uOutput, ivec3(base_image_coord, 1), v1.xxxx);
|
||||
imageStore(uOutput, ivec3(base_image_coord, 3), v1.yyyy);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
// Copyright (c) 2025 Hans-Kristian Arntzen
|
||||
// SPDX-License-Identifier: MIT
|
||||
#ifndef DWT_COMMON_H_
|
||||
#define DWT_COMMON_H_
|
||||
|
||||
const int APRON = 4;
|
||||
const int APRON_HALF = APRON / 2;
|
||||
const int BLOCK_SIZE = 32;
|
||||
const int BLOCK_SIZE_HALF = BLOCK_SIZE >> 1;
|
||||
|
||||
#if !FP16 && PRECISION == 0
|
||||
#undef PRECISION
|
||||
#define PRECISION 1
|
||||
#endif
|
||||
|
||||
#if PRECISION == 2
|
||||
#define FLOAT float
|
||||
#define VEC2 vec2
|
||||
#define VEC4 vec4
|
||||
#define SHARED_VEC2 vec2
|
||||
#elif PRECISION == 1
|
||||
#define FLOAT float
|
||||
#define VEC2 vec2
|
||||
#define VEC4 vec4
|
||||
#if FP16
|
||||
#define SHARED_VEC2 f16vec2
|
||||
#else
|
||||
#define SHARED_VEC2 uint
|
||||
#endif
|
||||
#else
|
||||
#define FLOAT float16_t
|
||||
#define VEC2 f16vec2
|
||||
#define VEC4 f16vec4
|
||||
#define SHARED_VEC2 f16vec2
|
||||
#endif
|
||||
|
||||
const FLOAT ALPHA = FLOAT(-1.586134342059924);
|
||||
const FLOAT BETA = FLOAT(-0.052980118572961);
|
||||
const FLOAT GAMMA = FLOAT(0.882911075530934);
|
||||
const FLOAT DELTA = FLOAT(0.443506852043971);
|
||||
const FLOAT K = FLOAT(1.230174104914001);
|
||||
const FLOAT inv_K = FLOAT(1.0 / 1.230174104914001);
|
||||
|
||||
shared SHARED_VEC2 shared_block[(BLOCK_SIZE + 2 * APRON) / 2][(BLOCK_SIZE + 2 * APRON) + 1];
|
||||
#if !FP16 && PRECISION == 1
|
||||
VEC2 load_shared(uint y, uint x) { return unpackHalf2x16(shared_block[y][x]); }
|
||||
void store_shared(uint y, uint x, VEC2 v) { shared_block[y][x] = packHalf2x16(v); }
|
||||
#else
|
||||
VEC2 load_shared(uint y, uint x) { return VEC2(shared_block[y][x]); }
|
||||
void store_shared(uint y, uint x, VEC2 v) { shared_block[y][x] = SHARED_VEC2(v); }
|
||||
#endif
|
||||
|
||||
bvec2 band(bvec2 a, bvec2 b)
|
||||
{
|
||||
return bvec2(a.x && b.x, a.y && b.y);
|
||||
}
|
||||
|
||||
#include "dwt_swizzle.h"
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,21 @@
|
||||
// Copyright (c) 2025 Hans-Kristian Arntzen
|
||||
// SPDX-License-Identifier: MIT
|
||||
#ifndef DWT_QUANT_SCALE_H_
|
||||
#define DWT_QUANT_SCALE_H_
|
||||
|
||||
float decode_quant_scale(uint code)
|
||||
{
|
||||
// Minimum scale: 0.25
|
||||
// Maximum scale: ~2.2
|
||||
return float(code) / 8.0 + 0.25;
|
||||
}
|
||||
|
||||
const uint ENCODE_QUANT_IDENTITY = 6;
|
||||
|
||||
uint encode_quant_scale(float scale)
|
||||
{
|
||||
// Round the quant scale FP up so that the quantizer scale effectively rounds down.
|
||||
return uint(ceil((scale - 0.25) * 8.0));
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,23 @@
|
||||
// Copyright (c) 2025 Hans-Kristian Arntzen
|
||||
// SPDX-License-Identifier: MIT
|
||||
#ifndef DWT_SWIZZLE_H_
|
||||
#define DWT_SWIZZLE_H_
|
||||
|
||||
ivec2 unswizzle4x8(uint index)
|
||||
{
|
||||
uint y = bitfieldExtract(index, 0, 1);
|
||||
uint x = bitfieldExtract(index, 1, 2);
|
||||
y |= bitfieldExtract(index, 3, 2) << 1;
|
||||
return ivec2(x, y);
|
||||
}
|
||||
|
||||
ivec2 unswizzle8x8(uint index)
|
||||
{
|
||||
uint y = bitfieldExtract(index, 0, 1);
|
||||
uint x = bitfieldExtract(index, 1, 2);
|
||||
y |= bitfieldExtract(index, 3, 2) << 1;
|
||||
x |= bitfieldExtract(index, 5, 1) << 2;
|
||||
return ivec2(x, y);
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,214 @@
|
||||
#version 450
|
||||
// Copyright (c) 2025 Hans-Kristian Arntzen
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#if FP16
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
|
||||
#endif
|
||||
|
||||
layout(local_size_x = 64) in;
|
||||
layout(constant_id = 0) const bool DCShift = false;
|
||||
|
||||
uint local_index;
|
||||
|
||||
#include "dwt_common.h"
|
||||
|
||||
layout(set = 0, binding = 0) uniform mediump sampler2DArray uTexture;
|
||||
layout(set = 0, binding = 1) writeonly mediump uniform image2D uOutput;
|
||||
|
||||
layout(push_constant) uniform Registers
|
||||
{
|
||||
ivec2 resolution;
|
||||
vec2 inv_resolution;
|
||||
};
|
||||
|
||||
vec2 generate_mirror_uv(ivec2 coord, bool even_x, bool even_y)
|
||||
{
|
||||
coord -= ivec2(band(bvec2(even_x, even_y), lessThan(coord, ivec2(0))));
|
||||
coord += 1;
|
||||
coord += ivec2(band(bvec2(!even_x, !even_y), greaterThanEqual(coord, resolution)));
|
||||
vec2 uv = vec2(coord) * inv_resolution;
|
||||
return uv.yx; // Transpose on load.
|
||||
}
|
||||
|
||||
void write_shared_4x4(ivec2 coord, VEC4 texels0, VEC4 texels1, VEC4 texels2, VEC4 texels3)
|
||||
{
|
||||
store_shared(coord.y + 0, 2 * coord.x + 0, VEC2(texels0.x, texels2.x));
|
||||
store_shared(coord.y + 0, 2 * coord.x + 1, VEC2(texels1.x, texels3.x));
|
||||
store_shared(coord.y + 0, 2 * coord.x + 2, VEC2(texels0.y, texels2.y));
|
||||
store_shared(coord.y + 0, 2 * coord.x + 3, VEC2(texels1.y, texels3.y));
|
||||
store_shared(coord.y + 1, 2 * coord.x + 0, VEC2(texels0.z, texels2.z));
|
||||
store_shared(coord.y + 1, 2 * coord.x + 1, VEC2(texels1.z, texels3.z));
|
||||
store_shared(coord.y + 1, 2 * coord.x + 2, VEC2(texels0.w, texels2.w));
|
||||
store_shared(coord.y + 1, 2 * coord.x + 3, VEC2(texels1.w, texels3.w));
|
||||
}
|
||||
|
||||
void load_image_with_apron()
|
||||
{
|
||||
ivec2 base_coord = ivec2(gl_WorkGroupID.xy) * ivec2(BLOCK_SIZE_HALF) - APRON_HALF;
|
||||
ivec2 local_coord0 = 2 * unswizzle8x8(local_index);
|
||||
ivec2 coord0 = base_coord + local_coord0;
|
||||
|
||||
// Transpose on load.
|
||||
VEC4 texels0 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(coord0, true, true), 0.0), 0)).wxzy;
|
||||
VEC4 texels1 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(coord0, false, true), 2.0), 0)).wxzy;
|
||||
VEC4 texels2 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(coord0, true, false), 1.0), 0)).wxzy;
|
||||
VEC4 texels3 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(coord0, false, false), 3.0), 0)).wxzy;
|
||||
write_shared_4x4(local_coord0, texels0, texels1, texels2, texels3);
|
||||
|
||||
ivec2 local_coord_horiz = ivec2(BLOCK_SIZE_HALF + 2 * (local_index % 2u), 2 * (local_index / 2u));
|
||||
if (local_coord_horiz.y < BLOCK_SIZE_HALF + 2 * APRON_HALF)
|
||||
{
|
||||
texels0 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(base_coord + local_coord_horiz, true, true), 0.0), 0)).wxzy;
|
||||
texels1 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(base_coord + local_coord_horiz, false, true), 2.0), 0)).wxzy;
|
||||
texels2 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(base_coord + local_coord_horiz, true, false), 1.0), 0)).wxzy;
|
||||
texels3 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(base_coord + local_coord_horiz, false, false), 3.0), 0)).wxzy;
|
||||
write_shared_4x4(local_coord_horiz, texels0, texels1, texels2, texels3);
|
||||
}
|
||||
|
||||
ivec2 local_coord_vert = local_coord_horiz.yx;
|
||||
if (local_coord_vert.x < BLOCK_SIZE_HALF)
|
||||
{
|
||||
texels0 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(base_coord + local_coord_vert, true, true), 0.0), 0)).wxzy;
|
||||
texels1 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(base_coord + local_coord_vert, false, true), 2.0), 0)).wxzy;
|
||||
texels2 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(base_coord + local_coord_vert, true, false), 1.0), 0)).wxzy;
|
||||
texels3 = VEC4(textureGather(uTexture, vec3(generate_mirror_uv(base_coord + local_coord_vert, false, false), 3.0), 0)).wxzy;
|
||||
write_shared_4x4(local_coord_vert, texels0, texels1, texels2, texels3);
|
||||
}
|
||||
|
||||
barrier();
|
||||
}
|
||||
|
||||
void inverse_transform8x2()
|
||||
{
|
||||
const int SIZE = 8;
|
||||
const int PADDED_SIZE = SIZE + 2 * APRON;
|
||||
const int PADDED_SIZE_HALF = PADDED_SIZE / 2;
|
||||
VEC2 values[PADDED_SIZE];
|
||||
|
||||
ivec2 local_coord = ivec2(8 * (local_index % 4u), local_index / 4u);
|
||||
|
||||
for (int i = 0; i < PADDED_SIZE; i += 2)
|
||||
{
|
||||
VEC2 v0 = load_shared(local_coord.y, local_coord.x + i + 0);
|
||||
VEC2 v1 = load_shared(local_coord.y, local_coord.x + i + 1);
|
||||
values[i + 0] = v0 * K;
|
||||
values[i + 1] = v1 * inv_K;
|
||||
}
|
||||
|
||||
// CDF 9/7 lifting steps.
|
||||
// Arith go brrr.
|
||||
for (int i = 2; i < PADDED_SIZE - 1; i += 2)
|
||||
values[i] -= DELTA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 3; i < PADDED_SIZE - 2; i += 2)
|
||||
values[i] -= GAMMA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 4; i < PADDED_SIZE - 3; i += 2)
|
||||
values[i] -= BETA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 5; i < PADDED_SIZE - 4; i += 2)
|
||||
values[i] -= ALPHA * (values[i - 1] + values[i + 1]);
|
||||
|
||||
// Avoid WAR hazard.
|
||||
barrier();
|
||||
|
||||
for (int i = APRON_HALF; i < PADDED_SIZE_HALF - APRON_HALF; i++)
|
||||
{
|
||||
VEC2 a = values[2 * i + 0];
|
||||
VEC2 b = values[2 * i + 1];
|
||||
|
||||
// Transpose the 2x2 block.
|
||||
VEC2 t0 = VEC2(a.x, b.x);
|
||||
VEC2 t1 = VEC2(a.y, b.y);
|
||||
|
||||
// Transpose write
|
||||
int y_coord = (local_coord.x >> 1) + (i - APRON_HALF);
|
||||
store_shared(y_coord, 2 * local_coord.y + 0, t0);
|
||||
store_shared(y_coord, 2 * local_coord.y + 1, t1);
|
||||
}
|
||||
}
|
||||
|
||||
void inverse_transform4x2(bool active_lane, int y_offset)
|
||||
{
|
||||
const int SIZE = 4;
|
||||
const int PADDED_SIZE = SIZE + 2 * APRON;
|
||||
const int PADDED_SIZE_HALF = PADDED_SIZE / 2;
|
||||
VEC2 values[PADDED_SIZE];
|
||||
|
||||
ivec2 local_coord = ivec2(4 * (local_index % 8u), local_index / 8u + y_offset);
|
||||
|
||||
if (active_lane)
|
||||
{
|
||||
for (int i = 0; i < PADDED_SIZE; i += 2)
|
||||
{
|
||||
VEC2 v0 = load_shared(local_coord.y, local_coord.x + i + 0);
|
||||
VEC2 v1 = load_shared(local_coord.y, local_coord.x + i + 1);
|
||||
values[i + 0] = v0 * K;
|
||||
values[i + 1] = v1 * inv_K;
|
||||
}
|
||||
|
||||
// CDF 9/7 lifting steps.
|
||||
// Arith go brrr.
|
||||
for (int i = 2; i < PADDED_SIZE - 1; i += 2)
|
||||
values[i] -= DELTA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 3; i < PADDED_SIZE - 2; i += 2)
|
||||
values[i] -= GAMMA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 4; i < PADDED_SIZE - 3; i += 2)
|
||||
values[i] -= BETA * (values[i - 1] + values[i + 1]);
|
||||
for (int i = 5; i < PADDED_SIZE - 4; i += 2)
|
||||
values[i] -= ALPHA * (values[i - 1] + values[i + 1]);
|
||||
}
|
||||
|
||||
// Avoid WAR hazard.
|
||||
barrier();
|
||||
|
||||
if (active_lane)
|
||||
{
|
||||
for (int i = APRON_HALF; i < PADDED_SIZE_HALF - APRON_HALF; i++)
|
||||
{
|
||||
VEC2 a = values[2 * i + 0];
|
||||
VEC2 b = values[2 * i + 1];
|
||||
|
||||
// Transpose the 2x2 block.
|
||||
VEC2 t0 = VEC2(a.x, b.x);
|
||||
VEC2 t1 = VEC2(a.y, b.y);
|
||||
|
||||
// Transpose write
|
||||
int y_coord = (local_coord.x >> 1) + (i - APRON_HALF);
|
||||
store_shared(y_coord, 2 * local_coord.y + 0, t0);
|
||||
store_shared(y_coord, 2 * local_coord.y + 1, t1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void main()
|
||||
{
|
||||
local_index = gl_LocalInvocationIndex;
|
||||
|
||||
load_image_with_apron();
|
||||
|
||||
// Horizontal transform.
|
||||
inverse_transform8x2();
|
||||
|
||||
// Also need to transform the apron.
|
||||
inverse_transform4x2(local_index < 32, BLOCK_SIZE_HALF);
|
||||
|
||||
barrier();
|
||||
|
||||
// Vertical transform.
|
||||
inverse_transform8x2();
|
||||
|
||||
barrier();
|
||||
|
||||
ivec2 local_coord = unswizzle8x8(local_index);
|
||||
|
||||
for (int y = local_coord.y; y < BLOCK_SIZE_HALF; y += 8)
|
||||
{
|
||||
for (int x = local_coord.x; x < BLOCK_SIZE; x += 8)
|
||||
{
|
||||
VEC2 v = load_shared(y, x);
|
||||
if (DCShift)
|
||||
v += FLOAT(0.5);
|
||||
imageStore(uOutput, ivec2(2 * y + 0, x) + BLOCK_SIZE * ivec2(gl_WorkGroupID.yx), v.xxxx);
|
||||
imageStore(uOutput, ivec2(2 * y + 1, x) + BLOCK_SIZE * ivec2(gl_WorkGroupID.yx), v.yyyy);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
#version 450
|
||||
// Copyright (c) 2025 Hans-Kristian Arntzen
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
|
||||
layout(location = 0) in vec2 vUV;
|
||||
layout(location = 0, component = 2) in float vIntCoord;
|
||||
|
||||
#if CHROMA_CONFIG == 0
|
||||
#define OUTPUT_PLANES 1
|
||||
#define INPUT_PLANES 1
|
||||
#elif CHROMA_CONFIG == 1
|
||||
#define OUTPUT_PLANES 2
|
||||
#define INPUT_PLANES 3
|
||||
#elif CHROMA_CONFIG == 2
|
||||
#define OUTPUT_PLANES 3
|
||||
#define INPUT_PLANES 2
|
||||
#else
|
||||
#error "Invalid chroma config"
|
||||
#endif
|
||||
|
||||
layout(location = 0) out mediump float oY;
|
||||
#if OUTPUT_PLANES == 2
|
||||
layout(location = 1) out mediump vec2 oCbCr;
|
||||
#elif OUTPUT_PLANES == 3
|
||||
layout(location = 1) out mediump float oCb;
|
||||
layout(location = 2) out mediump float oCr;
|
||||
#endif
|
||||
|
||||
layout(set = 0, binding = 0) uniform mediump texture2D uYEven;
|
||||
layout(set = 0, binding = 1) uniform mediump texture2D uYOdd;
|
||||
layout(set = 0, binding = 2) uniform mediump sampler uSampler;
|
||||
#if INPUT_PLANES == 3
|
||||
layout(set = 0, binding = 3) uniform mediump texture2D uCbEven;
|
||||
layout(set = 0, binding = 4) uniform mediump texture2D uCbOdd;
|
||||
layout(set = 0, binding = 5) uniform mediump texture2D uCrEven;
|
||||
layout(set = 0, binding = 6) uniform mediump texture2D uCrOdd;
|
||||
#elif INPUT_PLANES == 2
|
||||
layout(set = 0, binding = 3) uniform mediump texture2D uCbCrEven;
|
||||
layout(set = 0, binding = 4) uniform mediump texture2D uCbCrOdd;
|
||||
#endif
|
||||
|
||||
// Direct and naive implementing of the CDF 9/7 synthesis filters.
|
||||
// Optimized for the mobile GPUs which don't have any
|
||||
// competent compute/shared memory performance whatsoever,
|
||||
// i.e. anything not AMD/NV/Intel.
|
||||
|
||||
layout(constant_id = 0) const bool VERTICAL = false;
|
||||
layout(constant_id = 1) const bool FINAL_Y = false;
|
||||
layout(constant_id = 2) const bool FINAL_CBCR = false;
|
||||
layout(constant_id = 3) const int EDGE_CONDITION = 0;
|
||||
const ivec2 OFFSET_M2 = VERTICAL ? ivec2(0, 0) : ivec2(0, 0);
|
||||
const ivec2 OFFSET_M1 = VERTICAL ? ivec2(0, 1) : ivec2(1, 0);
|
||||
const ivec2 OFFSET_C = VERTICAL ? ivec2(0, 2) : ivec2(2, 0);
|
||||
const ivec2 OFFSET_P1 = VERTICAL ? ivec2(0, 3) : ivec2(3, 0);
|
||||
const ivec2 OFFSET_P2 = VERTICAL ? ivec2(0, 4) : ivec2(4, 0);
|
||||
|
||||
const float SYNTHESIS_LP_0 = 1.11508705;
|
||||
const float SYNTHESIS_LP_1 = 0.591271763114;
|
||||
const float SYNTHESIS_LP_2 = -0.057543526229;
|
||||
const float SYNTHESIS_LP_3 = -0.091271763114;
|
||||
|
||||
const float SYNTHESIS_HP_0 = 0.602949018236;
|
||||
const float SYNTHESIS_HP_1 = -0.266864118443;
|
||||
const float SYNTHESIS_HP_2 = -0.078223266529;
|
||||
const float SYNTHESIS_HP_3 = 0.016864118443;
|
||||
const float SYNTHESIS_HP_4 = 0.026748757411;
|
||||
|
||||
layout(push_constant) uniform Registers
|
||||
{
|
||||
vec2 uv_offset;
|
||||
vec2 half_texel_offset;
|
||||
float res_scale;
|
||||
int aligned_transform_size;
|
||||
};
|
||||
|
||||
float[10] sample_component_gather(mediump texture2D tex_even, mediump texture2D tex_odd)
|
||||
{
|
||||
float components[10];
|
||||
vec2 gather_uv = vUV + half_texel_offset;
|
||||
vec2 even0, even1, odd0, odd1;
|
||||
|
||||
if (VERTICAL)
|
||||
{
|
||||
even0 = textureGatherOffset(sampler2D(tex_even, uSampler), gather_uv, OFFSET_M1).wx;
|
||||
even1 = textureGatherOffset(sampler2D(tex_even, uSampler), gather_uv, OFFSET_P1).wx;
|
||||
odd0 = textureGatherOffset(sampler2D(tex_odd, uSampler), gather_uv, OFFSET_M2).wx;
|
||||
odd1 = textureGatherOffset(sampler2D(tex_odd, uSampler), gather_uv, OFFSET_C).wx;
|
||||
}
|
||||
else
|
||||
{
|
||||
even0 = textureGatherOffset(sampler2D(tex_even, uSampler), gather_uv, OFFSET_M1).wz;
|
||||
even1 = textureGatherOffset(sampler2D(tex_even, uSampler), gather_uv, OFFSET_P1).wz;
|
||||
odd0 = textureGatherOffset(sampler2D(tex_odd, uSampler), gather_uv, OFFSET_M2).wz;
|
||||
odd1 = textureGatherOffset(sampler2D(tex_odd, uSampler), gather_uv, OFFSET_C).wz;
|
||||
}
|
||||
|
||||
components[0] = 0.0;
|
||||
components[1] = odd0.x;
|
||||
components[2] = even0.x;
|
||||
components[3] = odd0.y;
|
||||
components[4] = even0.y;
|
||||
components[5] = odd1.x;
|
||||
components[6] = even1.x;
|
||||
components[7] = odd1.y;
|
||||
components[8] = even1.y;
|
||||
components[9] = textureLodOffset(sampler2D(tex_odd, uSampler), vUV, 0.0, OFFSET_P2).x;
|
||||
|
||||
return components;
|
||||
}
|
||||
|
||||
vec2[10] sample_component_gather2(mediump texture2D tex_even, mediump texture2D tex_odd)
|
||||
{
|
||||
vec2 components[10];
|
||||
|
||||
// Little point in using gather here, at least for now.
|
||||
components[0] = vec2(0.0);
|
||||
components[1] = textureLodOffset(sampler2D(tex_odd, uSampler), vUV, 0.0, OFFSET_M2).xy;
|
||||
components[2] = textureLodOffset(sampler2D(tex_even, uSampler), vUV, 0.0, OFFSET_M1).xy;
|
||||
components[3] = textureLodOffset(sampler2D(tex_odd, uSampler), vUV, 0.0, OFFSET_M1).xy;
|
||||
components[4] = textureLodOffset(sampler2D(tex_even, uSampler), vUV, 0.0, OFFSET_C).xy;
|
||||
components[5] = textureLodOffset(sampler2D(tex_odd, uSampler), vUV, 0.0, OFFSET_C).xy;
|
||||
components[6] = textureLodOffset(sampler2D(tex_even, uSampler), vUV, 0.0, OFFSET_P1).xy;
|
||||
components[7] = textureLodOffset(sampler2D(tex_odd, uSampler), vUV, 0.0, OFFSET_P1).xy;
|
||||
components[8] = textureLodOffset(sampler2D(tex_even, uSampler), vUV, 0.0, OFFSET_P2).xy;
|
||||
components[9] = textureLodOffset(sampler2D(tex_odd, uSampler), vUV, 0.0, OFFSET_P2).xy;
|
||||
|
||||
return components;
|
||||
}
|
||||
|
||||
void main()
|
||||
{
|
||||
bool is_odd = (int(vIntCoord) & 1) != 0;
|
||||
|
||||
float Y[10] = sample_component_gather(uYEven, uYOdd);
|
||||
#if INPUT_PLANES == 2
|
||||
vec2 CbCr[10] = sample_component_gather2(uCbCrEven, uCbCrOdd);
|
||||
#elif INPUT_PLANES == 3
|
||||
float Cb[10] = sample_component_gather(uCbEven, uCbOdd);
|
||||
float Cr[10] = sample_component_gather(uCrEven, uCrOdd);
|
||||
vec2 CbCr[10];
|
||||
[[unroll]]
|
||||
for (int i = 0; i < 10; i++)
|
||||
CbCr[i] = vec2(Cb[i], Cr[i]);
|
||||
#endif
|
||||
|
||||
if (EDGE_CONDITION < 0)
|
||||
{
|
||||
// The mirroring rules are particular.
|
||||
// For odd inputs we can rely on the mirrored sampling to get intended behavior.
|
||||
if (vIntCoord < 1.0)
|
||||
{
|
||||
// Y4 is the pivot.
|
||||
Y[2] = Y[6];
|
||||
#if INPUT_SAMPLES > 1
|
||||
CbCr[2] = CbCr[6];
|
||||
#endif
|
||||
}
|
||||
}
|
||||
else if (EDGE_CONDITION > 0)
|
||||
{
|
||||
if (vIntCoord + 2.0 > aligned_transform_size)
|
||||
{
|
||||
// We're on the last two pixels.
|
||||
// Y5 is the pivot. LP inputs behave as expected when using mirroring.
|
||||
Y[7] = Y[3];
|
||||
Y[9] = Y[1];
|
||||
#if INPUT_SAMPLES > 1
|
||||
CbCr[7] = CbCr[3];
|
||||
CbCr[9] = CbCr[1];
|
||||
#endif
|
||||
}
|
||||
else if (vIntCoord + 4.0 >= aligned_transform_size)
|
||||
{
|
||||
// Y7 is the pivot.
|
||||
Y[9] = Y[5];
|
||||
#if INPUT_SAMPLES > 1
|
||||
CbCr[9] = CbCr[5];
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
#if INPUT_PLANES > 1
|
||||
#define AccumT vec3
|
||||
#define GenInput(comp) vec3(Y[comp], CbCr[comp])
|
||||
#else
|
||||
#define AccumT float
|
||||
#define GenInput(comp) Y[comp]
|
||||
#endif
|
||||
|
||||
AccumT C0, C1, C2, C3, C4;
|
||||
float W0, W1, W2, W3, W4;
|
||||
|
||||
// Not ideal, but gotta do what we gotta do.
|
||||
// GPU will have to take both paths here,
|
||||
// but at least we avoid dynamic load-store which is RIP perf on these chips ...
|
||||
if (is_odd)
|
||||
{
|
||||
C0 = GenInput(5);
|
||||
C1 = GenInput(4) + GenInput(6);
|
||||
C2 = GenInput(3) + GenInput(7);
|
||||
C3 = GenInput(2) + GenInput(8);
|
||||
C4 = GenInput(1) + GenInput(9);
|
||||
|
||||
W0 = SYNTHESIS_HP_0;
|
||||
W1 = SYNTHESIS_LP_1;
|
||||
W2 = SYNTHESIS_HP_2;
|
||||
W3 = SYNTHESIS_LP_3;
|
||||
W4 = SYNTHESIS_HP_4;
|
||||
}
|
||||
else
|
||||
{
|
||||
C0 = GenInput(4);
|
||||
C1 = GenInput(3) + GenInput(5);
|
||||
C2 = GenInput(2) + GenInput(6);
|
||||
C3 = GenInput(1) + GenInput(7);
|
||||
C4 = AccumT(0.0);
|
||||
|
||||
W0 = SYNTHESIS_LP_0;
|
||||
W1 = SYNTHESIS_HP_1;
|
||||
W2 = SYNTHESIS_LP_2;
|
||||
W3 = SYNTHESIS_HP_3;
|
||||
W4 = 0.0;
|
||||
}
|
||||
|
||||
AccumT result = C0 * W0 + C1 * W1 + C2 * W2 + C3 * W3 + C4 * W4;
|
||||
|
||||
#if OUTPUT_PLANES == 3
|
||||
oY = result.x;
|
||||
oCb = result.y;
|
||||
oCr = result.z;
|
||||
#elif OUTPUT_PLANES == 2
|
||||
oY = result.x;
|
||||
oCbCr = result.yz;
|
||||
#else
|
||||
oY = result;
|
||||
#endif
|
||||
|
||||
if (FINAL_Y)
|
||||
oY += 0.5;
|
||||
|
||||
if (FINAL_CBCR)
|
||||
{
|
||||
#if OUTPUT_PLANES == 3
|
||||
oCb += 0.5;
|
||||
oCr += 0.5;
|
||||
#elif OUTPUT_PLANES == 2
|
||||
oCbCr += 0.5;
|
||||
#endif
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
#version 450
|
||||
// Copyright (c) 2025 Hans-Kristian Arntzen
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
layout(location = 0) out vec2 vUV;
|
||||
layout(location = 0, component = 2) out float vIntCoord;
|
||||
|
||||
layout(push_constant) uniform Registers
|
||||
{
|
||||
vec2 uv_offset;
|
||||
vec2 half_texel_offset;
|
||||
float res_scale;
|
||||
};
|
||||
|
||||
layout(constant_id = 0) const bool VERTICAL = false;
|
||||
|
||||
void main()
|
||||
{
|
||||
if (gl_VertexIndex == 0)
|
||||
vUV = vec2(0.0, 0.0);
|
||||
else if (gl_VertexIndex == 1)
|
||||
vUV = vec2(0.0, 2.0);
|
||||
else
|
||||
vUV = vec2(2.0, 0.0);
|
||||
|
||||
gl_Position = vec4(vUV * 2.0 - 1.0, 0.0, 1.0);
|
||||
|
||||
if (VERTICAL)
|
||||
vIntCoord = vUV.y * res_scale;
|
||||
else
|
||||
vIntCoord = vUV.x * res_scale;
|
||||
|
||||
vUV += uv_offset;
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
#version 450
|
||||
// Copyright (c) 2025 Hans-Kristian Arntzen
|
||||
// SPDX-License-Identifier: MIT
|
||||
#extension GL_EXT_samplerless_texture_functions : require
|
||||
|
||||
layout(set = 0, binding = 0) uniform texture2D uTex;
|
||||
layout(set = 0, binding = 1) writeonly uniform image2D uImage;
|
||||
|
||||
layout(local_size_x = 8, local_size_y = 8) in;
|
||||
|
||||
float power_to_db(float p)
|
||||
{
|
||||
return max(10.0 * log2(p) / log2(10.0), -100.0);
|
||||
}
|
||||
|
||||
void main()
|
||||
{
|
||||
float input_power = 0.0;
|
||||
const int Stride = 8;
|
||||
for (int y = 0; y < Stride; y++)
|
||||
{
|
||||
for (int x = 0; x < Stride; x++)
|
||||
{
|
||||
if (any(notEqual(ivec4(gl_GlobalInvocationID.xy, x, y), ivec4(0))))
|
||||
{
|
||||
vec2 c = texelFetch(uTex, ivec2(gl_GlobalInvocationID.xy) * Stride + ivec2(x, y), 0).xy;
|
||||
input_power += dot(c, c);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
imageStore(uImage, ivec2(gl_GlobalInvocationID.xy), vec4(power_to_db(input_power)));
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
#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
|
||||
|
||||
layout(local_size_x_id = 0) in;
|
||||
|
||||
struct RDOperation
|
||||
{
|
||||
int quant;
|
||||
uint block_offset_saving;
|
||||
};
|
||||
|
||||
const int BLOCK_SPACE_SUBDIVISION = 16;
|
||||
|
||||
layout(set = 0, binding = 0) readonly buffer Buckets
|
||||
{
|
||||
layout(offset = 4) int consumed_payload;
|
||||
layout(offset = 64) int total_savings_per_bucket[128 * BLOCK_SPACE_SUBDIVISION];
|
||||
RDOperation rdo_operations[];
|
||||
} buckets;
|
||||
|
||||
layout(set = 0, binding = 1) buffer QuantList
|
||||
{
|
||||
int data[];
|
||||
} quant_data;
|
||||
|
||||
layout(push_constant) uniform Registers
|
||||
{
|
||||
uint target_payload_size;
|
||||
uint num_blocks_per_subdivision;
|
||||
} registers;
|
||||
|
||||
void main()
|
||||
{
|
||||
int required_savings_per_bucket = int(buckets.consumed_payload) - int(registers.target_payload_size);
|
||||
if (gl_WorkGroupID.x != 0)
|
||||
{
|
||||
int prev_bucket_total = buckets.total_savings_per_bucket[gl_WorkGroupID.x - 1];
|
||||
// This bucket is empty.
|
||||
if (buckets.total_savings_per_bucket[gl_WorkGroupID.x] == prev_bucket_total)
|
||||
return;
|
||||
|
||||
required_savings_per_bucket -= prev_bucket_total;
|
||||
}
|
||||
else
|
||||
{
|
||||
// This bucket is empty.
|
||||
if (buckets.total_savings_per_bucket[gl_WorkGroupID.x] == 0)
|
||||
return;
|
||||
}
|
||||
|
||||
// If all previous buckets can complete the job, skip.
|
||||
if (required_savings_per_bucket <= 0)
|
||||
return;
|
||||
|
||||
uint total_saved = 0;
|
||||
|
||||
for (uint i = 0; i < registers.num_blocks_per_subdivision && total_saved < required_savings_per_bucket; i += gl_SubgroupSize)
|
||||
{
|
||||
RDOperation op = RDOperation(0, 0);
|
||||
if (i + gl_SubgroupInvocationID < registers.num_blocks_per_subdivision)
|
||||
op = buckets.rdo_operations[gl_WorkGroupID.x * registers.num_blocks_per_subdivision + i + gl_SubgroupInvocationID];
|
||||
|
||||
uint saving = bitfieldExtract(op.block_offset_saving, 16, 16);
|
||||
uint block_offset = bitfieldExtract(op.block_offset_saving, 0, 16);
|
||||
|
||||
uint scan_saving = subgroupInclusiveAdd(saving);
|
||||
|
||||
bool should_apply_quant = total_saved + scan_saving - saving < required_savings_per_bucket;
|
||||
if (should_apply_quant && saving != 0)
|
||||
atomicMax(quant_data.data[block_offset], op.quant);
|
||||
|
||||
total_saved += subgroupShuffle(scan_saving, gl_SubgroupSize - 1);
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,71 @@
|
||||
{
|
||||
"shaders": [
|
||||
{
|
||||
"name": "dwt",
|
||||
"compute": true,
|
||||
"path": "dwt.comp",
|
||||
"variants": [
|
||||
{ "define": "PRECISION", "count": 3, "resolve": false },
|
||||
{ "define": "FP16", "count": 2, "resolve": true }
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "block_packing",
|
||||
"compute": true,
|
||||
"path": "block_packing.comp"
|
||||
},
|
||||
{
|
||||
"name": "idwt",
|
||||
"compute": true,
|
||||
"path": "idwt.comp",
|
||||
"variants": [
|
||||
{ "define": "PRECISION", "count": 3, "resolve": false },
|
||||
{ "define": "FP16", "count": 2, "resolve": true }
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "idwt_vs",
|
||||
"path": "idwt.vert"
|
||||
},
|
||||
{
|
||||
"name": "idwt_fs",
|
||||
"path": "idwt.frag",
|
||||
"variants": [
|
||||
{ "define": "CHROMA_CONFIG", "count": 3 }
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "power_to_db",
|
||||
"compute": true,
|
||||
"path": "power_to_db.comp"
|
||||
},
|
||||
{
|
||||
"name": "analyze_rate_control",
|
||||
"compute": true,
|
||||
"path": "analyze_rate_control.comp"
|
||||
},
|
||||
{
|
||||
"name": "analyze_rate_control_finalize",
|
||||
"compute": true,
|
||||
"path": "analyze_rate_control_finalize.comp"
|
||||
},
|
||||
{
|
||||
"name": "resolve_rate_control",
|
||||
"compute": true,
|
||||
"path": "resolve_rate_control.comp"
|
||||
},
|
||||
{
|
||||
"name": "wavelet_quant",
|
||||
"compute": true,
|
||||
"path": "wavelet_quant.comp"
|
||||
},
|
||||
{
|
||||
"name": "wavelet_dequant",
|
||||
"compute": true,
|
||||
"path": "wavelet_dequant.comp",
|
||||
"variants": [
|
||||
{ "define": "STORAGE_MODE", "count" : 3 }
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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]));
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
#version 450
|
||||
|
||||
#if DELTA
|
||||
layout(set = 0, binding = 0) uniform texture2D Y0;
|
||||
layout(set = 0, binding = 1) uniform texture2D Y1;
|
||||
#else
|
||||
layout(set = 0, binding = 0) uniform texture2D Y;
|
||||
layout(set = 0, binding = 1) uniform texture2D Cb;
|
||||
layout(set = 0, binding = 2) uniform texture2D Cr;
|
||||
#endif
|
||||
layout(set = 0, binding = 3) uniform sampler Samp;
|
||||
|
||||
layout(location = 0) out vec3 FragColor;
|
||||
layout(location = 0) in vec2 vUV;
|
||||
|
||||
layout(constant_id = 0) const bool BT2020 = false;
|
||||
layout(constant_id = 1) const bool FullRange = false;
|
||||
|
||||
const mat3 yuv2rgb_bt709 = mat3(
|
||||
vec3(1.0, 1.0, 1.0),
|
||||
vec3(0.0, -0.13397432 / 0.7152, 1.8556),
|
||||
vec3(1.5748, -0.33480248 / 0.7152, 0.0));
|
||||
|
||||
const mat3 yuv2rgb_bt2020 = mat3(
|
||||
vec3(1.0, 1.0, 1.0),
|
||||
vec3(0.0, -0.202008 / 0.587, 1.772),
|
||||
vec3(1.402, -0.419198 / 0.587, 0.0));
|
||||
|
||||
void main()
|
||||
{
|
||||
#if DELTA
|
||||
float y0 = textureLod(sampler2D(Y0, Samp), vUV, 0.0).x;
|
||||
float y1 = textureLod(sampler2D(Y1, Samp), vUV, 0.0).x;
|
||||
FragColor = vec3(abs(y0 - y1) * 10.0);
|
||||
#else
|
||||
float y = textureLod(sampler2D(Y, Samp), vUV, 0.0).x;
|
||||
float cb = textureLod(sampler2D(Cb, Samp), vUV, 0.0).x;
|
||||
float cr = textureLod(sampler2D(Cr, Samp), vUV, 0.0).x;
|
||||
|
||||
cb -= 0.5;
|
||||
cr -= 0.5;
|
||||
|
||||
if (!FullRange)
|
||||
{
|
||||
y -= 16.0 / 255.0;
|
||||
y *= 255.0 / 219.0;
|
||||
const float ChromaScale = 255.0 / 224.0;
|
||||
cb *= ChromaScale;
|
||||
cr *= ChromaScale;
|
||||
y = clamp(y, 0.0, 1.0);
|
||||
cb = clamp(cb, -0.5, 0.5);
|
||||
cr = clamp(cr, -0.5, 0.5);
|
||||
}
|
||||
|
||||
if (BT2020)
|
||||
FragColor = yuv2rgb_bt2020 * vec3(y, cb, cr);
|
||||
else
|
||||
FragColor = yuv2rgb_bt709 * vec3(y, cb, cr);
|
||||
#endif
|
||||
}
|
||||
Reference in New Issue
Block a user