Files
punktfunk/crates/pyrowave-sys/vendor/pyrowave/pyrowave_encoder.cpp
T
enricobuehler 4c3b11445c feat(host): vendor PyroWave + minimal Granite subset as crates/pyrowave-sys
Phase 0 of design/pyrowave-codec-plan.md — the opt-in wired-LAN ultra-low-
latency codec. Vendored at upstream 509e4f88 (API 0.4.0, Granite 44362775,
volk + vulkan-headers pins in PUNKTFUNK-VENDOR.txt), pruned to the 6.6 MB
the standalone no-renderer build needs; scripts/vendor-pyrowave.sh
reproduces the tree (a pin bump is protocol-affecting, plan §4.2).

build.rs drives the wrapper CMakeLists (static archives incl. a static
C-API lib upstream only ships shared) + bindgen over pyrowave.h; Linux and
Windows only, empty stub elsewhere (Apple gets a native Metal port, §4.7).
Offline-safe by construction: no network, no system lib, vendored Vulkan
headers — same model as the opus dep (flatpak builder has no network).

Phase-0 validation on .21 (RTX 5070 Ti, driver 610.43.03):
- upstream pyrowave-c-test + interop test (incl. dmabuf/DRM-modifier
  Vulkan<->Vulkan) pass, from the pristine AND the pruned tree
- GPU kernel times at ~1.6 bpp noise: encode/decode 0.090/0.042 ms @800p,
  0.146/0.067 @1080p, 0.226/0.103 @1440p, 0.477/0.201 @4K — order of
  magnitude under NVENC's 1-2 ms retrieve, CBR lands within ~100 B of
  target
- cargo test -p pyrowave-sys green (static link + API-version pin check)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-15 00:35:10 +02:00

1264 lines
41 KiB
C++

// Copyright (c) 2025 Hans-Kristian Arntzen
// SPDX-License-Identifier: MIT
#include "pyrowave_encoder.hpp"
#include "device.hpp"
#include "buffer.hpp"
#include "image.hpp"
#include "math.hpp"
#include "pyrowave_common.hpp"
#include <algorithm>
#include <cmath>
namespace PyroWave
{
using namespace Granite;
using namespace Vulkan;
static constexpr int BlockSpaceSubdivision = 16;
static constexpr int NumRDOBuckets = 128;
static constexpr int RDOBucketOffset = 64;
static int compute_block_count_per_subdivision(int num_blocks)
{
int per_subdivision = align(num_blocks, BlockSpaceSubdivision) / BlockSpaceSubdivision;
per_subdivision = int(Util::next_pow2(per_subdivision));
return per_subdivision;
}
struct QuantizerPushData
{
ivec2 resolution;
ivec2 resolution_8x8_blocks;
vec2 inv_resolution;
float input_layer;
float quant_resolution;
int32_t block_offset;
int32_t block_stride;
float rdo_distortion_scale;
};
struct BlockPackingPushData
{
ivec2 resolution;
ivec2 resolution_32x32_blocks;
ivec2 resolution_8x8_blocks;
uint32_t quant_resolution_code;
uint32_t sequence_count;
uint32_t block_offset_32x32;
uint32_t block_stride_32x32;
uint32_t block_offset_8x8;
uint32_t block_stride_8x8;
};
struct AnalyzeRateControlPushData
{
ivec2 resolution;
ivec2 resolution_8x8_blocks;
int32_t block_offset_8x8;
int32_t block_stride_8x8;
int32_t block_offset_32x32;
int32_t block_stride_32x32;
uint32_t total_wg_count;
uint32_t num_blocks_aligned;
uint32_t block_index_shamt;
};
struct RDOperation
{
int32_t quant;
uint16_t block_offset;
uint16_t block_saving;
};
struct Encoder::Impl final : public WaveletBuffers
{
BufferHandle bucket_buffer, meta_buffer, block_stat_buffer, payload_data, quant_buffer;
bool encode(CommandBuffer &cmd, const ViewBuffers &views, const BitstreamBuffers &buffers);
bool encode_pre_transformed(CommandBuffer &cmd, const BitstreamBuffers &buffers, float quant_scale);
bool encode_quant_and_coding(CommandBuffer &cmd, const BitstreamBuffers &buffers, float quant_scale);
bool dwt(CommandBuffer &cmd, const ViewBuffers &views);
bool quant(CommandBuffer &cmd, float quant_scale);
bool analyze_rdo(CommandBuffer &cmd);
bool resolve_rdo(CommandBuffer &cmd, size_t target_payload_size);
bool block_packing(CommandBuffer &cmd, const BitstreamBuffers &buffers, float quant_scale);
float get_noise_power_normalized_quant_resolution(int level, int component, int band) const;
float get_quant_resolution(int level, int component, int band) const;
float get_quant_rdo_distortion_scale(int level, int component, int band) const;
void init_block_meta() override;
size_t compute_num_packets(const void *meta, size_t packet_boundary) const;
size_t packetize(Packet *packets, size_t packet_boundary,
void *bitstream, size_t size,
const void *mapped_meta, const void *mapped_bitstream) const;
void report_stats(const void *mapped_meta, const void *mapped_bitstream) const;
void analyze_alternative_packing(const void *mapped_meta, const void *mapped_bitstream) const;
bool validate_bitstream(const uint32_t *bitstream_u32, const BitstreamPacket *meta, uint32_t block_index) const;
uint32_t sequence_count = 0;
};
float Encoder::Impl::get_quant_rdo_distortion_scale(int level, int component, int band) const
{
// From my Linelet master thesis. Copy paste 11 years later, ah yes :D
float horiz_midpoint = (band & 1) ? 0.75f : 0.25f;
float vert_midpoint = (band & 2) ? 0.75f : 0.25f;
// Normal PC monitors.
constexpr float dpi = 96.0f;
// Compromise between couch gaming and desktop.
constexpr float viewing_distance = 1.0f;
constexpr float cpd_nyquist = 0.34f * viewing_distance * dpi;
float cpd = std::sqrt(horiz_midpoint * horiz_midpoint + vert_midpoint * vert_midpoint) *
cpd_nyquist * std::exp2(-float(level));
// Don't allow a situation where we're quantizing LL band hard.
cpd = std::max(cpd, 8.0f);
float csf = 2.6f * (0.0192f + 0.114f * cpd) * std::exp(-std::pow(0.114f * cpd, 1.1f));
// Heavily discount chroma quality.
if (component != 0 && level != DecompositionLevels - 1)
{
// Consider chroma a little more important if we're not subsampling.
if (chroma == ChromaSubsampling::Chroma420)
csf *= 0.6f;
}
// Due to filtering, distortion in lower bands will result in more noise power.
// By scaling the distortion by this factor, we ensure uniform results.
float resolution = get_noise_power_normalized_quant_resolution(level, component, band);
float weighted_resolution = csf * resolution;
// The distortion is scaled in terms of power, not amplitude.
return weighted_resolution * weighted_resolution;
}
float Encoder::Impl::get_quant_resolution(int level, int component, int band) const
{
// FP16 range is limited, and this is more than a good enough initial estimate.
return std::min<float>(
Configuration::get().get_precision() >= 1 ? 4096.0f : 512.0f,
get_noise_power_normalized_quant_resolution(level, component, band));
}
float Encoder::Impl::get_noise_power_normalized_quant_resolution(int level, int component, int band) const
{
// The initial quantization resolution aims for a flat spectrum with noise power normalization.
// The low-pass gain for CDF 9/7 is 6 dB (1 bit). Every decomposition level subtracts 6 dB.
// Maybe make this based on the max rate to have a decent initial estimate.
int bits = Configuration::get().get_precision() >= 1 ? 8 : 6;
if (band == 0)
bits += 2;
else if (band < 3)
bits += 1;
bits += level;
// Chroma starts at level 1, subtract one bit.
if (component != 0)
bits--;
return float(1 << bits);
}
void Encoder::Impl::init_block_meta()
{
WaveletBuffers::init_block_meta();
BufferCreateInfo info;
info.domain = BufferDomain::Device;
info.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
info.size = block_count_8x8 * sizeof(BlockStats);
block_stat_buffer = device->create_buffer(info);
device->set_name(*block_stat_buffer, "block-stat-buffer");
info.size = block_count_8x8 * sizeof(BlockMeta);
meta_buffer = device->create_buffer(info);
device->set_name(*meta_buffer, "meta-buffer");
// Worst case estimate.
info.size = aligned_width * aligned_height * 2;
payload_data = device->create_buffer(info);
device->set_name(*payload_data, "payload-data");
info.size = block_count_32x32 * sizeof(uint32_t);
quant_buffer = device->create_buffer(info);
device->set_name(*quant_buffer, "quant-buffer");
info.size = RDOBucketOffset;
info.size += NumRDOBuckets * BlockSpaceSubdivision * sizeof(uint32_t);
info.size += NumRDOBuckets * compute_block_count_per_subdivision(block_count_32x32) *
BlockSpaceSubdivision * sizeof(RDOperation);
bucket_buffer = device->create_buffer(info);
device->set_name(*bucket_buffer, "bucket-buffer");
}
bool Encoder::Impl::block_packing(CommandBuffer &cmd, const BitstreamBuffers &buffers, float quant_scale)
{
cmd.begin_region("DWT block packing");
auto start_packing = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
cmd.set_program(shaders.block_packing);
cmd.set_storage_buffer(0, 0, *buffers.bitstream.buffer, buffers.bitstream.offset, buffers.bitstream.size);
cmd.set_storage_buffer(0, 1, *buffers.meta.buffer, buffers.meta.offset, buffers.meta.size);
cmd.set_storage_buffer(0, 2, *meta_buffer);
cmd.set_storage_buffer(0, 3, *payload_data);
cmd.set_storage_buffer(0, 4, *block_stat_buffer);
cmd.set_storage_buffer(0, 5, *quant_buffer);
if (device->supports_subgroup_size_log2(true, 4, 6))
{
cmd.set_subgroup_size_log2(true, 4, 6);
}
else
{
LOGI("No compatible subgroup size config.\n");
return false;
}
for (int level = 0; level < DecompositionLevels; level++)
{
auto level_width = wavelet_img_high_res->get_width(level);
auto level_height = wavelet_img_high_res->get_height(level);
for (int component = 0; component < NumComponents; component++)
{
// Ignore top-level CbCr when doing 420 subsampling.
if (level == 0 && component != 0 && chroma == ChromaSubsampling::Chroma420)
continue;
char label[128];
snprintf(label, sizeof(label), "level %d, component %d", level, component);
cmd.begin_region(label);
for (int band = (level == DecompositionLevels - 1 ? 0 : 1); band < 4; band++)
{
BlockPackingPushData packing_push = {};
packing_push.resolution = ivec2(level_width, level_height);
packing_push.resolution_32x32_blocks = ivec2((level_width + 31) / 32, (level_height + 31) / 32);
packing_push.resolution_8x8_blocks = ivec2((level_width + 7) / 8, (level_height + 7) / 8);
auto quant_res = quant_scale < 0.0f ? get_quant_resolution(level, component, band) : quant_scale;
packing_push.quant_resolution_code = encode_quant(1.0f / quant_res);
packing_push.sequence_count = sequence_count;
auto &meta = block_meta[component][level][band];
packing_push.block_offset_32x32 = meta.block_offset_32x32;
packing_push.block_stride_32x32 = meta.block_stride_32x32;
packing_push.block_offset_8x8 = meta.block_offset_8x8;
packing_push.block_stride_8x8 = meta.block_stride_8x8;
cmd.push_constants(&packing_push, 0, sizeof(packing_push));
cmd.dispatch((packing_push.resolution_32x32_blocks.x + 1) / 2,
(packing_push.resolution_32x32_blocks.y + 1) / 2,
1);
}
cmd.end_region();
}
}
auto end_packing = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
cmd.barrier(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_2_CLEAR_BIT | VK_PIPELINE_STAGE_2_COPY_BIT,
VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT | VK_ACCESS_2_TRANSFER_WRITE_BIT | VK_ACCESS_2_TRANSFER_READ_BIT);
device->register_time_interval("GPU", std::move(start_packing), std::move(end_packing), "Packing");
cmd.end_region();
return true;
}
bool Encoder::Impl::resolve_rdo(CommandBuffer &cmd, size_t target_payload_size)
{
cmd.begin_region("DWT resolve");
auto start_resolve = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
if (target_payload_size >= sizeof(BitstreamSequenceHeader))
target_payload_size -= sizeof(BitstreamSequenceHeader);
cmd.set_specialization_constant_mask(1);
if (device->supports_subgroup_size_log2(true, 6, 6))
{
cmd.set_specialization_constant(0, 64);
cmd.set_subgroup_size_log2(true, 6, 6);
}
else if (device->supports_subgroup_size_log2(true, 4, 4))
{
cmd.set_specialization_constant(0, 16);
cmd.set_subgroup_size_log2(true, 4, 4);
}
else if (device->supports_subgroup_size_log2(true, 5, 5))
{
cmd.set_specialization_constant(0, 32);
cmd.set_subgroup_size_log2(true, 5, 5);
}
else
{
LOGI("No compatible subgroup size config.\n");
return false;
}
cmd.set_program(shaders.resolve_rate_control);
struct
{
uint32_t target_payload_size;
uint32_t num_blocks_per_subdivision;
} push = {};
push.target_payload_size = target_payload_size / sizeof(uint32_t);
push.num_blocks_per_subdivision = compute_block_count_per_subdivision(block_count_32x32);
cmd.push_constants(&push, 0, sizeof(push));
cmd.set_storage_buffer(0, 0, *bucket_buffer);
cmd.set_storage_buffer(0, 1, *quant_buffer);
cmd.dispatch(NumRDOBuckets * BlockSpaceSubdivision, 1, 1);
cmd.barrier(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_READ_BIT);
cmd.end_region();
auto end_resolve = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
device->register_time_interval("GPU", std::move(start_resolve), std::move(end_resolve), "Resolve");
cmd.set_specialization_constant_mask(0);
return true;
}
bool Encoder::Impl::analyze_rdo(CommandBuffer &cmd)
{
auto start_analyze = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
cmd.begin_region("DWT analyze");
cmd.set_program(shaders.analyze_rate_control);
if (device->supports_subgroup_size_log2(true, 4, 6))
{
cmd.set_subgroup_size_log2(true, 4, 6);
}
else
{
LOGI("No compatible subgroup size config.\n");
return false;
}
// Quantize
for (int level = 0; level < DecompositionLevels; level++)
{
for (int component = 0; component < NumComponents; component++)
{
// Ignore top-level CbCr when doing 420 subsampling.
if (level == 0 && component != 0 && chroma == ChromaSubsampling::Chroma420)
continue;
AnalyzeRateControlPushData push = {};
char label[128];
snprintf(label, sizeof(label), "level %d, component %d", level, component);
cmd.begin_region(label);
for (int band = (level == DecompositionLevels - 1 ? 0 : 1); band < 4; band++)
{
auto level_width = wavelet_img_high_res->get_width(level);
auto level_height = wavelet_img_high_res->get_height(level);
push.resolution.x = level_width;
push.resolution.y = level_height;
push.resolution_8x8_blocks.x = (level_width + 7) / 8;
push.resolution_8x8_blocks.y = (level_height + 7) / 8;
push.block_offset_8x8 = block_meta[component][level][band].block_offset_8x8;
push.block_stride_8x8 = block_meta[component][level][band].block_stride_8x8;
push.block_offset_32x32 = block_meta[component][level][band].block_offset_32x32;
push.block_stride_32x32 = block_meta[component][level][band].block_stride_32x32;
push.total_wg_count = block_count_32x32;
push.num_blocks_aligned = compute_block_count_per_subdivision(block_count_32x32) * BlockSpaceSubdivision;
push.block_index_shamt = Util::floor_log2(compute_block_count_per_subdivision(block_count_32x32));
cmd.push_constants(&push, 0, sizeof(push));
cmd.set_storage_buffer(0, 0, *bucket_buffer);
cmd.set_storage_buffer(0, 1, *block_stat_buffer);
cmd.dispatch((level_width + 31) / 32, (level_height + 31) / 32, 1);
}
cmd.end_region();
}
}
cmd.barrier(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_READ_BIT);
cmd.set_program(shaders.analyze_rate_control_finalize);
cmd.dispatch(1, 1, 1);
cmd.barrier(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_READ_BIT);
cmd.end_region();
auto end_analyze = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
device->register_time_interval("GPU", std::move(start_analyze), std::move(end_analyze), "Analyze");
return true;
}
bool Encoder::Impl::quant(CommandBuffer &cmd, float quant_scale)
{
auto start_quant = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
cmd.begin_region("DWT quantize");
cmd.set_program(shaders.wavelet_quant);
cmd.set_specialization_constant_mask(0);
if (device->supports_subgroup_size_log2(true, 3, 7))
{
cmd.set_subgroup_size_log2(true, 3, 7);
}
else
{
LOGI("No compatible subgroup size config.\n");
return false;
}
// Quantize
for (int level = 0; level < DecompositionLevels; level++)
{
for (int component = 0; component < NumComponents; component++)
{
// Ignore top-level CbCr when doing 420 subsampling.
if (level == 0 && component != 0 && chroma == ChromaSubsampling::Chroma420)
continue;
QuantizerPushData push = {};
char label[128];
snprintf(label, sizeof(label), "DWT quant, level %d, component %d", level, component);
cmd.begin_region(label);
for (int band = (level == DecompositionLevels - 1 ? 0 : 1); band < 4; band++)
{
float quant_res = quant_scale < 0.0f ? get_quant_resolution(level, component, band) : quant_scale;
push.resolution.x = wavelet_img_high_res->get_width(level);
push.resolution.y = wavelet_img_high_res->get_height(level);
push.resolution_8x8_blocks.x = (push.resolution.x + 7) / 8;
push.resolution_8x8_blocks.y = (push.resolution.y + 7) / 8;
push.inv_resolution.x = 1.0f / float(push.resolution.x);
push.inv_resolution.y = 1.0f / float(push.resolution.y);
push.input_layer = float(band);
push.quant_resolution = 1.0f / decode_quant(encode_quant(1.0f / quant_res));
push.rdo_distortion_scale = get_quant_rdo_distortion_scale(level, component, band) * (1.0f / 256.0f);
int blocks_x = (push.resolution.x + 31) / 32;
int blocks_y = (push.resolution.y + 31) / 32;
push.block_offset = block_meta[component][level][band].block_offset_8x8;
push.block_stride = block_meta[component][level][band].block_stride_8x8;
cmd.push_constants(&push, 0, sizeof(push));
cmd.set_texture(0, 0, *component_layer_views[component][level], *border_sampler);
cmd.set_storage_buffer(0, 1, *meta_buffer);
cmd.set_storage_buffer(0, 2, *block_stat_buffer);
cmd.set_storage_buffer(0, 3, *payload_data);
cmd.dispatch(blocks_x, blocks_y, 1);
}
cmd.end_region();
}
}
cmd.end_region();
cmd.barrier(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_READ_BIT);
auto end_quant = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
device->register_time_interval("GPU", std::move(start_quant), std::move(end_quant), "Quant");
return true;
}
bool Encoder::Impl::dwt(CommandBuffer &cmd, const ViewBuffers &views)
{
struct Push
{
uvec2 resolution;
vec2 inv_resolution;
uvec2 aligned_resolution;
} push = {};
// Forward transforms.
cmd.set_program(shaders.dwt[PYROWAVE_PRECISION]);
// Only need simple 2-lane swaps.
cmd.set_subgroup_size_log2(true, 2, 7);
cmd.set_specialization_constant_mask(1);
cmd.set_specialization_constant(0, false);
auto start_dwt = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
for (int output_level = 0; output_level < DecompositionLevels; output_level++)
{
if (output_level > 0)
{
push.resolution = uvec2(component_ll_views[0][output_level - 1]->get_view_width(),
component_ll_views[0][output_level - 1]->get_view_height());
push.aligned_resolution = push.resolution;
}
else
{
push.resolution = uvec2(views.planes[0]->get_view_width(), views.planes[0]->get_view_height());
push.aligned_resolution.x = aligned_width;
push.aligned_resolution.y = aligned_height;
}
push.inv_resolution.x = 1.0f / float(push.resolution.x);
push.inv_resolution.y = 1.0f / float(push.resolution.y);
cmd.push_constants(&push, 0, sizeof(push));
if (output_level == 0)
{
cmd.set_specialization_constant(0, true);
if (chroma == ChromaSubsampling::Chroma444)
{
for (int c = 0; c < NumComponents; c++)
{
char label[64];
snprintf(label, sizeof(label), "DWT level 0, component %u", c);
cmd.begin_region(label);
cmd.set_texture(0, 0, *views.planes[c], *mirror_repeat_sampler);
cmd.set_storage_texture(0, 1, *component_layer_views[c][output_level]);
cmd.dispatch((push.aligned_resolution.x + 31) / 32, (push.aligned_resolution.y + 31) / 32, 1);
cmd.end_region();
}
}
else
{
cmd.set_texture(0, 0, *views.planes[0], *mirror_repeat_sampler);
cmd.set_storage_texture(0, 1, *component_layer_views[0][output_level]);
cmd.begin_region("DWT level 0 Y");
cmd.dispatch((push.aligned_resolution.x + 31) / 32, (push.aligned_resolution.y + 31) / 32, 1);
cmd.end_region();
}
}
else
{
for (int c = 0; c < NumComponents; c++)
{
if (chroma == ChromaSubsampling::Chroma420 && c != 0 && output_level == 1)
{
push.resolution = uvec2(views.planes[c]->get_view_width(), views.planes[c]->get_view_height());
push.aligned_resolution.x = aligned_width >> output_level;
push.aligned_resolution.y = aligned_height >> output_level;
push.inv_resolution.x = 1.0f / float(push.resolution.x);
push.inv_resolution.y = 1.0f / float(push.resolution.y);
cmd.push_constants(&push, 0, sizeof(push));
cmd.set_texture(0, 0, *views.planes[c], *mirror_repeat_sampler);
cmd.set_specialization_constant(0, true);
}
else
{
cmd.set_texture(0, 0, *component_ll_views[c][output_level - 1], *mirror_repeat_sampler);
}
cmd.set_storage_texture(0, 1, *component_layer_views[c][output_level]);
char label[64];
snprintf(label, sizeof(label), "DWT level %u, component %u", output_level, c);
cmd.begin_region(label);
cmd.dispatch((push.aligned_resolution.x + 31) / 32, (push.aligned_resolution.y + 31) / 32, 1);
cmd.end_region();
}
}
cmd.barrier(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_SAMPLED_READ_BIT);
cmd.set_specialization_constant(0, false);
}
auto end_dwt = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
device->register_time_interval("GPU", std::move(start_dwt), std::move(end_dwt), "DWT");
cmd.set_specialization_constant_mask(0);
return true;
}
size_t Encoder::Impl::compute_num_packets(const void *meta_, size_t packet_boundary) const
{
auto *meta = static_cast<const BitstreamPacket *>(meta_);
size_t num_packets = 0;
size_t size_in_packet = 0;
size_in_packet += sizeof(BitstreamSequenceHeader);
for (int i = 0; i < block_count_32x32; i++)
{
size_t packet_size = meta[i].num_words * sizeof(uint32_t);
if (!packet_size)
continue;
if (size_in_packet + packet_size > packet_boundary)
{
size_in_packet = 0;
num_packets++;
}
size_in_packet += packet_size;
}
if (size_in_packet)
num_packets++;
return num_packets;
}
#if 0
static int max_magnitude(const int (&values)[64][64], int off_x, int off_y, int w, int h)
{
int max_magnitude = 0;
for (int y = 0; y < h; y++)
for (int x = 0; x < w; x++)
max_magnitude = std::max<int>(max_magnitude, std::abs(values[off_y + y][off_x + x]));
return max_magnitude;
}
static int num_significant_values(const int (&values)[64][64], int off_x, int off_y, int w, int h)
{
int num_significant = 0;
for (int y = 0; y < h; y++)
for (int x = 0; x < w; x++)
if (values[off_y + y][off_x + x] != 0)
num_significant++;
return num_significant;
}
static bool has_significant_value(const int (&values)[64][64], int off_x, int off_y, int w, int h)
{
return max_magnitude(values, off_x, off_y, w, h) != 0;
}
template <int PacketBlockWidth, int PacketBlockHeight, int SubBlockWidth, int SubBlockHeight>
static int analyze_cost(const int (&values)[64][64])
{
int cost = 0;
for (int y = 0; y < 64; y += PacketBlockHeight)
for (int x = 0; x < 64; x += PacketBlockWidth)
if (has_significant_value(values, x, y, PacketBlockWidth, PacketBlockHeight))
cost += 8;
constexpr int BlockWidth = PacketBlockWidth / 4;
constexpr int BlockHeight = PacketBlockHeight / 4;
for (int y = 0; y < 64; y += BlockHeight)
{
for (int x = 0; x < 64; x += BlockWidth)
{
int mag = max_magnitude(values, x, y, BlockWidth, BlockHeight);
constexpr int NumSubBlocksX = BlockWidth / SubBlockWidth;
constexpr int NumSubBlocksY = BlockHeight / SubBlockHeight;
if (mag != 0)
{
cost += 2 * NumSubBlocksX * NumSubBlocksY / 8; // 2 bits to encode planes.
cost += 1; // 4 bits to encode Q_bits, 4 bits to encode quant scale per block.
}
else
{
continue;
}
constexpr int MaxDeltaQ = 3;
uint32_t q_bits = 0;
{
uint32_t num_magnitude_bits = 32 - Util::leading_zeroes(mag);
if (num_magnitude_bits > MaxDeltaQ)
q_bits = num_magnitude_bits - MaxDeltaQ;
}
int weight_bits = 0;
for (int subblock_y = 0; subblock_y < NumSubBlocksY; subblock_y++)
{
for (int subblock_x = 0; subblock_x < NumSubBlocksX; subblock_x++)
{
int subblock_mag = max_magnitude(values,
x + subblock_x * SubBlockWidth,
y + subblock_y * SubBlockHeight,
SubBlockWidth, SubBlockHeight);
uint32_t num_magnitude_bits = subblock_mag != 0 ? (32 - Util::leading_zeroes(subblock_mag)) : 0;
num_magnitude_bits = std::max(num_magnitude_bits, q_bits);
weight_bits += SubBlockWidth * SubBlockHeight * num_magnitude_bits;
}
}
int num_sign_bits = num_significant_values(values, x, y, BlockWidth, BlockHeight);
//cost += 4 * ((weight_bits + num_sign_bits + 31) / 32);
cost += ((weight_bits + num_sign_bits + 7) / 8);
}
}
return cost;
}
void Encoder::Impl::analyze_alternative_packing(const void *mapped_meta, const void *mapped_bitstream) const
{
auto *meta = static_cast<const BitstreamPacket *>(mapped_meta);
auto *bitstream = static_cast<const uint32_t *>(mapped_bitstream);
int cost_32x32_quad = 0;
int cost_32x32_horiz = 0;
int cost_32x32_vert = 0;
int cost_64x32 = 0;
int cost_64x64 = 0;
for (int component = 0; component < NumComponents; component++)
{
for (int level = 0; level < DecompositionLevels; level++)
{
// Ignore top-level CbCr when doing 420 subsampling.
if (level == 0 && component != 0)
continue;
auto band_width = wavelet_img->get_width(level);
auto band_height = wavelet_img->get_height(level);
int blocks_x_64x64 = (band_width + 63) / 64;
int blocks_y_64x64 = (band_height + 63) / 64;
for (int band = 3; band >= (level == DecompositionLevels - 1 ? 0 : 1); band--)
{
auto &block_mapping = block_meta[component][level][band];
for (int block_y = 0; block_y < blocks_y_64x64; block_y++)
{
for (int block_x = 0; block_x < blocks_x_64x64; block_x++)
{
int block_index = block_mapping.block_offset_64x64 + block_y * block_mapping.block_stride_64x64 + block_x;
if (meta[block_index].num_words == 0)
continue;
int dequant_values[64][64] = {};
const auto &mapping = block_64x64_to_16x16_mapping[block_index];
auto *header = reinterpret_cast<const BitstreamHeader *>(bitstream + meta[block_index].offset_u32);
int blocks_16x16 = int(Util::popcount32(header->ballot));
auto *control_words = bitstream + meta[block_index].offset_u32 + 2;
auto *payload_words = control_words + blocks_16x16;
Util::for_each_bit(header->ballot, [&](unsigned bit) {
int block_16x16_x = int(bit & 3);
int block_16x16_y = int(bit >> 2);
int block_16x16 = mapping.block_offset_16x16 + mapping.block_stride_16x16 * block_16x16_y + block_16x16_x;
auto &mapping_16x16 = block_meta_16x16[block_16x16];
auto q_bits = (*control_words >> 16) & 0xf;
Util::for_each_bit(mapping_16x16.block_mask, [&](unsigned bit_offset) {
auto num_planes = q_bits + ((*control_words >> bit_offset) & 0x3);
if (num_planes == 0)
return;
num_planes++;
int subblock_x = int(bit_offset >> 3u) & 1;
int subblock_y = int(bit_offset >> 1u) & 3;
int base_x = block_16x16_x * 16 + subblock_x * 8;
int base_y = block_16x16_y * 16 + subblock_y * 4;
for (int y = 0; y < 4; y++)
{
for (int x = 0; x < 8; x++)
{
int swizzled = 0;
swizzled |= ((x >> 0) & 1) << 0;
swizzled |= ((y >> 0) & 3) << 1;
swizzled |= ((x >> 1) & 3) << 3;
assert(swizzled < 32);
for (uint32_t plane = 1; plane < num_planes; plane++)
{
dequant_values[base_y + y][base_x + x] <<= 1;
dequant_values[base_y + y][base_x + x] |= int(payload_words[plane] >> swizzled) & 1;
}
if ((payload_words[0] & (1u << swizzled)) != 0)
dequant_values[base_y + y][base_x + x] *= -1;
}
}
payload_words += num_planes;
});
control_words++;
});
cost_32x32_quad += analyze_cost<32, 32, 2, 2>(dequant_values);
cost_32x32_horiz += analyze_cost<32, 32, 4, 2>(dequant_values);
cost_32x32_vert += analyze_cost<32, 32, 2, 4>(dequant_values);
cost_64x32 += analyze_cost<64, 32, 4, 4>(dequant_values);
cost_64x64 += analyze_cost<64, 64, 8, 4>(dequant_values);
auto payload_offset = payload_words - (bitstream + meta[block_index].offset_u32 + meta[block_index].num_words);
if (payload_offset != 0)
abort();
}
}
}
}
}
LOGI("32x32 (2x2) cost: %d bytes\n", cost_32x32_quad);
LOGI("32x32 (4x2) cost: %d bytes\n", cost_32x32_horiz);
LOGI("32x32 (2x4) cost: %d bytes\n", cost_32x32_vert);
LOGI("64x32 cost: %d bytes\n", cost_64x32);
LOGI("64x64 cost: %d bytes\n", cost_64x64);
}
void Encoder::Impl::report_stats(const void *mapped_meta, const void *mapped_bitstream) const
{
auto *meta = static_cast<const BitstreamPacket *>(mapped_meta);
auto *bitstream = static_cast<const uint32_t *>(mapped_bitstream);
int total_pixels = 0;
int total_words = 0;
static const char *components[] = { "Y", "Cb", "Cr" };
static const char *bands[] = { "LL", "HL", "LH", "HH" };
constexpr int MaxPlanes = 16;
int plane_histogram[MaxPlanes][256] = {};
int total_planes[MaxPlanes] = {};
for (int component = 0; component < NumComponents; component++)
{
for (int level = 0; level < DecompositionLevels; level++)
{
int total_words_in_level = 0;
// Ignore top-level CbCr when doing 420 subsampling.
if (level == 0 && component != 0)
continue;
auto band_width = wavelet_img->get_width(level);
auto band_height = wavelet_img->get_height(level);
int blocks_x_64x64 = (band_width + 63) / 64;
int blocks_y_64x64 = (band_height + 63) / 64;
for (int band = 3; band >= (level == DecompositionLevels - 1 ? 0 : 1); band--)
{
auto &block_mapping = block_meta[component][level][band];
int words = 0;
for (int block_y = 0; block_y < blocks_y_64x64; block_y++)
{
for (int block_x = 0; block_x < blocks_x_64x64; block_x++)
{
int block_index = block_mapping.block_offset_64x64 + block_y * block_mapping.block_stride_64x64 + block_x;
if (meta[block_index].num_words == 0)
continue;
const auto &mapping = block_64x64_to_16x16_mapping[block_index];
words += meta[block_index].num_words;
auto *header = reinterpret_cast<const BitstreamHeader *>(bitstream + meta[block_index].offset_u32);
int blocks_16x16 = int(Util::popcount32(header->ballot));
auto *control_words = bitstream + meta[block_index].offset_u32 + 2;
auto *payload_words = control_words + blocks_16x16;
Util::for_each_bit(header->ballot, [&](unsigned bit) {
int x = int(bit & 3);
int y = int(bit >> 2);
int block_16x16 = mapping.block_offset_16x16 + mapping.block_stride_16x16 * y + x;
auto &mapping_16x16 = block_meta_16x16[block_16x16];
auto q_bits = (*control_words >> 16) & 0xf;
Util::for_each_bit(mapping_16x16.block_mask, [&](unsigned bit_offset) {
auto num_planes = q_bits + ((*control_words >> bit_offset) & 0x3);
if (num_planes != 0)
num_planes++;
for (uint32_t j = 0; j < num_planes; j++)
{
plane_histogram[j][(payload_words[j] >> 0) & 0xff]++;
plane_histogram[j][(payload_words[j] >> 8) & 0xff]++;
plane_histogram[j][(payload_words[j] >> 16) & 0xff]++;
plane_histogram[j][(payload_words[j] >> 24) & 0xff]++;
total_planes[j] += 4;
}
payload_words += num_planes;
});
control_words++;
});
auto payload_offset = payload_words - (bitstream + meta[block_index].offset_u32 + meta[block_index].num_words);
if (payload_offset != 0)
abort();
}
}
int bytes = words * 4;
double bpp = (bytes * 8.0) / (band_width * band_height);
LOGI("%s: decomposition level %d, band %s: %.3f bpp\n",
components[component], level, bands[band], bpp);
total_words += words;
if (component == 0)
total_pixels += band_width * band_height;
total_words_in_level += words;
}
LOGI("%s: decomposition level %d: %d bytes\n", components[component], level, total_words_in_level * 4);
}
}
double plane_entropy[MaxPlanes] = {};
for (int i = 0; i < 256; i++)
{
for (int j = 0; j < MaxPlanes; j++)
{
if (total_planes[j] && plane_histogram[j][i])
{
auto p = double(plane_histogram[j][i]) / double(total_planes[j]);
plane_entropy[j] -= p * log2(p);
}
}
}
for (int i = 0; i < MaxPlanes; i++)
{
LOGI(" Plane %d entropy: %.3f %%\n", i, 100.0 * plane_entropy[i] / 8.0);
LOGI(" Plane %d bytes: %d\n", i, total_planes[i]);
}
LOGI("Overall: %.3f bpp\n", (total_words * 32.0) / total_pixels);
analyze_alternative_packing(mapped_meta, mapped_bitstream);
}
#endif
bool Encoder::Impl::validate_bitstream(
const uint32_t *bitstream_u32, const BitstreamPacket *meta, uint32_t block_index) const
{
if (meta[block_index].num_words == 0)
return true;
bitstream_u32 += meta[block_index].offset_u32;
auto *header = reinterpret_cast<const BitstreamHeader *>(bitstream_u32);
if (header->block_index != block_index)
{
LOGI("Mismatch in block index. header: %u, meta: %u\n", header->block_index, block_index);
return false;
}
if (header->payload_words != meta[block_index].num_words)
{
LOGI("Mismatch in payload words, header: %u, meta: %u\n", header->payload_words, meta[block_index].num_words);
return false;
}
// 32x32 block layout:
// N = popcount(ballot)
// N * u16 control words. 2 bits per active 4x2 block.
// N * u8 control words. 4 bits Q, 4 bits quant scale.
// Plane data: M * u8.
// Tightly packed sign data follows. Depends on number of significant values while decoding plane data.
int blocks_8x8 = int(Util::popcount32(header->ballot));
auto *bitstream_u8 = reinterpret_cast<const uint8_t *>(bitstream_u32);
auto *block_control_words = reinterpret_cast<const uint16_t *>(bitstream_u32 + 2);
auto *q_control_words = reinterpret_cast<const uint8_t *>(block_control_words + blocks_8x8);
uint32_t offset = sizeof(BitstreamHeader) + 3 * blocks_8x8;
if (offset > header->payload_words * 4)
{
LOGE("payload_words is not large enough.\n");
return false;
}
const auto &mapping = block_32x32_to_8x8_mapping[header->block_index];
bool invalid_packet = false;
int num_significant_values = 0;
Util::for_each_bit(header->ballot, [&](unsigned bit) {
int x = int(bit & 3);
int y = int(bit >> 2);
if (x < mapping.block_width_8x8 && y < mapping.block_height_8x8)
{
auto q_bits = *q_control_words & 0xf;
for (int subblock_offset = 0; subblock_offset < 16; subblock_offset += 2)
{
int num_planes = q_bits + ((*block_control_words >> subblock_offset) & 3);
int plane_significance = 0;
for (int plane = 0; plane < num_planes; plane++)
plane_significance |= bitstream_u8[offset++];
num_significant_values += int(Util::popcount32(plane_significance));
}
block_control_words++;
q_control_words++;
}
else
{
LOGE("block_index %u: 8x8 block is out of bounds. (%d, %d) >= (%d, %d)\n",
block_index, x, y, mapping.block_width_8x8, mapping.block_height_8x8);
invalid_packet = true;
}
});
if (invalid_packet)
return false;
// We expect this many sign bits to have come through.
offset += (num_significant_values + 7) / 8;
auto offset_words = (offset + 3) / 4;
if (offset_words != header->payload_words)
{
LOGE("Block index %u, offset %u != %u\n", block_index, offset_words, header->payload_words);
return false;
}
return true;
}
size_t Encoder::Impl::packetize(Packet *packets, size_t packet_boundary,
void *output_bitstream_, size_t size,
const void *mapped_meta,
const void *mapped_bitstream) const
{
size_t num_packets = 0;
size_t size_in_packet = 0;
size_t packet_offset = 0;
size_t output_offset = 0;
auto *meta = static_cast<const BitstreamPacket *>(mapped_meta);
auto *input_bitstream = static_cast<const uint32_t *>(mapped_bitstream);
auto *output_bitstream = static_cast<uint8_t *>(output_bitstream_);
(void)size;
size_t num_non_zero_blocks = 0;
for (int i = 0; i < block_count_32x32; i++)
if (meta[i].num_words != 0)
num_non_zero_blocks++;
BitstreamSequenceHeader header = {};
header.width_minus_1 = width - 1;
header.height_minus_1 = height - 1;
header.sequence = reinterpret_cast<const BitstreamHeader *>(input_bitstream + meta[0].offset_u32)->sequence;
header.extended = 1;
header.code = BITSTREAM_EXTENDED_CODE_START_OF_FRAME;
header.total_blocks = num_non_zero_blocks;
header.chroma_resolution = chroma == ChromaSubsampling::Chroma444 ? CHROMA_RESOLUTION_444 : CHROMA_RESOLUTION_420;
assert(sizeof(header) <= size);
memcpy(output_bitstream, &header, sizeof(header));
output_offset += sizeof(header);
size_in_packet += sizeof(header);
//for (int i = 0; i < block_count_32x32; i++)
// if (!validate_bitstream(input_bitstream, meta, i))
// return false;
for (int i = 0; i < block_count_32x32; i++)
{
size_t packet_size = meta[i].num_words * sizeof(uint32_t);
if (!packet_size)
continue;
if (size_in_packet + packet_size > packet_boundary)
{
packets[num_packets++] = { packet_offset, size_in_packet };
size_in_packet = 0;
packet_offset = output_offset;
}
assert(output_offset + packet_size <= size);
assert(packet_size >= sizeof(BitstreamHeader) / sizeof(uint32_t));
uint16_t block = reinterpret_cast<const BitstreamHeader *>(input_bitstream + meta[i].offset_u32)->block_index;
(void)block;
assert(block == i);
memcpy(output_bitstream + output_offset, input_bitstream + meta[i].offset_u32, packet_size);
output_offset += packet_size;
size_in_packet += packet_size;
}
if (size_in_packet)
packets[num_packets++] = { packet_offset, size_in_packet };
return num_packets;
}
bool Encoder::Impl::encode_quant_and_coding(
Vulkan::CommandBuffer &cmd, const BitstreamBuffers &buffers, float quant_scale)
{
cmd.enable_subgroup_size_control(true);
if (!quant(cmd, quant_scale))
return false;
if (!analyze_rdo(cmd))
return false;
if (!resolve_rdo(cmd, buffers.target_size))
return false;
if (!block_packing(cmd, buffers, quant_scale))
return false;
cmd.enable_subgroup_size_control(false);
return true;
}
bool Encoder::Impl::encode_pre_transformed(
Vulkan::CommandBuffer &cmd, const BitstreamBuffers &buffers, float quant_scale)
{
cmd.fill_buffer(*payload_data, 0, 0, 2 * sizeof(uint32_t));
cmd.fill_buffer(*bucket_buffer, 0);
cmd.fill_buffer(*quant_buffer, 0);
// Don't need to read the payload offset counter until quantizer.
cmd.barrier(VK_PIPELINE_STAGE_2_CLEAR_BIT, VK_ACCESS_TRANSFER_WRITE_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
VK_ACCESS_2_SHADER_STORAGE_READ_BIT | VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT);
return encode_quant_and_coding(cmd, buffers, quant_scale);
}
bool Encoder::Impl::encode(CommandBuffer &cmd, const ViewBuffers &views, const BitstreamBuffers &buffers)
{
sequence_count = (sequence_count + 1) & SequenceCountMask;
cmd.image_barrier(*wavelet_img_high_res, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT);
if (wavelet_img_low_res)
{
cmd.image_barrier(*wavelet_img_low_res, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT);
}
cmd.fill_buffer(*payload_data, 0, 0, 2 * sizeof(uint32_t));
cmd.fill_buffer(*bucket_buffer, 0);
cmd.fill_buffer(*quant_buffer, 0);
cmd.enable_subgroup_size_control(true);
if (!dwt(cmd, views))
return false;
cmd.enable_subgroup_size_control(false);
// Don't need to read the payload offset counter until quantizer.
cmd.barrier(VK_PIPELINE_STAGE_2_CLEAR_BIT, VK_ACCESS_TRANSFER_WRITE_BIT,
VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
VK_ACCESS_2_SHADER_STORAGE_READ_BIT | VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT);
return encode_quant_and_coding(cmd, buffers, -1.0f);
}
Encoder::Encoder()
{
impl.reset(new Impl);
}
bool Encoder::init(Device *device, int width_, int height_, ChromaSubsampling chroma_)
{
auto ops = device->get_device_features().vk11_props.subgroupSupportedOperations;
constexpr VkSubgroupFeatureFlags required_features =
VK_SUBGROUP_FEATURE_ARITHMETIC_BIT |
VK_SUBGROUP_FEATURE_SHUFFLE_BIT |
VK_SUBGROUP_FEATURE_SHUFFLE_RELATIVE_BIT |
VK_SUBGROUP_FEATURE_VOTE_BIT |
VK_SUBGROUP_FEATURE_BALLOT_BIT |
VK_SUBGROUP_FEATURE_CLUSTERED_BIT |
VK_SUBGROUP_FEATURE_BASIC_BIT;
if ((ops & required_features) != required_features)
{
LOGE("There are missing subgroup features. Device supports #%x, but requires #%x.\n",
ops, required_features);
return false;
}
if (!device->get_device_features().vk12_features.storageBuffer8BitAccess)
return false;
if (!device->get_device_features().enabled_features.shaderInt16)
return false;
// This should cover any HW I care about.
if (!device->supports_subgroup_size_log2(true, 4, 4) &&
!device->supports_subgroup_size_log2(true, 5, 5) &&
!device->supports_subgroup_size_log2(true, 6, 6))
return false;
return impl->init(device, width_, height_, chroma_, false);
}
bool Encoder::encode(CommandBuffer &cmd, const ViewBuffers &views, const BitstreamBuffers &buffers)
{
return impl->encode(cmd, views, buffers);
}
const Vulkan::ImageView &Encoder::get_wavelet_band(int component, int level)
{
return *impl->component_layer_views[component][level];
}
bool Encoder::encode_pre_transformed(Vulkan::CommandBuffer &cmd, const BitstreamBuffers &buffers, float quant_scale)
{
return impl->encode_pre_transformed(cmd, buffers, quant_scale);
}
size_t Encoder::compute_num_packets(const void *meta, size_t packet_boundary) const
{
return impl->compute_num_packets(meta, packet_boundary);
}
size_t Encoder::packetize(Packet *packets, size_t packet_boundary,
void *bitstream, size_t size,
const void *mapped_meta, const void *mapped_bitstream) const
{
return impl->packetize(packets, packet_boundary, bitstream, size, mapped_meta, mapped_bitstream);
}
void Encoder::report_stats(const void *, const void *) const
{
//impl->report_stats(mapped_meta, mapped_bitstream);
}
uint64_t Encoder::get_meta_required_size() const
{
return impl->block_count_32x32 * sizeof(BitstreamPacket);
}
Encoder::~Encoder()
{
}
}