lavc/apv_encode_vulkan: add a Vulkan APV encoder

This commit adds a compliant Vulkan compute APV encoder.
This commit is contained in:
Lynne
2026-06-10 15:35:56 +09:00
parent c5591a91eb
commit 876d10f683
8 changed files with 1868 additions and 0 deletions

1
configure vendored
View File

@@ -3156,6 +3156,7 @@ apng_encoder_select="deflate_wrapper llvidencdsp"
aptx_encoder_select="audio_frame_queue"
aptx_hd_encoder_select="audio_frame_queue"
apv_decoder_select="cbs_apv"
apv_vulkan_encoder_select="vulkan spirv_compiler cbs_apv"
asv1_decoder_select="blockdsp bswapdsp idctdsp"
asv1_encoder_select="aandcttables bswapdsp fdctdsp pixblockdsp"
asv2_decoder_select="blockdsp bswapdsp idctdsp"

View File

@@ -252,6 +252,7 @@ OBJS-$(CONFIG_APTX_HD_ENCODER) += aptxenc.o aptx.o
OBJS-$(CONFIG_APNG_DECODER) += png.o pngdec.o pngdsp.o
OBJS-$(CONFIG_APNG_ENCODER) += png.o pngenc.o
OBJS-$(CONFIG_APV_DECODER) += apv_decode.o apv_entropy.o apv_dsp.o
OBJS-$(CONFIG_APV_VULKAN_ENCODER) += apv_encode_vulkan.o
OBJS-$(CONFIG_ARBC_DECODER) += arbc.o
OBJS-$(CONFIG_ARGO_DECODER) += argo.o
OBJS-$(CONFIG_SSA_DECODER) += assdec.o ass.o

View File

@@ -46,6 +46,7 @@ extern const FFCodec ff_ansi_decoder;
extern const FFCodec ff_apng_encoder;
extern const FFCodec ff_apng_decoder;
extern const FFCodec ff_apv_decoder;
extern const FFCodec ff_apv_vulkan_encoder;
extern const FFCodec ff_arbc_decoder;
extern const FFCodec ff_argo_decoder;
extern const FFCodec ff_asv1_encoder;

File diff suppressed because it is too large Load Diff

View File

@@ -4,6 +4,10 @@ clean::
OBJS-$(CONFIG_APV_VULKAN_HWACCEL) += vulkan/apv_decode.comp.spv.o \
vulkan/apv_idct.comp.spv.o
OBJS-$(CONFIG_APV_VULKAN_ENCODER) += vulkan/apv_encode_dct.comp.spv.o \
vulkan/apv_encode_tiles.comp.spv.o \
vulkan/seg_gather.comp.spv.o
OBJS-$(CONFIG_FFV1_VULKAN_ENCODER) += vulkan/ffv1_enc_setup.comp.spv.o \
vulkan/ffv1_enc_reset.comp.spv.o \
vulkan/ffv1_enc_reset_golomb.comp.spv.o \

View File

@@ -0,0 +1,204 @@
/*
* Copyright (c) 2026 Lynne <dev@lynne.ee>
*
* This file is part of FFmpeg.
*
* FFmpeg is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* FFmpeg is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with FFmpeg; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
#version 460
#pragma shader_stage(compute)
#extension GL_GOOGLE_include_directive : require
#include "common.glsl"
#include "dct.glsl"
#define APV_MAX_NUM_COMP 4
#define APV_MAX_TILE_COLS 20
#define APV_MAX_TILE_ROWS 20
#define APV_MAX_TILE_COUNT (APV_MAX_TILE_COLS * APV_MAX_TILE_ROWS)
#define APV_MIN_TRANS_COEFF -32768
#define APV_MAX_TRANS_COEFF 32767
#define APV_TR_SIZE 8
#define APV_BLK_COEFFS (APV_TR_SIZE * APV_TR_SIZE)
#define APV_MB_SIZE 16
/*
* Buffer holding per-tile, per-component coefficient blocks.
* Layout (linear):
* tile_y * tile_cols * num_comp * blocks_per_tile * 64
* + tile_x * num_comp * blocks_per_tile * 64
* + comp * blocks_per_tile * 64
* + block_in_tile * 64
* + coeff_in_block
*
* blocks_per_tile is computed by the host as:
* mbs_per_tile_x * mbs_per_tile_y * blocks_per_mb[comp]
* where blocks_per_mb is 4 (luma) or 4 (chroma in 444), etc.
*/
layout (set = 0, binding = 0, scalar) writeonly buffer coeffs_buf {
int16_t coeffs[];
};
layout (set = 0, binding = 1) uniform readonly iimage2D src[];
layout (push_constant, scalar) uniform pushConstants {
ivec2 frame_dim; /* in pixels */
ivec2 tile_count; /* number of tile columns/rows */
ivec2 tile_mb_dim; /* MBs per tile (cols, rows) */
ivec2 log2_chroma_sub; /* 0/0 for 444, 1/0 for 422, etc. */
int num_comp;
int bit_depth;
/* Per-component quant scale fact/(level_scale*2^qp_shift). The encoder
* uses one QP per component, so it never varies by tile. */
float qf[APV_MAX_NUM_COMP];
/* The quantisation matrix (raster order), the same one signalled in the
* frame header. Staged into shared memory at the top of main(). */
uint8_t qmat[64];
};
/* Workgroup-local copy of qmat, filled once per workgroup (prores_raw style). */
shared uint8_t qmat_buf[64];
void main(void)
{
/* Workgroup grid:
* x: total MB columns over the frame (frame_mb_x)
* y: total MB rows over the frame (frame_mb_y)
* z: component index [0..num_comp)
*
* Local size (8, 4, 1):
* gl_LocalInvocationID.x in [0..7] = row index inside an 8x8 block
* gl_LocalInvocationID.y in [0..3] = which 8x8 block within the MB
*
* Luma and 4:4:4 chroma use all 4 blocks; 4:2:2 chroma uses 2 (a vertical
* pair) and the surplus two invocations early-out before the store.
*/
/* Stage the quantisation matrix into shared memory once, the same way
* prores_raw does -- one workgroup-wide copy instead of a push-constant
* read for every coefficient. */
if (gl_LocalInvocationIndex == 0u) {
[[unroll]]
for (uint i = 0u; i < 64u; i++)
qmat_buf[i] = qmat[i];
}
barrier();
const uint comp = gl_WorkGroupID.z;
const uint mb_x_lin = gl_WorkGroupID.x;
const uint mb_y_lin = gl_WorkGroupID.y;
/* Map workgroup to its tile (luma coords) */
const ivec2 sub_shift = (comp == 0u) ? ivec2(0) : log2_chroma_sub;
/* Compute which tile this MB belongs to in MB units */
const int tx = int(mb_x_lin) / tile_mb_dim.x;
const int ty = int(mb_y_lin) / tile_mb_dim.y;
const int tile_idx = ty * tile_count.x + tx;
const int mb_x_in_tile = int(mb_x_lin) - tx * tile_mb_dim.x;
const int mb_y_in_tile = int(mb_y_lin) - ty * tile_mb_dim.y;
/* Remainder tiles: the rightmost tile column may be narrower than
* tile_mb_dim.x. Pack block indices by the tile's ACTUAL width so the
* entropy pass sees them contiguously. */
const int frame_mb_x = (frame_dim.x + APV_MB_SIZE - 1) / APV_MB_SIZE;
const int actual_tw = min(tile_mb_dim.x, frame_mb_x - tx * tile_mb_dim.x);
const int mb_in_tile = mb_y_in_tile * actual_tw + mb_x_in_tile;
/*
* Per-MB block grid for this component: luma and 4:4:4 chroma are 2x2
* (4 blocks); 4:2:2 chroma is 1 wide x 2 tall (2 blocks). Derived from
* the chroma sub-sampling shift.
*/
const uint bw = 2u >> uint(sub_shift.x);
const uint bh = 2u >> uint(sub_shift.y);
const uint nb_blk = bw * bh;
/*
* Uniform coefficient stride: every tile-component is allocated the luma
* block count (4 per MB). A sub-sampled chroma component leaves the rest
* of its region unused, which keeps the buffer layout flat.
*/
const uint blocks_per_tile = uint(tile_mb_dim.x * tile_mb_dim.y) * 4u;
/* Block index within the workgroup -> position inside the MB */
const uint blk = gl_LocalInvocationID.y;
const uint row = gl_LocalInvocationID.x;
/* Block coordinate inside the (possibly sub-sampled) macroblock */
const ivec2 block_offset = ivec2(int(blk % bw), int(blk / bw));
/* Compute pixel coordinate for this thread's row of samples */
ivec2 mb_origin_luma = ivec2(int(mb_x_lin), int(mb_y_lin)) * APV_MB_SIZE;
ivec2 mb_origin = mb_origin_luma >> sub_shift;
ivec2 block_origin = mb_origin + block_offset * APV_TR_SIZE;
ivec2 coord = block_origin + ivec2(0, int(row));
/* Clamp to image bounds, in case frame dimensions are not aligned */
ivec2 img_dim = imageSize(src[comp]);
coord = min(coord, img_dim - ivec2(1));
const float bias = float(1 << (bit_depth - 1));
const float fact = bias;
/* Load 8 horizontal samples, subtract bias, normalize to ~[-1,1] */
[[unroll]]
for (int i = 0; i < 8; i++) {
int s = imageLoad(src[comp], coord + ivec2(i, 0)).x;
blocks[blk][row * 9u + uint(i)] = (float(s) - bias) / fact;
}
barrier();
/* Column DCT (offset varies on x-axis, traverses rows via stride 9) */
fdct8(blk, row, 9);
barrier();
/* Row DCT (offset varies on y-axis, traverses cols via stride 1) */
fdct8(blk, row * 9u, 1);
barrier();
/*
* Quantize and store. Each thread writes its row.
* level = round( fdct2d((sample - bias)/fact) * qf * 1024/qmat[i] )
* fdct8() is the exact orthonormal forward DCT, the reciprocal of the
* spec iDCT (apv_decode_transquant_c), which reconstructs
* sample - bias = (qmat[i]*level_scale*2^qp_shift / 1024)
* * iDCT_ortho(level).
* qf[comp] carries fact/(level_scale*2^qp_shift); the per-coefficient
* factor 1024/qmat[i] inverts the decoder's per-coefficient dequant.
*/
const float scale_const = qf[comp];
/* Compute coefficient base offset in the buffer */
const uint tile_lin = uint(tile_idx);
const uint coeff_base =
(tile_lin * uint(num_comp) + comp) * blocks_per_tile * 64u +
(uint(mb_in_tile) * nb_blk + blk) * 64u;
/* Surplus invocations of a sub-sampled component (blk >= nb_blk) took
* part in the barriers above but must not write any coefficients. */
if (blk < nb_blk) {
[[unroll]]
for (int i = 0; i < 8; i++) {
float v = blocks[blk][row * 9u + uint(i)];
float pf = 1024.0f / float(qmat_buf[row * 8u + uint(i)]);
int lvl = int(round(v * scale_const * pf));
lvl = clamp(lvl, APV_MIN_TRANS_COEFF, APV_MAX_TRANS_COEFF);
coeffs[coeff_base + row * 8u + uint(i)] = int16_t(lvl);
}
}
}

View File

@@ -0,0 +1,422 @@
/*
* Copyright (c) 2026 Lynne <dev@lynne.ee>
*
* This file is part of FFmpeg.
*
* FFmpeg is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* FFmpeg is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with FFmpeg; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
#version 460
#pragma shader_stage(compute)
#extension GL_GOOGLE_include_directive : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_arithmetic : require
#extension GL_KHR_shader_subgroup_rotate : require
#define PB_UNALIGNED
#include "common.glsl"
#define APV_BLK_COEFFS 64
/*
* One workgroup encodes one tile-component. The workgroup size (set by the
* host via spec constants 253/254/255) equals the number of transform blocks
* in the tile-component, so there is exactly one invocation per block.
*
* The minimum APV tile is 16x8 MBs -> 16*8*4 = 512 blocks. The buffers below
* are sized for 1024 (a 2x tile). With a guaranteed subgroup size of >= 32,
* a 1024-invocation workgroup has at most 1024/32 = 32 subgroups.
*/
#define MAX_BLOCKS 1024
#define MAX_SUBGROUPS 32
/* ff_zigzag_direct, packed: each byte is the raster index (y*8 + x). */
const uint8_t zigzag[64] = {
uint8_t( 0), uint8_t( 1), uint8_t( 8), uint8_t(16),
uint8_t( 9), uint8_t( 2), uint8_t( 3), uint8_t(10),
uint8_t(17), uint8_t(24), uint8_t(32), uint8_t(25),
uint8_t(18), uint8_t(11), uint8_t( 4), uint8_t( 5),
uint8_t(12), uint8_t(19), uint8_t(26), uint8_t(33),
uint8_t(40), uint8_t(48), uint8_t(41), uint8_t(34),
uint8_t(27), uint8_t(20), uint8_t(13), uint8_t( 6),
uint8_t( 7), uint8_t(14), uint8_t(21), uint8_t(28),
uint8_t(35), uint8_t(42), uint8_t(49), uint8_t(56),
uint8_t(57), uint8_t(50), uint8_t(43), uint8_t(36),
uint8_t(29), uint8_t(22), uint8_t(15), uint8_t(23),
uint8_t(30), uint8_t(37), uint8_t(44), uint8_t(51),
uint8_t(58), uint8_t(59), uint8_t(52), uint8_t(45),
uint8_t(38), uint8_t(31), uint8_t(39), uint8_t(46),
uint8_t(53), uint8_t(60), uint8_t(61), uint8_t(54),
uint8_t(47), uint8_t(55), uint8_t(62), uint8_t(63),
};
/* Coefficients are int16, accessed through a packed u32 view: the mask
* build streams whole words, and the few nonzero values are extracted
* on demand. */
layout (set = 0, binding = 0, scalar) readonly buffer coeffs_buf {
uint32_t coeffs32[];
};
int coeff_at(uint base16, uint i)
{
uint idx = base16 + i;
uint w = coeffs32[idx >> 1];
return (int(w << ((1u - (idx & 1u)) << 4))) >> 16;
}
/*
* Inverse zig-zag: zz_inv[raster] = scan position. Indexed only by
* unrolled-constant indices, so it folds into immediates.
*/
const uint8_t zz_inv[64] = {
uint8_t( 0), uint8_t( 1), uint8_t( 5), uint8_t( 6),
uint8_t(14), uint8_t(15), uint8_t(27), uint8_t(28),
uint8_t( 2), uint8_t( 4), uint8_t( 7), uint8_t(13),
uint8_t(16), uint8_t(26), uint8_t(29), uint8_t(42),
uint8_t( 3), uint8_t( 8), uint8_t(12), uint8_t(17),
uint8_t(25), uint8_t(30), uint8_t(41), uint8_t(43),
uint8_t( 9), uint8_t(11), uint8_t(18), uint8_t(24),
uint8_t(31), uint8_t(40), uint8_t(44), uint8_t(53),
uint8_t(10), uint8_t(19), uint8_t(23), uint8_t(32),
uint8_t(39), uint8_t(45), uint8_t(52), uint8_t(54),
uint8_t(20), uint8_t(22), uint8_t(33), uint8_t(38),
uint8_t(46), uint8_t(51), uint8_t(55), uint8_t(60),
uint8_t(21), uint8_t(34), uint8_t(37), uint8_t(47),
uint8_t(50), uint8_t(56), uint8_t(59), uint8_t(61),
uint8_t(35), uint8_t(36), uint8_t(48), uint8_t(49),
uint8_t(57), uint8_t(58), uint8_t(62), uint8_t(63),
};
/* Zig-zag-domain nonzero map of one block: bit s = coefficient at scan
* position s is nonzero. 32 sequential word loads, constant bit targets. */
uint64_t nz_mask(uint base16)
{
uint64_t mask = uint64_t(0);
uint b32 = base16 >> 1;
[[unroll]]
for (uint w = 0u; w < 32u; w++) {
uint v = coeffs32[b32 + w];
if ((v & 0xFFFFu) != 0u)
mask |= uint64_t(1) << zz_inv[2u * w];
if ((v >> 16) != 0u)
mask |= uint64_t(1) << zz_inv[2u * w + 1u];
}
return mask;
}
/* Index of the lowest set bit; mask must be nonzero. */
int findLSB64(uint64_t m)
{
u32vec2 h = unpack32(m);
return (h.x != 0u) ? findLSB(h.x) : 32 + findLSB(h.y);
}
layout (set = 0, binding = 1, scalar) writeonly buffer sizes_buf {
uint32_t tile_comp_sizes[];
};
layout (push_constant, scalar) uniform pushConstants {
u8buf bytestream; /* device address of the bytestream buffer */
ivec2 tile_count; /* number of tile columns/rows */
int num_comp;
uint slot_size; /* per-tile-component bytestream slot size */
uint comp_base; /* component index of this dispatch's z = 0 */
uint blocks_per_tile; /* uniform coeff stride, in blocks */
ivec2 frame_mb; /* frame size in MBs (luma basis) */
ivec2 tile_mb_dim; /* full-tile size in MBs */
uint blocks_per_mb; /* blocks per MB of this dispatch's components */
};
shared uint32_t sg_lasttail[MAX_SUBGROUPS]; /* last block's tail, per subgroup */
shared uint32_t sg_scan [MAX_SUBGROUPS]; /* per-subgroup scan totals */
/*
* Flush only the complete bytes held by the writer. The trailing partial
* byte is intentionally left unwritten: it is the block's tail, and the next
* block writes it (prepended to its own first byte). flush_put_bits() rounds
* the byte count up; this rounds it down.
*/
void flush_whole_bytes(inout PutBitContext pb)
{
if (pb.bit_left < BUF_BITS)
pb.bit_buf <<= pb.bit_left;
uint to_write = uint(BUF_BITS - pb.bit_left) >> 3;
u8buf bs = u8buf(pb.buf);
for (int i = 0; i < to_write; i++)
bs[i].v = BYTE_EXTRACT(pb.bit_buf, BUF_BYTES - uint8_t(1) - i);
pb.buf = uint64_t(bs) + to_write;
}
/*
* A block coder bundles the output writer, a running bit count, a rolling
* copy of the last emitted bits, and a flag selecting whether to actually
* write. The count and encode passes share the exact same code path, so the
* counted length always matches the written one.
*/
struct BlockCoder {
PutBitContext pb;
uint nbits;
uint roll;
bool wr;
};
void emit(inout BlockCoder bc, uint32_t n, uint32_t value)
{
bc.nbits += n;
bc.roll = (bc.roll << n) | value; /* keeps the last >= 8 emitted bits */
if (bc.wr)
put_bits(bc.pb, n, value);
}
/* Variable-length code from the APV spec, section 7.2.4. */
void write_vlc(inout BlockCoder bc, int kParam, int symbolVal)
{
int threshold1 = 1 << kParam;
int threshold2 = threshold1 << 1;
if (symbolVal < threshold1) {
emit(bc, 1u, 1u);
if (kParam > 0)
emit(bc, uint(kParam), uint(symbolVal));
} else if (symbolVal < threshold2) {
emit(bc, 2u, 0u);
if (kParam > 0)
emit(bc, uint(kParam), uint(symbolVal - threshold1));
} else {
emit(bc, 2u, 1u);
int adjusted = symbolVal - threshold1;
int n_plus_k = findMSB(adjusted);
int n = n_plus_k - kParam;
if (n > 0)
emit(bc, uint(n), 0u);
emit(bc, 1u, 1u);
int residual = adjusted - (1 << n_plus_k);
if (n_plus_k > 0)
emit(bc, uint(n_plus_k), uint(residual));
}
}
/* abs() of the first non-zero AC coefficient in zig-zag order, 0 if none. */
int first_ac_level(uint base16, uint64_t mask)
{
uint64_t mac = mask >> 1;
if (mac == uint64_t(0))
return 0;
int sp = 1 + findLSB64(mac);
return abs(coeff_at(base16, uint(zigzag[sp])));
}
/* Entropy-code one transform block given its predicted context. The nonzero
* map drives the zero-runs, so only the nonzero values are ever loaded. */
void process_block(inout BlockCoder bc, uint base16, uint64_t mask,
int prev_dc, int prev_k_dc, int prev_1st_ac)
{
/* DC */
int dc = coeff_at(base16, 0u);
int dc_diff = dc - prev_dc;
int abs_dc = abs(dc_diff);
write_vlc(bc, prev_k_dc, abs_dc);
if (abs_dc != 0)
emit(bc, 1u, dc_diff < 0 ? 1u : 0u);
/* AC, zig-zag from position 1. m's bit 0 = scan position scan_pos. */
uint64_t m = mask >> 1;
int scan_pos = 1;
int prev_level = prev_1st_ac;
int prev_run = 0;
while (scan_pos < APV_BLK_COEFFS) {
int run = (m == uint64_t(0)) ? APV_BLK_COEFFS - scan_pos
: findLSB64(m);
int k_run = clamp(prev_run >> 2, 0, 2);
write_vlc(bc, k_run, run);
scan_pos += run;
prev_run = run;
if (scan_pos < APV_BLK_COEFFS) {
int level = coeff_at(base16, uint(zigzag[scan_pos]));
int abs_level_m1 = abs(level) - 1;
int k_level = clamp(prev_level >> 2, 0, 4);
write_vlc(bc, k_level, abs_level_m1);
emit(bc, 1u, level < 0 ? 1u : 0u);
prev_level = abs_level_m1 + 1;
scan_pos++;
m >>= run + 1;
}
}
}
/* DC predictor context for block 'idx' -- pure function of preceding DCs. */
void dc_context(uint coeff_base, uint idx, out int prev_dc, out int prev_k_dc)
{
if (idx == 0u) {
prev_dc = 0;
prev_k_dc = 5;
return;
}
uint base = coeff_base + idx * APV_BLK_COEFFS;
int dc_m1 = coeff_at(base - APV_BLK_COEFFS, 0u);
int dc_m2 = (idx == 1u) ? 0 : coeff_at(base - 2u * APV_BLK_COEFFS, 0u);
prev_dc = dc_m1;
prev_k_dc = min(abs(dc_m1 - dc_m2) >> 1, 5);
}
void main(void)
{
const uint b = gl_LocalInvocationID.x;
const uint comp = comp_base + gl_WorkGroupID.z;
const uint tile_idx = gl_WorkGroupID.y * uint(tile_count.x) + gl_WorkGroupID.x;
const uint tile_comp = tile_idx * uint(num_comp) + comp;
/*
* Remainder tiles: the rightmost column / bottom row of tiles may be
* smaller than the full tile the workgroup was sized for. Invocations at
* b >= nb stay for the barriers and scans (contributing zeros) but never
* read coefficients or write bitstream. Active blocks are always the
* contiguous prefix [0, nb), so the tail-handoff chain below never
* crosses an inactive->active boundary.
*/
const int actual_tw = min(tile_mb_dim.x,
frame_mb.x - int(gl_WorkGroupID.x) * tile_mb_dim.x);
const int actual_th = min(tile_mb_dim.y,
frame_mb.y - int(gl_WorkGroupID.y) * tile_mb_dim.y);
const uint nb = uint(actual_tw * actual_th) * blocks_per_mb;
/*
* This tile-component's coefficient region. The coeff buffer uses a
* uniform per-tile-component stride (the luma block count); a chroma
* tile-component simply reads fewer blocks from its region.
*/
const uint coeff_base = tile_comp * blocks_per_tile * APV_BLK_COEFFS;
const uint blk = coeff_base + b * APV_BLK_COEFFS;
int prev_dc = 0, prev_k_dc = 5;
uint64_t mask = uint64_t(0);
if (b < nb) {
dc_context(coeff_base, b, prev_dc, prev_k_dc);
mask = nz_mask(blk);
}
/*
* Scan #1: prev_1st_ac_level is the first-AC level of the most recent
* *non-empty* block before this one. Pack (blockIndex+1, level) into a
* key (empty blocks -> 0) so a plain Max picks the highest-index prior
* non-empty block. Two levels: subgroup scan, then across subgroups.
*/
int my_first_ac = (b < nb) ? first_ac_level(blk, mask) : 0;
uint key = (my_first_ac == 0) ? 0u
: (((b + 1u) << 16) | uint(my_first_ac));
uint key_excl = subgroupExclusiveMax(key);
uint key_total = subgroupMax(key);
if (subgroupElect())
sg_scan[gl_SubgroupID] = key_total;
barrier();
uint carry = key_excl;
for (uint i = 0u; i < gl_NumSubgroups; i++) {
if (i >= gl_SubgroupID)
break;
carry = max(carry, sg_scan[i]);
}
int prev_1st_ac = (carry == 0u) ? 0 : int(carry & 0xFFFFu);
barrier(); /* sg_scan is reused by scan #2 */
/*
* Count pass: measure this block's coded length, and -- for free, since
* we walk every emitted bit anyway -- keep a rolling copy of the last
* bits in cnt.roll.
*/
BlockCoder cnt;
cnt.nbits = 0u;
cnt.roll = 0u;
cnt.wr = false;
if (b < nb)
process_block(cnt, blk, mask, prev_dc, prev_k_dc, prev_1st_ac);
/*
* Scan #2: exclusive prefix sum of the bit counts gives each block's
* start offset (in bits) within the tile-component bitstream.
*/
uint bits_excl = subgroupExclusiveAdd(cnt.nbits);
uint bits_total = subgroupAdd(cnt.nbits);
if (subgroupElect())
sg_scan[gl_SubgroupID] = bits_total;
barrier();
uint off = bits_excl;
for (uint i = 0u; i < gl_NumSubgroups; i++) {
if (i >= gl_SubgroupID)
break;
off += sg_scan[i];
}
/*
* This block's tail = its last (end_bit & 7) bits. The next block
* prepends these so its own writer starts on a byte boundary; that makes
* every block write a disjoint run of whole bytes -- no atomics, no
* shared bytes. end_bit is known only now (after scan #2).
*
* The predecessor's tail is just the left-neighbour lane, so a subgroup
* rotate fetches it from a register; only at a subgroup boundary (lane 0)
* does it fall back to shared memory -- the previous subgroup's last block.
*/
uint end_bit = off + cnt.nbits;
uint tail_n = end_bit & 7u;
uint my_tail = cnt.roll & ((1u << tail_n) - 1u);
uint prev_tail = subgroupRotate(my_tail, gl_SubgroupSize - 1u);
if (gl_SubgroupInvocationID == gl_SubgroupSize - 1u)
sg_lasttail[gl_SubgroupID] = my_tail;
barrier();
if (gl_SubgroupInvocationID == 0u && gl_SubgroupID > 0u)
prev_tail = sg_lasttail[gl_SubgroupID - 1u];
/* Write pass: each block writes bytes [off>>3, end_bit>>3). Inactive
* remainder-tile invocations write nothing. */
if (b >= nb)
return;
BlockCoder enc;
init_put_bits(enc.pb,
OFFBUF(u8buf, bytestream, tile_comp * slot_size + (off >> 3u)),
uint64_t(slot_size));
enc.nbits = 0u;
enc.roll = 0u;
enc.wr = true;
/* Prepend the predecessor's tail so this block's first byte comes out
* complete. (off & 7 == 0 for block 0, which has no predecessor.) */
if (b > 0u)
put_bits(enc.pb, off & 7u, prev_tail);
process_block(enc, blk, mask, prev_dc, prev_k_dc, prev_1st_ac);
if (b == nb - 1u) {
/* Last block: no successor, so flush the trailing partial byte too
* (the tile-component is byte-aligned, zero-padded), and record the
* total size. */
flush_put_bits(enc.pb);
tile_comp_sizes[tile_comp] = (end_bit + 7u) >> 3u;
} else {
/* Leave the trailing partial byte for the next block to prepend. */
flush_whole_bytes(enc.pb);
}
}

View File

@@ -0,0 +1,131 @@
/*
* Copyright (c) 2026 Lynne <dev@lynne.ee>
*
* This file is part of FFmpeg.
*
* FFmpeg is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* FFmpeg is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with FFmpeg; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
#version 460
#pragma shader_stage(compute)
#extension GL_GOOGLE_include_directive : require
#include "common.glsl"
/*
* Generic segment gather.
*
* A parallel encoder emits each output segment (an APV tile-component, an
* FFv1 slice, ...) into its own fixed-stride, device-local slot, because no
* encoder workgroup knows the others' final sizes. Run afterwards, this
* shader prefix-sums the per-segment sizes and packs the segments back to
* back into one contiguous, host-visible buffer -- so the device->host
* transfer is a single coalesced stream rather than a scattered dribble.
*
* One workgroup per segment. Inputs: the sparse slot buffer, the per-segment
* sizes, and the slot stride.
*/
layout (set = 0, binding = 0, scalar) readonly buffer sizes_buf {
uint32_t seg_sizes[];
};
layout (push_constant, scalar) uniform pushConstants {
u8buf sparse; /* device-local: one slot per segment */
u8buf compacted; /* host-visible: contiguous output */
uint slot_size; /* stride between sparse slots */
};
shared uint s_dst_off;
/*
* The 16-byte window starting `sh` bytes (0..15) into the 32-byte pair
* (lo, hi). This bridges the source/destination misalignment so both the
* loads and the stores in main() stay aligned. sh == 0 returns lo unchanged.
*/
u32vec4 funnel(u32vec4 lo, u32vec4 hi, uint sh)
{
if (sh == 0u)
return lo;
uint s[8] = uint[8](lo.x, lo.y, lo.z, lo.w, hi.x, hi.y, hi.z, hi.w);
uint uw = sh >> 2u; /* whole uints into the window */
uint bb = (sh & 3u) << 3u; /* remaining sub-uint shift, in bits */
if (bb == 0u)
return u32vec4(s[uw], s[uw + 1u], s[uw + 2u], s[uw + 3u]);
return u32vec4((s[uw ] >> bb) | (s[uw + 1u] << (32u - bb)),
(s[uw + 1u] >> bb) | (s[uw + 2u] << (32u - bb)),
(s[uw + 2u] >> bb) | (s[uw + 3u] << (32u - bb)),
(s[uw + 3u] >> bb) | (s[uw + 4u] << (32u - bb)));
}
void main(void)
{
const uint seg = gl_WorkGroupID.x;
const uint b = gl_LocalInvocationID.x;
const uint wg = gl_WorkGroupSize.x;
/*
* Destination offset: the sum of all preceding segment sizes. The output
* is packed tight -- segments back to back -- so it is usable directly as
* the assembled bitstream.
*/
if (b == 0u) {
uint o = 0u;
for (uint i = 0u; i < seg; i++)
o += seg_sizes[i];
s_dst_off = o;
}
barrier();
const uint n = seg_sizes[seg];
const uint64_t src_base = uint64_t(sparse) + seg * slot_size;
const uint64_t dst_base = uint64_t(compacted) + s_dst_off;
u8buf src8 = u8buf(src_base);
u8buf dst8 = u8buf(dst_base);
/*
* The destination is tightly packed, so it starts at an arbitrary byte
* offset; the source slots are 16-aligned. Copy a short head byte-wise to
* bring the destination to a 16-byte boundary, then the bulk as aligned
* u32vec4 stores -- the wide PCIe transactions the gather needs -- each
* fed from two aligned source loads via funnel(). A 16-aligned source
* makes that shift exactly `head`.
*/
const uint head = min((16u - (uint(dst_base) & 15u)) & 15u, n);
for (uint i = b; i < head; i += wg)
dst8[i].v = src8[i].v;
const uint rem = n - head;
const uint nbody = rem >> 4u;
/*
* The last body word's second source load would read past the slot;
* drop it from the loop and let the byte-wise tail below cover it.
*/
const uint nsafe = nbody > 0u ? nbody - 1u : 0u;
u32vec4buf srcw = u32vec4buf(src_base);
u32vec4buf dstw = u32vec4buf(dst_base + head);
for (uint w = b; w < nsafe; w += wg)
dstw[w].v = funnel(srcw[w].v, srcw[w + 1u].v, head);
for (uint i = head + (nsafe << 4u) + b; i < n; i += wg)
dst8[i].v = src8[i].v;
}