PR #25174 opened by Lynne URL: https://code.ffmpeg.org/FFmpeg/FFmpeg/pulls/25174 Patch URL: https://code.ffmpeg.org/FFmpeg/FFmpeg/pulls/25174.patch
12k FPS on 1080p 444, 718 FPS on 8K 444, on 6900XT. I don't think any ASIC-based implementation in production can beat it. >From 619e329813b078d0c75b88668007eda796935d84 Mon Sep 17 00:00:00 2001 From: Lynne <[email protected]> Date: Sun, 18 Oct 2026 17:39:41 +0900 Subject: [PATCH 1/2] lavc/mjpegenc_common: export the comment and Huffman table writers For the Vulkan MJPEG encoder to write its headers with them. --- libavcodec/mjpegenc_common.c | 26 +++++++++++++------------- libavcodec/mjpegenc_common.h | 4 ++++ 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/libavcodec/mjpegenc_common.c b/libavcodec/mjpegenc_common.c index bc58ca9903..c9fdf94f1e 100644 --- a/libavcodec/mjpegenc_common.c +++ b/libavcodec/mjpegenc_common.c @@ -37,8 +37,8 @@ #include "version.h" /* table_class: 0 = DC coef, 1 = AC coefs */ -static int put_huffman_table(PutBitContext *p, int table_class, int table_id, - const uint8_t *bits_table, const uint8_t *value_table) +int ff_mjpeg_put_huffman_table(PutBitContext *p, int table_class, int table_id, + const uint8_t *bits_table, const uint8_t *value_table) { int n = 0; @@ -106,24 +106,24 @@ static void jpeg_table_header(AVCodecContext *avctx, PutBitContext *p, // Only MJPEG can have a variable Huffman variable. All other // formats use the default Huffman table. if (m && m->huffman == HUFFMAN_TABLE_OPTIMAL) { - size += put_huffman_table(p, 0, 0, m->bits_dc_luminance, + size += ff_mjpeg_put_huffman_table(p,0, 0, m->bits_dc_luminance, m->val_dc_luminance); - size += put_huffman_table(p, 0, 1, m->bits_dc_chrominance, + size += ff_mjpeg_put_huffman_table(p,0, 1, m->bits_dc_chrominance, m->val_dc_chrominance); - size += put_huffman_table(p, 1, 0, m->bits_ac_luminance, + size += ff_mjpeg_put_huffman_table(p,1, 0, m->bits_ac_luminance, m->val_ac_luminance); - size += put_huffman_table(p, 1, 1, m->bits_ac_chrominance, + size += ff_mjpeg_put_huffman_table(p,1, 1, m->bits_ac_chrominance, m->val_ac_chrominance); } else { - size += put_huffman_table(p, 0, 0, ff_mjpeg_bits_dc_luminance, + size += ff_mjpeg_put_huffman_table(p,0, 0, ff_mjpeg_bits_dc_luminance, ff_mjpeg_val_dc); - size += put_huffman_table(p, 0, 1, ff_mjpeg_bits_dc_chrominance, + size += ff_mjpeg_put_huffman_table(p,0, 1, ff_mjpeg_bits_dc_chrominance, ff_mjpeg_val_dc); - size += put_huffman_table(p, 1, 0, ff_mjpeg_bits_ac_luminance, + size += ff_mjpeg_put_huffman_table(p,1, 0, ff_mjpeg_bits_ac_luminance, ff_mjpeg_val_ac_luminance); - size += put_huffman_table(p, 1, 1, ff_mjpeg_bits_ac_chrominance, + size += ff_mjpeg_put_huffman_table(p,1, 1, ff_mjpeg_bits_ac_chrominance, ff_mjpeg_val_ac_chrominance); } AV_WB16(ptr, size); @@ -202,8 +202,8 @@ int ff_mjpeg_add_gain_map_size(AVCodecContext *avctx, struct MJpegContext *m, return 0; } -static void jpeg_put_comments(AVCodecContext *avctx, PutBitContext *p, - const AVFrame *frame, const struct MJpegContext *m) +void ff_mjpeg_put_comments(AVCodecContext *avctx, PutBitContext *p, + const AVFrame *frame, const struct MJpegContext *m) { const AVFrameSideData *sd = NULL; int size; @@ -371,7 +371,7 @@ void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb, if (avctx->codec_id == AV_CODEC_ID_AMV) return; - jpeg_put_comments(avctx, pb, frame, m); + ff_mjpeg_put_comments(avctx, pb, frame, m); chroma_matrix = !lossless && !!memcmp(luma_intra_matrix, chroma_intra_matrix, diff --git a/libavcodec/mjpegenc_common.h b/libavcodec/mjpegenc_common.h index 20e3902c03..ebca39967b 100644 --- a/libavcodec/mjpegenc_common.h +++ b/libavcodec/mjpegenc_common.h @@ -32,6 +32,8 @@ int ff_mjpeg_add_icc_profile_size(AVCodecContext *avctx, const AVFrame *frame, size_t *max_pkt_size); int ff_mjpeg_add_gain_map_size(AVCodecContext *avctx, struct MJpegContext *m, const AVFrame *frame, size_t *max_pkt_size); +void ff_mjpeg_put_comments(AVCodecContext *avctx, PutBitContext *p, + const AVFrame *frame, const struct MJpegContext *m); void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb, const AVFrame *frame, const struct MJpegContext *m, const uint8_t intra_matrix_permutation[64], @@ -40,6 +42,8 @@ void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb, const uint16_t chroma_intra_matrix[64], int use_slices); void ff_mjpeg_encode_picture_trailer(PutBitContext *pb, int header_bits); +int ff_mjpeg_put_huffman_table(PutBitContext *p, int table_class, int table_id, + const uint8_t *bits_table, const uint8_t *value_table); void ff_mjpeg_escape_FF(PutBitContext *pb, int start); void ff_mjpeg_build_huffman_codes(uint8_t *huff_size, uint16_t *huff_code, const uint8_t *bits_table, -- 2.54.0 >From bd5c9be6f5e14ea23f4ebc2a82e96cb2c62d2e15 Mon Sep 17 00:00:00 2001 From: Lynne <[email protected]> Date: Sun, 18 Oct 2026 17:39:41 +0900 Subject: [PATCH 2/2] lavc/mjpegenc_vulkan: add a Vulkan compute MJPEG encoder This commit adds a compute-based baseline MJPEG encoder. Blocks are transformed and quantised with the same matrices and -q:v scale as the software encoder. Each restart interval is Huffman coded by one workgroup, with byte stuffing and restart markers done in the same pass, and the intervals are then packed with the shared segment gather. By default, as in the software encoder, optimal Huffman tables are built on the GPU from a histogram of each frame. Gray, 4:2:0, 4:2:2 and 4:4:4 are supported, planar or semi-planar, 8-bit. --- Changelog | 1 + configure | 1 + libavcodec/Makefile | 2 + libavcodec/allcodecs.c | 1 + libavcodec/mjpegenc_vulkan.c | 951 ++++++++++++++++++ libavcodec/vulkan/Makefile | 5 + libavcodec/vulkan/mjpeg_encode_block.glsl | 162 +++ libavcodec/vulkan/mjpeg_encode_dct.comp.glsl | 157 +++ .../vulkan/mjpeg_encode_histogram.comp.glsl | 66 ++ .../vulkan/mjpeg_encode_huffman.comp.glsl | 212 ++++ .../vulkan/mjpeg_encode_tables.comp.glsl | 227 +++++ 11 files changed, 1785 insertions(+) create mode 100644 libavcodec/mjpegenc_vulkan.c create mode 100644 libavcodec/vulkan/mjpeg_encode_block.glsl create mode 100644 libavcodec/vulkan/mjpeg_encode_dct.comp.glsl create mode 100644 libavcodec/vulkan/mjpeg_encode_histogram.comp.glsl create mode 100644 libavcodec/vulkan/mjpeg_encode_huffman.comp.glsl create mode 100644 libavcodec/vulkan/mjpeg_encode_tables.comp.glsl diff --git a/Changelog b/Changelog index 32e7c8bfe7..00e9b1374a 100644 --- a/Changelog +++ b/Changelog @@ -19,6 +19,7 @@ version <next>: - Panasonic RW2 decoder - xHE-AAC encoding via libfdk-xhe-aac - jpegapp bitstream filter +- MJPEG Vulkan encoder version 9.0: diff --git a/configure b/configure index c97ddb186f..f51d255d5e 100755 --- a/configure +++ b/configure @@ -3203,6 +3203,7 @@ metasound_decoder_select="lsp sinewin" mimic_decoder_select="blockdsp bswapdsp hpeldsp idctdsp" mjpeg_decoder_select="blockdsp idctdsp jpegtables" mjpeg_encoder_select="jpegtables mpegvideoenc" +mjpeg_vulkan_encoder_select="jpegtables vulkan spirv_compiler vulkan_seg_gather" mjpegb_decoder_select="hpeldsp mjpeg_decoder" mlp_decoder_select="mlp_parser" mlp_encoder_select="lpc audio_frame_queue" diff --git a/libavcodec/Makefile b/libavcodec/Makefile index 249b0fe0b8..fa694a84c3 100644 --- a/libavcodec/Makefile +++ b/libavcodec/Makefile @@ -538,6 +538,8 @@ OBJS-$(CONFIG_MJPEG_DECODER) += mjpegdec.o mjpegdec_common.o OBJS-$(CONFIG_MJPEG_QSV_DECODER) += qsvdec.o OBJS-$(CONFIG_MJPEG_ENCODER) += mjpegenc.o mjpegenc_common.o \ mjpegenc_huffman.o gain_map.o +OBJS-$(CONFIG_MJPEG_VULKAN_ENCODER) += mjpegenc_vulkan.o mjpegenc_common.o \ + gain_map.o mpeg12data.o OBJS-$(CONFIG_MJPEGB_DECODER) += mjpegbdec.o OBJS-$(CONFIG_MJPEG_CUVID_DECODER) += cuviddec.o OBJS-$(CONFIG_MJPEG_QSV_ENCODER) += qsvenc_jpeg.o diff --git a/libavcodec/allcodecs.c b/libavcodec/allcodecs.c index c344cb539e..28f10077b5 100644 --- a/libavcodec/allcodecs.c +++ b/libavcodec/allcodecs.c @@ -198,6 +198,7 @@ extern const FFCodec ff_mdec_decoder; extern const FFCodec ff_media100_decoder; extern const FFCodec ff_mimic_decoder; extern const FFCodec ff_mjpeg_encoder; +extern const FFCodec ff_mjpeg_vulkan_encoder; extern const FFCodec ff_mjpeg_decoder; extern const FFCodec ff_mjpegb_decoder; extern const FFCodec ff_mmvideo_decoder; diff --git a/libavcodec/mjpegenc_vulkan.c b/libavcodec/mjpegenc_vulkan.c new file mode 100644 index 0000000000..05981944a6 --- /dev/null +++ b/libavcodec/mjpegenc_vulkan.c @@ -0,0 +1,951 @@ +/* + * Copyright (c) 2026 Lynne <[email protected]> + * + * 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 + */ + +#include "libavutil/intreadwrite.h" +#include "libavutil/mem.h" +#include "libavutil/opt.h" +#include "libavutil/pixdesc.h" +#include "libavutil/vulkan.h" + +#include "avcodec.h" +#include "codec_internal.h" +#include "encode.h" +#include "hwconfig.h" +#include "jpegtables.h" +#include "mathops.h" +#include "mjpeg.h" +#include "mjpegenc.h" +#include "mjpegenc_common.h" +#include "mpeg12data.h" +#include "put_bits.h" +#include "vulkan_video.h" + +extern const unsigned char ff_mjpeg_encode_dct_comp_spv_data[]; +extern const unsigned int ff_mjpeg_encode_dct_comp_spv_len; + +extern const unsigned char ff_mjpeg_encode_histogram_comp_spv_data[]; +extern const unsigned int ff_mjpeg_encode_histogram_comp_spv_len; + +extern const unsigned char ff_mjpeg_encode_tables_comp_spv_data[]; +extern const unsigned int ff_mjpeg_encode_tables_comp_spv_len; + +extern const unsigned char ff_mjpeg_encode_huffman_comp_spv_data[]; +extern const unsigned int ff_mjpeg_encode_huffman_comp_spv_len; + +#define MAX_COMPONENTS 3 +#define MAX_MCU_BLOCKS 10 + +/* DC luma, DC chroma, AC luma, AC chroma, 256 entries each */ +#define NB_TABLES 4 +#define NB_TAB_ENTRIES (NB_TABLES*256) + +/* Blocks per workgroup of the transform */ +#define DCT_WG_BLOCKS 8 + +/* A block takes at most a 16-bit code and 11 extra bits for its DC + * difference, and 16 + 10 bits for each of its 63 AC coefficients. Byte + * stuffing can double that. */ +#define MAX_BLOCK_BYTES (2*((16 + 11 + 63*(16 + 10) + 7) / 8)) + +/* Slots start out sized for this many bytes per block, which few intervals + * need, and grow when one overflows */ +#define INIT_BLOCK_BYTES 64 + +/* Blocks per restart interval, when left to pick it */ +#define AUTO_SEG_BLOCKS 512 + +/* Size of the headers, without any ICC profile or gain map */ +#define MAX_BASE_HEADER_SIZE 1024 + +/* Bytes are packed in words, see mjpeg_encode_dct.comp.glsl */ +#define PACK8(a, b, c, d) ((a) | (b) << 8 | (c) << 16 | (uint32_t)(d) << 24) + +typedef struct DCTPushData { + int32_t plane_size[MAX_COMPONENTS][2]; + uint32_t comp_info[MAX_COMPONENTS]; + uint32_t comp_offset[MAX_COMPONENTS]; + uint32_t mcu_width; + uint32_t blocks_per_mcu; + uint32_t nb_blocks; + uint32_t comp; +} DCTPushData; + +typedef struct BlockPushData { + VkDeviceAddress bytestream; + uint32_t slot_size; + uint32_t nb_blocks; + uint32_t nb_segs; + uint32_t mcu_blocks[MAX_MCU_BLOCKS]; +} BlockPushData; + +/* Huffman tables, as written in the DHT segment */ +typedef struct DHTData { + uint8_t bits[NB_TABLES][17]; /* codes per length, from index 1 */ + uint8_t vals[NB_TABLES][256]; +} DHTData; + +typedef struct VulkanEncodeMJPEGFrameData { + /* Each restart interval is coded into its own slot, and the slots are + * gathered into the output. Both grow when a slot overflows. */ + FFVkBuffer slots; + FFVkBuffer out; + size_t slot_size; + + /* The source frame, kept for the side data written in the headers, and + * to code it again if an interval overflows its slot */ + AVFrame *frame; +} VulkanEncodeMJPEGFrameData; + +typedef struct VulkanEncodeMJPEGContext { + const AVClass *class; + + FFVulkanContext s; + AVVulkanDeviceQueueFamily *qf; + FFVkExecPool exec_pool; + FFVkEncodeLoop loop; + VulkanEncodeMJPEGFrameData *exec_ctx_info; + + FFVulkanShader shd_dct; + FFVulkanShader shd_histogram; + FFVulkanShader shd_tables; + FFVulkanShader shd_huffman; + FFVulkanShader shd_gather; + + FFVkBuffer quant_buf; + FFVkBuffer codes_buf; /* the standard tables */ + DHTData fixed_dht; + + /* Per context: the coefficients, then with optimal tables the codes + * and the histogram they are built from, on the device; the sizes of + * the intervals, the packed size and the size of the largest interval + * that overflowed its slot, then the DHT, on the host */ + FFVkBuffer dev; + FFVkBuffer host; + size_t dev_stride, codes_offset; + size_t host_stride, dht_offset; + + DCTPushData dct_push; + BlockPushData block_push; + + /* Frame layout */ + int nb_components; + int hsample[MAX_COMPONENTS]; + int vsample[MAX_COMPONENTS]; + int mcu_width, mcu_height; + int blocks_per_mcu; + int nb_blocks; + int seg_mcus; + int seg_blocks; + int nb_segs; + size_t slot_size; + size_t max_slot_size; + + int qscale; + int nb_quant_matrices; + uint8_t quant_matrix[2][64]; /* raster order */ + + /* Options */ + int huffman; + int restart_interval; + int async_depth; +} VulkanEncodeMJPEGContext; + +static void write_headers(AVCodecContext *avctx, PutBitContext *pb, + const AVFrame *frame, const DHTData *dht) +{ + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + const int nb_ids = 1 + (ev->nb_components > 1); /* luma, chroma */ + + put_marker(pb, SOI); + ff_mjpeg_put_comments(avctx, pb, frame, NULL); + + put_marker(pb, DQT); + put_bits(pb, 16, 2 + ev->nb_quant_matrices*65); + for (int t = 0; t < ev->nb_quant_matrices; t++) { + put_bits(pb, 8, t); /* 8-bit precision */ + for (int i = 0; i < 64; i++) + put_bits(pb, 8, ev->quant_matrix[t][ff_zigzag_direct[i]]); + } + + put_marker(pb, DRI); + put_bits(pb, 16, 4); + put_bits(pb, 16, ev->seg_mcus); + + put_marker(pb, DHT); + flush_put_bits(pb); + uint8_t *size = put_bits_ptr(pb); + put_bits(pb, 16, 0); /* patched below */ + int n = 2; + for (int t = 0; t < NB_TABLES; t++) + if ((t & 1) < nb_ids) + n += ff_mjpeg_put_huffman_table(pb, t >> 1, t & 1, + dht->bits[t], dht->vals[t]); + AV_WB16(size, n); + + put_marker(pb, SOF0); + put_bits(pb, 16, 8 + 3*ev->nb_components); + put_bits(pb, 8, 8); /* 8 bits/component */ + put_bits(pb, 16, avctx->height); + put_bits(pb, 16, avctx->width); + put_bits(pb, 8, ev->nb_components); + for (int c = 0; c < ev->nb_components; c++) { + put_bits(pb, 8, c + 1); + put_bits(pb, 4, ev->hsample[c]); + put_bits(pb, 4, ev->vsample[c]); + put_bits(pb, 8, c && ev->nb_quant_matrices > 1); + } + + put_marker(pb, SOS); + put_bits(pb, 16, 6 + 2*ev->nb_components); + put_bits(pb, 8, ev->nb_components); + for (int c = 0; c < ev->nb_components; c++) { + put_bits(pb, 8, c + 1); + put_bits(pb, 4, !!c); /* DC table */ + put_bits(pb, 4, !!c); /* AC table */ + } + put_bits(pb, 8, 0); /* Ss */ + put_bits(pb, 8, 63); /* Se */ + put_bits(pb, 8, 0); /* Ah, Al */ +} + +static void compute_barrier(FFVulkanFunctions *vk, FFVkExecContext *exec, + VkPipelineStageFlags2 src_stage, + VkAccessFlags2 src_access) +{ + vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .pMemoryBarriers = &(VkMemoryBarrier2) { + .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER_2, + .srcStageMask = src_stage, + .srcAccessMask = src_access, + .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + .dstAccessMask = VK_ACCESS_2_SHADER_READ_BIT | + VK_ACCESS_2_SHADER_WRITE_BIT, + }, + .memoryBarrierCount = 1, + }); +} + +/* The slots of a context, and the output they are gathered into, for the + * current slot size */ +static int alloc_slots(VulkanEncodeMJPEGContext *ev, + VulkanEncodeMJPEGFrameData *fd) +{ + int err; + size_t size = ev->nb_segs*ev->slot_size; + uint8_t *mem; + + ff_vk_free_buf(&ev->s, &fd->slots); + ff_vk_free_buf(&ev->s, &fd->out); + + RET(ff_vk_create_buf(&ev->s, &fd->slots, size, NULL, NULL, + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | + VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)); + RET(ff_vk_create_buf(&ev->s, &fd->out, size, NULL, NULL, + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | + VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | + ev->s.host_cached_flag)); + RET(ff_vk_map_buffer(&ev->s, &fd->out, &mem, 0)); + fd->slot_size = ev->slot_size; + +fail: + return err; +} + +static int record_frame(AVCodecContext *avctx, FFVkExecContext *exec) +{ + int err; + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + FFVulkanFunctions *vk = &ev->s.vkfn; + VulkanEncodeMJPEGFrameData *fd = exec->opaque; + AVFrame *frame = fd->frame; + const int optimal = ev->huffman == HUFFMAN_TABLE_OPTIMAL; + const size_t dev_base = exec->idx*ev->dev_stride; + const size_t host_base = exec->idx*ev->host_stride; + FFVkBuffer *codes_buf = &ev->codes_buf; + size_t codes_offset = 0; + VkImageView views[AV_NUM_DATA_POINTERS]; + VkImageMemoryBarrier2 img_bar[AV_NUM_DATA_POINTERS]; + int nb_img_bar = 0; + + if (fd->slot_size != ev->slot_size) { + err = alloc_slots(ev, fd); + if (err < 0) + return err; + } + + RET(ff_vk_exec_start(&ev->s, exec)); + + /* The size of the largest overflow, and the histogram, start at 0 */ + vk->CmdFillBuffer(exec->buf, ev->host.buf, + host_base + (ev->nb_segs + 1)*sizeof(uint32_t), + sizeof(uint32_t), 0); + if (optimal) { + codes_buf = &ev->dev; + codes_offset = dev_base + ev->codes_offset; + vk->CmdFillBuffer(exec->buf, ev->dev.buf, + codes_offset + NB_TAB_ENTRIES*sizeof(uint32_t), + NB_TAB_ENTRIES*sizeof(uint32_t), 0); + } + + RET(ff_vk_exec_add_dep_frame(&ev->s, exec, frame, + VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT, + VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT)); + RET(ff_vk_create_imageviews(&ev->s, exec, views, frame, FF_VK_REP_UINT)); + ff_vk_frame_barrier(&ev->s, exec, frame, img_bar, &nb_img_bar, + VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT, + VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + VK_ACCESS_SHADER_READ_BIT, + VK_IMAGE_LAYOUT_GENERAL, + VK_QUEUE_FAMILY_IGNORED); + vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .pImageMemoryBarriers = img_bar, + .imageMemoryBarrierCount = nb_img_bar, + }); + + /* Transform and quantise, component by component */ + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_dct, 0, 0, 0, + &ev->dev, dev_base, ev->codes_offset, + VK_FORMAT_UNDEFINED); + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_dct, 0, 1, 0, + &ev->quant_buf, 0, VK_WHOLE_SIZE, + VK_FORMAT_UNDEFINED); + ff_vk_shader_update_img_array(&ev->s, exec, &ev->shd_dct, frame, views, + 0, 2, VK_IMAGE_LAYOUT_GENERAL, + VK_NULL_HANDLE); + ff_vk_exec_bind_shader(&ev->s, exec, &ev->shd_dct); + ff_vk_shader_update_push_const(&ev->s, exec, &ev->shd_dct, + VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(ev->dct_push), &ev->dct_push); + for (uint32_t c = 0; c < ev->nb_components; c++) { + ff_vk_shader_update_push_const(&ev->s, exec, &ev->shd_dct, + VK_SHADER_STAGE_COMPUTE_BIT, + offsetof(DCTPushData, comp), sizeof(c), &c); + vk->CmdDispatch(exec->buf, + (ev->mcu_width*ev->hsample[c] + DCT_WG_BLOCKS - 1) / DCT_WG_BLOCKS, + ev->mcu_height*ev->vsample[c], 1); + } + + compute_barrier(vk, exec, + VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT | + VK_PIPELINE_STAGE_2_CLEAR_BIT, + VK_ACCESS_2_SHADER_WRITE_BIT | + VK_ACCESS_2_TRANSFER_WRITE_BIT); + + if (optimal) { + /* Count the symbols of the frame */ + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_histogram, 0, 0, 0, + &ev->dev, dev_base, ev->codes_offset, + VK_FORMAT_UNDEFINED); + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_histogram, 0, 1, 0, + &ev->dev, codes_offset + NB_TAB_ENTRIES*sizeof(uint32_t), + NB_TAB_ENTRIES*sizeof(uint32_t), + VK_FORMAT_UNDEFINED); + ff_vk_exec_bind_shader(&ev->s, exec, &ev->shd_histogram); + ff_vk_shader_update_push_const(&ev->s, exec, &ev->shd_histogram, + VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(ev->block_push), &ev->block_push); + vk->CmdDispatch(exec->buf, (ev->nb_blocks + 255) / 256, 1, 1); + + compute_barrier(vk, exec, VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + VK_ACCESS_2_SHADER_WRITE_BIT); + + /* Build the tables from it */ + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_tables, 0, 0, 0, + &ev->dev, codes_offset, + 2*NB_TAB_ENTRIES*sizeof(uint32_t), + VK_FORMAT_UNDEFINED); + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_tables, 0, 1, 0, + &ev->host, host_base + ev->dht_offset, + sizeof(DHTData), VK_FORMAT_UNDEFINED); + ff_vk_exec_bind_shader(&ev->s, exec, &ev->shd_tables); + vk->CmdDispatch(exec->buf, NB_TABLES, 1, 1); + + compute_barrier(vk, exec, VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + VK_ACCESS_2_SHADER_WRITE_BIT); + } + + /* Code each restart interval into its slot */ + ev->block_push.bytestream = fd->slots.address; + ev->block_push.slot_size = fd->slot_size; + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_huffman, 0, 0, 0, + &ev->dev, dev_base, ev->codes_offset, + VK_FORMAT_UNDEFINED); + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_huffman, 0, 1, 0, + codes_buf, codes_offset, + NB_TAB_ENTRIES*sizeof(uint32_t), + VK_FORMAT_UNDEFINED); + ff_vk_shader_update_desc_buffer(&ev->s, exec, &ev->shd_huffman, 0, 2, 0, + &ev->host, host_base, + (ev->nb_segs + 2)*sizeof(uint32_t), + VK_FORMAT_UNDEFINED); + ff_vk_exec_bind_shader(&ev->s, exec, &ev->shd_huffman); + ff_vk_shader_update_push_const(&ev->s, exec, &ev->shd_huffman, + VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(ev->block_push), &ev->block_push); + vk->CmdDispatch(exec->buf, ev->nb_segs, 1, 1); + + /* Pack the slots */ + RET(ff_vk_seg_gather(&ev->s, exec, &ev->shd_gather, + &ev->host, host_base, ev->nb_segs, + &fd->slots, fd->slot_size, + &fd->out, 0, 0)); + + RET(ff_vk_exec_submit(&ev->s, exec)); + + return 0; + +fail: + ff_vk_exec_discard(&ev->s, exec); + return err; +} + +static int submit_frame(AVCodecContext *avctx, FFVkExecContext *exec, + AVFrame *frame) +{ + VulkanEncodeMJPEGFrameData *fd = exec->opaque; + int err; + + fd->frame = av_frame_clone(frame); + if (!fd->frame) + return AVERROR(ENOMEM); + + err = record_frame(avctx, exec); + if (err < 0) + av_frame_free(&fd->frame); + return err; +} + +/* Invalidates the first size bytes of a host-visible buffer, if needed */ +static void invalidate_buffer(VulkanEncodeMJPEGContext *ev, FFVkBuffer *buf, + size_t size) +{ + FFVulkanFunctions *vk = &ev->s.vkfn; + size_t atom = ev->s.props.properties.limits.nonCoherentAtomSize; + + if (buf->flags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT) + return; + + size = FFALIGN(size, atom); + vk->InvalidateMappedMemoryRanges(ev->s.hwctx->act_dev, 1, + &(VkMappedMemoryRange) { + .sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE, + .memory = buf->mem, + .size = size < buf->size ? size : VK_WHOLE_SIZE, + }); +} + +static int get_packet(AVCodecContext *avctx, FFVkExecContext *exec, + AVPacket *pkt) +{ + int err; + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + VulkanEncodeMJPEGFrameData *fd = exec->opaque; + const uint8_t *host = ev->host.mapped_mem + exec->idx*ev->host_stride; + const uint32_t *sizes = (const uint32_t *)host; + + while (1) { + ff_vk_exec_wait(&ev->s, exec); + invalidate_buffer(ev, &ev->host, ev->host.size); + uint32_t needed = sizes[ev->nb_segs + 1]; + if (!needed) + break; + + /* An interval overflowed its slot: grow them, with some room for + * the frames to come, unless a frame before did, and code the frame + * again */ + if (needed > ev->slot_size) { + av_log(avctx, AV_LOG_VERBOSE, "Restart interval of %u bytes " + "overflowed its slot, growing slots from %zu bytes\n", + needed, ev->slot_size); + ev->slot_size = FFMIN(FFALIGN(needed + needed / 4, 16), + ev->max_slot_size); + } + RET(record_frame(avctx, exec)); + } + + size_t payload_size = sizes[ev->nb_segs]; + invalidate_buffer(ev, &fd->out, payload_size); + + const DHTData *dht = &ev->fixed_dht; + if (ev->huffman == HUFFMAN_TABLE_OPTIMAL) + dht = (const DHTData *)(host + ev->dht_offset); + + size_t max_header_size = MAX_BASE_HEADER_SIZE; + RET(ff_mjpeg_add_icc_profile_size(avctx, fd->frame, &max_header_size)); + RET(ff_mjpeg_add_gain_map_size(avctx, NULL, fd->frame, &max_header_size)); + RET(ff_get_encode_buffer(avctx, pkt, max_header_size + payload_size, 0)); + + PutBitContext pb; + init_put_bits(&pb, pkt->data, max_header_size); + write_headers(avctx, &pb, fd->frame, dht); + flush_put_bits(&pb); + size_t header_size = put_bytes_output(&pb); + + memcpy(pkt->data + header_size, fd->out.mapped_mem, payload_size); + av_shrink_packet(pkt, header_size + payload_size); + + av_log(avctx, AV_LOG_DEBUG, "Encoded frame: %d bytes\n", pkt->size); + +fail: + av_frame_free(&fd->frame); + return err; +} + +static int vulkan_encode_mjpeg_receive_packet(AVCodecContext *avctx, + AVPacket *pkt) +{ + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + return ff_vk_encode_loop_receive_packet(avctx, &ev->loop, pkt); +} + +static av_cold void vulkan_encode_mjpeg_flush(AVCodecContext *avctx) +{ + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + ff_vk_encode_loop_flush(avctx, &ev->loop); +} + +static av_cold int load_shader(VulkanEncodeMJPEGContext *ev, FFVulkanShader *shd, + const uint8_t *spv, size_t spv_len, + uint32_t wg_size[3], VkSpecializationInfo *spec, + size_t push_size, + const FFVulkanDescriptorSetBinding *desc, int nb_desc) +{ + ff_vk_shader_load(shd, VK_SHADER_STAGE_COMPUTE_BIT, spec, wg_size, 0); + if (push_size) + ff_vk_shader_add_push_const(shd, 0, push_size, VK_SHADER_STAGE_COMPUTE_BIT); + ff_vk_shader_add_descriptor_set(&ev->s, shd, desc, nb_desc, 0); + + int err = ff_vk_shader_link(&ev->s, shd, spv, spv_len, "main"); + if (err < 0) + return err; + return ff_vk_shader_register_exec(&ev->s, &ev->exec_pool, shd); +} + +#define BUF(n) { .name = n, .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, \ + .stages = VK_SHADER_STAGE_COMPUTE_BIT } + +static av_cold int init_shaders(AVCodecContext *avctx) +{ + int err; + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + FFVulkanContext *s = &ev->s; + uint32_t min_sg = s->subgroup_props.minSubgroupSize; + + /* Transform: one invocation per row of a block */ + SPEC_LIST_CREATE(dct_sl, 2, 2*sizeof(uint32_t)) + SPEC_LIST_ADD(dct_sl, 16, 32, DCT_WG_BLOCKS); + SPEC_LIST_ADD(dct_sl, 17, 32, 1); + RET(load_shader(ev, &ev->shd_dct, + ff_mjpeg_encode_dct_comp_spv_data, + ff_mjpeg_encode_dct_comp_spv_len, + (uint32_t []) { 8, DCT_WG_BLOCKS, 1 }, dct_sl, sizeof(DCTPushData), + (FFVulkanDescriptorSetBinding []) { + BUF("coeffs_buf"), BUF("quant_buf"), + { .name = "src", .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, + .stages = VK_SHADER_STAGE_COMPUTE_BIT, + .elems = av_pix_fmt_count_planes(s->frames->sw_format) }, + }, 3)); + + /* Huffman coding: one workgroup per restart interval, one invocation + * per block, as for the symbol counting below */ + SPEC_LIST_CREATE(blk_sl, 3, 3*sizeof(uint32_t)) + SPEC_LIST_ADD(blk_sl, 0, 32, FFALIGN(ev->seg_blocks, min_sg) / min_sg); + SPEC_LIST_ADD(blk_sl, 1, 32, ev->seg_blocks); + SPEC_LIST_ADD(blk_sl, 2, 32, ev->blocks_per_mcu); + RET(load_shader(ev, &ev->shd_huffman, + ff_mjpeg_encode_huffman_comp_spv_data, + ff_mjpeg_encode_huffman_comp_spv_len, + (uint32_t []) { ev->seg_blocks, 1, 1 }, blk_sl, sizeof(BlockPushData), + (FFVulkanDescriptorSetBinding []) { + BUF("coeffs_buf"), BUF("huff_buf"), BUF("sizes_buf"), + }, 3)); + + if (ev->huffman == HUFFMAN_TABLE_OPTIMAL) { + RET(load_shader(ev, &ev->shd_histogram, + ff_mjpeg_encode_histogram_comp_spv_data, + ff_mjpeg_encode_histogram_comp_spv_len, + (uint32_t []) { 256, 1, 1 }, blk_sl, sizeof(BlockPushData), + (FFVulkanDescriptorSetBinding []) { + BUF("coeffs_buf"), BUF("hist_buf"), + }, 2)); + + /* One workgroup per table, one invocation per symbol */ + SPEC_LIST_CREATE(tab_sl, 1, sizeof(uint32_t)) + SPEC_LIST_ADD(tab_sl, 0, 32, 256 / min_sg); + RET(load_shader(ev, &ev->shd_tables, + ff_mjpeg_encode_tables_comp_spv_data, + ff_mjpeg_encode_tables_comp_spv_len, + (uint32_t []) { 256, 1, 1 }, tab_sl, 0, + (FFVulkanDescriptorSetBinding []) { + BUF("tables_buf"), BUF("dht_buf"), + }, 2)); + } + + RET(ff_vk_seg_gather_init(s, &ev->exec_pool, &ev->shd_gather)); + +fail: + return err; +} + +/* Quantisers and the standard Huffman tables, which do not change */ +static av_cold int init_tables(AVCodecContext *avctx) +{ + int err; + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + const uint16_t *matrix[2] = { ff_mpeg1_default_intra_matrix, + ff_mpeg1_default_intra_matrix }; + /* Code counts per length, from index 1, and symbols */ + const uint8_t *const bits[NB_TABLES] = { ff_mjpeg_bits_dc_luminance, + ff_mjpeg_bits_dc_chrominance, + ff_mjpeg_bits_ac_luminance, + ff_mjpeg_bits_ac_chrominance }; + const uint8_t *const vals[NB_TABLES] = { ff_mjpeg_val_dc, ff_mjpeg_val_dc, + ff_mjpeg_val_ac_luminance, + ff_mjpeg_val_ac_chrominance }; + float *quant; + uint32_t *codes; + + /* As in the software encoder, the scale is part of the matrix, and the + * DC quantiser stays fixed. */ + RET(ff_check_codec_matrices(avctx, FF_MATRIX_TYPE_INTRA | FF_MATRIX_TYPE_CHROMA_INTRA, + (7 + ev->qscale) / ev->qscale, 65535)); + if (avctx->intra_matrix) + matrix[0] = matrix[1] = avctx->intra_matrix; + if (avctx->chroma_intra_matrix) + matrix[1] = avctx->chroma_intra_matrix; + + for (int t = 0; t < 2; t++) { + ev->quant_matrix[t][0] = 8; + for (int i = 1; i < 64; i++) + ev->quant_matrix[t][i] = av_clip_uint8((matrix[t][i]*ev->qscale) >> 3); + } + ev->nb_quant_matrices = 1 + (ev->nb_components > 1 && + memcmp(ev->quant_matrix[0], ev->quant_matrix[1], 64)); + + RET(ff_vk_create_buf(&ev->s, &ev->quant_buf, 2*64*sizeof(float), NULL, NULL, + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT)); + RET(ff_vk_map_buffer(&ev->s, &ev->quant_buf, (uint8_t **)&quant, 0)); + for (int t = 0; t < 2; t++) + for (int i = 0; i < 64; i++) + quant[t*64 + i] = 1.0f / ev->quant_matrix[t][i]; + RET(ff_vk_unmap_buffer(&ev->s, &ev->quant_buf, 1)); + + RET(ff_vk_create_buf(&ev->s, &ev->codes_buf, NB_TAB_ENTRIES*sizeof(uint32_t), + NULL, NULL, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT)); + RET(ff_vk_map_buffer(&ev->s, &ev->codes_buf, (uint8_t **)&codes, 0)); + for (int t = 0; t < NB_TABLES; t++) { + uint8_t huff_size[256]; + uint16_t huff_code[256]; + int nb_vals = 0; + + for (int l = 1; l <= 16; l++) + nb_vals += bits[t][l]; + + ff_mjpeg_build_huffman_codes(huff_size, huff_code, bits[t], vals[t]); + for (int i = 0; i < nb_vals; i++) { + int sym = vals[t][i]; + codes[t*256 + sym] = huff_size[sym] << 16 | huff_code[sym]; + } + + memcpy(ev->fixed_dht.bits[t], bits[t], 17); + memcpy(ev->fixed_dht.vals[t], vals[t], nb_vals); + } + RET(ff_vk_unmap_buffer(&ev->s, &ev->codes_buf, 1)); + +fail: + return err; +} + +static av_cold int init_layout(AVCodecContext *avctx) +{ + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + const VkPhysicalDeviceLimits *limits = &ev->s.props.properties.limits; + enum AVPixelFormat sw_format = ev->s.frames->sw_format; + const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(sw_format); + DCTPushData *dp = &ev->dct_push; + BlockPushData *bp = &ev->block_push; + + /* Luma takes as many blocks per MCU as chroma is subsampled by */ + ev->nb_components = desc->nb_components; + ev->blocks_per_mcu = (1 << (desc->log2_chroma_w + desc->log2_chroma_h)) + + ev->nb_components - 1; + if (desc->comp[0].depth != 8 || desc->flags & AV_PIX_FMT_FLAG_RGB || + (ev->nb_components != 1 && ev->nb_components != 3) || + ev->blocks_per_mcu > MAX_MCU_BLOCKS) { + av_log(avctx, AV_LOG_ERROR, "Unsupported input format: %s\n", + av_get_pix_fmt_name(sw_format)); + return AVERROR(ENOTSUP); + } + + if (avctx->width > 65535 || avctx->height > 65535) { + av_log(avctx, AV_LOG_ERROR, "Dimensions %dx%d too large for JPEG\n", + avctx->width, avctx->height); + return AVERROR(EINVAL); + } + + ev->mcu_width = AV_CEIL_RSHIFT(avctx->width, 3 + desc->log2_chroma_w); + ev->mcu_height = AV_CEIL_RSHIFT(avctx->height, 3 + desc->log2_chroma_h); + ev->nb_blocks = ev->mcu_width*ev->mcu_height*ev->blocks_per_mcu; + + for (int c = 0, b = 0; c < ev->nb_components; c++) { + int log2_h = c ? 0 : desc->log2_chroma_w; + int log2_v = c ? 0 : desc->log2_chroma_h; + int nb = 1 << (log2_h + log2_v); + + ev->hsample[c] = 1 << log2_h; + ev->vsample[c] = 1 << log2_v; + dp->plane_size[c][0] = AV_CEIL_RSHIFT(avctx->width, c ? desc->log2_chroma_w : 0); + dp->plane_size[c][1] = AV_CEIL_RSHIFT(avctx->height, c ? desc->log2_chroma_h : 0); + dp->comp_info[c] = PACK8(desc->comp[c].plane, desc->comp[c].offset, + log2_h, log2_v); + dp->comp_offset[c] = b; + + /* The first block of a component in an MCU is predicted from the + * last one of the component in the previous MCU */ + for (int i = 0; i < nb; i++, b++) + bp->mcu_blocks[b] = PACK8(c > 0, i ? 1 : ev->blocks_per_mcu - nb + 1, 0, 0); + } + + /* Every block of a restart interval gets an invocation */ + int max_mcus = FFMIN(limits->maxComputeWorkGroupInvocations, + limits->maxComputeWorkGroupSize[0]) / ev->blocks_per_mcu; + ev->seg_mcus = ev->restart_interval ? ev->restart_interval + : AUTO_SEG_BLOCKS / ev->blocks_per_mcu; + ev->seg_mcus = FFMIN(ev->seg_mcus, ev->mcu_width*ev->mcu_height); + if (ev->seg_mcus > max_mcus) { + av_log(avctx, AV_LOG_ERROR, "Restart interval of %d MCUs too long, " + "the maximum is %d\n", ev->seg_mcus, max_mcus); + return AVERROR(EINVAL); + } + ev->seg_blocks = ev->seg_mcus*ev->blocks_per_mcu; + ev->nb_segs = (ev->mcu_width*ev->mcu_height + ev->seg_mcus - 1) / ev->seg_mcus; + + /* An interval ends with a 2-byte marker */ + ev->max_slot_size = FFALIGN((size_t)ev->seg_blocks*MAX_BLOCK_BYTES + 2, 16); + ev->slot_size = FFALIGN((size_t)ev->seg_blocks*INIT_BLOCK_BYTES + 2, 16); + + /* 64 planes of coefficients and 4 of mask halfwords per block, see + * mjpeg_encode_block.glsl */ + size_t align = limits->minStorageBufferOffsetAlignment; + ev->codes_offset = FFALIGN((size_t)ev->nb_blocks*(64 + 4)*sizeof(int16_t), align); + ev->dev_stride = ev->codes_offset; + if (ev->huffman == HUFFMAN_TABLE_OPTIMAL) + ev->dev_stride += 2*NB_TAB_ENTRIES*sizeof(uint32_t); + ev->dht_offset = FFALIGN((ev->nb_segs + 2)*sizeof(uint32_t), align); + ev->host_stride = FFALIGN(ev->dht_offset + sizeof(DHTData), align); + + dp->mcu_width = ev->mcu_width; + dp->blocks_per_mcu = ev->blocks_per_mcu; + dp->nb_blocks = ev->nb_blocks; + bp->nb_blocks = ev->nb_blocks; + bp->nb_segs = ev->nb_segs; + + return 0; +} + +static av_cold int init_buffers(AVCodecContext *avctx) +{ + int err; + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + uint8_t *mem; + + RET(ff_vk_create_buf(&ev->s, &ev->dev, ev->async_depth*ev->dev_stride, NULL, NULL, + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | + VK_BUFFER_USAGE_TRANSFER_DST_BIT, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)); + RET(ff_vk_create_buf(&ev->s, &ev->host, ev->async_depth*ev->host_stride, NULL, NULL, + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | + VK_BUFFER_USAGE_TRANSFER_DST_BIT, + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | + ev->s.host_cached_flag)); + RET(ff_vk_map_buffer(&ev->s, &ev->host, &mem, 0)); + + for (int i = 0; i < ev->async_depth; i++) + RET(alloc_slots(ev, &ev->exec_ctx_info[i])); + +fail: + return err; +} + +static av_cold int vulkan_encode_mjpeg_close(AVCodecContext *avctx) +{ + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + + ff_vk_exec_pool_free(&ev->s, &ev->exec_pool); + + ff_vk_shader_free(&ev->s, &ev->shd_dct); + ff_vk_shader_free(&ev->s, &ev->shd_histogram); + ff_vk_shader_free(&ev->s, &ev->shd_tables); + ff_vk_shader_free(&ev->s, &ev->shd_huffman); + ff_vk_shader_free(&ev->s, &ev->shd_gather); + + ff_vk_free_buf(&ev->s, &ev->quant_buf); + ff_vk_free_buf(&ev->s, &ev->codes_buf); + ff_vk_free_buf(&ev->s, &ev->dev); + ff_vk_free_buf(&ev->s, &ev->host); + + if (ev->exec_ctx_info) { + for (int i = 0; i < ev->async_depth; i++) { + VulkanEncodeMJPEGFrameData *fd = &ev->exec_ctx_info[i]; + ff_vk_free_buf(&ev->s, &fd->slots); + ff_vk_free_buf(&ev->s, &fd->out); + av_frame_free(&fd->frame); + } + av_freep(&ev->exec_ctx_info); + } + + ff_vk_encode_loop_uninit(&ev->loop); + ff_vk_uninit(&ev->s); + + return 0; +} + +static av_cold int vulkan_encode_mjpeg_init(AVCodecContext *avctx) +{ + int err; + VulkanEncodeMJPEGContext *ev = avctx->priv_data; + + if (!avctx->hw_frames_ctx) { + av_log(avctx, AV_LOG_ERROR, "An AVHWFramesContext is required.\n"); + return AVERROR(EINVAL); + } + + err = ff_mjpeg_encode_check_pix_fmt(avctx); + if (err < 0) + return err; + + /* -q:v selects the quantiser scale, within qmin and qmax, as in the + * software encoder */ + int qmin = FFMAX(avctx->qmin, 1); + ev->qscale = 3; + if (avctx->global_quality > 0) + ev->qscale = (avctx->global_quality + FF_QP2LAMBDA / 2) / FF_QP2LAMBDA; + ev->qscale = av_clip(ev->qscale, qmin, FFMAX(avctx->qmax, qmin)); + + err = ff_vk_init(&ev->s, avctx, NULL, avctx->hw_frames_ctx); + if (err < 0) + return err; + + err = init_layout(avctx); + if (err < 0) + return err; + + ev->qf = ff_vk_qf_find(&ev->s, VK_QUEUE_COMPUTE_BIT, 0); + if (!ev->qf) { + av_log(avctx, AV_LOG_ERROR, "Device has no compute queues!\n"); + return AVERROR(ENOTSUP); + } + + err = ff_vk_exec_pool_init(&ev->s, ev->qf, &ev->exec_pool, + ev->async_depth, 0, 0, 0, NULL); + if (err < 0) + return err; + + ev->async_depth = ev->exec_pool.pool_size; + ev->exec_ctx_info = av_calloc(ev->async_depth, sizeof(*ev->exec_ctx_info)); + if (!ev->exec_ctx_info) + return AVERROR(ENOMEM); + for (int i = 0; i < ev->async_depth; i++) + ev->exec_pool.contexts[i].opaque = &ev->exec_ctx_info[i]; + + err = init_tables(avctx); + if (err < 0) + return err; + + err = init_shaders(avctx); + if (err < 0) + return err; + + err = init_buffers(avctx); + if (err < 0) + return err; + + err = ff_vk_encode_loop_init(&ev->s, &ev->exec_pool, &ev->loop, + submit_frame, get_packet); + if (err < 0) + return err; + + av_log(avctx, AV_LOG_VERBOSE, "%dx%d, %d components, %d blocks per MCU, " + "restart interval of %d MCUs, %d intervals, qscale %d\n", + avctx->width, avctx->height, ev->nb_components, ev->blocks_per_mcu, + ev->seg_mcus, ev->nb_segs, ev->qscale); + + return 0; +} + +#define OFFSET(x) offsetof(VulkanEncodeMJPEGContext, x) +#define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM +static const AVOption vulkan_encode_mjpeg_options[] = { + { "huffman", "Huffman table strategy", OFFSET(huffman), AV_OPT_TYPE_INT, + { .i64 = HUFFMAN_TABLE_OPTIMAL }, 0, HUFFMAN_TABLE_OPTIMAL, VE, .unit = "huffman" }, + { "default", "Standard tables (Annex K)", 0, AV_OPT_TYPE_CONST, { .i64 = HUFFMAN_TABLE_DEFAULT }, 0, 0, VE, .unit = "huffman" }, + { "optimal", "Tables built for each frame", 0, AV_OPT_TYPE_CONST, { .i64 = HUFFMAN_TABLE_OPTIMAL }, 0, 0, VE, .unit = "huffman" }, + { "restart_interval", "Restart interval in MCUs (0 = auto)", OFFSET(restart_interval), + AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 65535, VE }, + { "async_depth", "Internal parallelization depth", OFFSET(async_depth), + AV_OPT_TYPE_INT, { .i64 = 2 }, 1, INT_MAX, VE }, + { NULL } +}; + +static const AVClass vulkan_encode_mjpeg_class = { + .class_name = "mjpeg_vulkan", + .item_name = av_default_item_name, + .option = vulkan_encode_mjpeg_options, + .version = LIBAVUTIL_VERSION_INT, +}; + +static const AVCodecHWConfigInternal *const vulkan_encode_mjpeg_hw_configs[] = { + HW_CONFIG_ENCODER_FRAMES(VULKAN, VULKAN), + NULL, +}; + +const FFCodec ff_mjpeg_vulkan_encoder = { + .p.name = "mjpeg_vulkan", + CODEC_LONG_NAME("MJPEG (Motion JPEG) (Vulkan)"), + .p.type = AVMEDIA_TYPE_VIDEO, + .p.id = AV_CODEC_ID_MJPEG, + .priv_data_size = sizeof(VulkanEncodeMJPEGContext), + .init = &vulkan_encode_mjpeg_init, + FF_CODEC_RECEIVE_PACKET_CB(&vulkan_encode_mjpeg_receive_packet), + .flush = &vulkan_encode_mjpeg_flush, + .close = &vulkan_encode_mjpeg_close, + .p.priv_class = &vulkan_encode_mjpeg_class, + .p.capabilities = AV_CODEC_CAP_DELAY | + AV_CODEC_CAP_HARDWARE | + AV_CODEC_CAP_DR1 | + AV_CODEC_CAP_ENCODER_FLUSH | + AV_CODEC_CAP_ENCODER_REORDERED_OPAQUE, + .caps_internal = FF_CODEC_CAP_INIT_CLEANUP | + FF_CODEC_CAP_ICC_PROFILES, + CODEC_PIXFMTS(AV_PIX_FMT_VULKAN), + .hw_configs = vulkan_encode_mjpeg_hw_configs, + /* Frames cannot be converted to full range once uploaded: limited range + * is refused at init instead, unless strictness allows it */ + .color_ranges = AVCOL_RANGE_MPEG | AVCOL_RANGE_JPEG, + .p.wrapper_name = "vulkan", +}; diff --git a/libavcodec/vulkan/Makefile b/libavcodec/vulkan/Makefile index 6b97605e08..8f78cc7633 100644 --- a/libavcodec/vulkan/Makefile +++ b/libavcodec/vulkan/Makefile @@ -36,6 +36,11 @@ OBJS-$(CONFIG_FFV1_VULKAN_HWACCEL) += vulkan/ffv1_dec_setup.comp.spv.o \ vulkan/ffv1_dec_bayer.comp.spv.o \ vulkan/ffv1_dec_bayer_golomb.comp.spv.o +OBJS-$(CONFIG_MJPEG_VULKAN_ENCODER) += vulkan/mjpeg_encode_dct.comp.spv.o \ + vulkan/mjpeg_encode_histogram.comp.spv.o \ + vulkan/mjpeg_encode_tables.comp.spv.o \ + vulkan/mjpeg_encode_huffman.comp.spv.o + OBJS-$(CONFIG_PRORES_KS_VULKAN_ENCODER) += vulkan/prores_ks_alpha_data.comp.spv.o \ vulkan/prores_ks_slice_data.comp.spv.o \ vulkan/prores_ks_estimate_slice.comp.spv.o \ diff --git a/libavcodec/vulkan/mjpeg_encode_block.glsl b/libavcodec/vulkan/mjpeg_encode_block.glsl new file mode 100644 index 0000000000..d3a95a16bd --- /dev/null +++ b/libavcodec/vulkan/mjpeg_encode_block.glsl @@ -0,0 +1,162 @@ +/* + * Copyright (c) 2026 Lynne <[email protected]> + * + * 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 + */ + +#ifndef VULKAN_MJPEG_ENCODE_BLOCK_H +#define VULKAN_MJPEG_ENCODE_BLOCK_H + +/* + * The symbol walk over one quantised block, shared by the histogram and the + * Huffman passes, which define the Coder struct and put_sym(). + * + * Blocks are numbered in coding order: MCU by MCU, and within an MCU, the + * blocks of each component in turn. Coefficients are stored packed and + * planar: for block blk of nb_blocks, coeffs[blk] is its DC, + * coeffs[k*nb_blocks + blk] its k-th nonzero AC coefficient in zig-zag + * order, and the 4 halfwords from plane 64 on the mask of its nonzero AC + * coefficients, bit n for zig-zag position n. So neighbouring invocations, + * which code neighbouring blocks, load neighbouring values. + */ + +layout (constant_id = 1) const uint seg_blocks = 512; /* per restart interval */ +layout (constant_id = 2) const uint blocks_per_mcu = 6; + +layout (set = 0, binding = 0, scalar) readonly buffer coeffs_buf { + int16_t coeffs[]; +}; + +layout (push_constant, scalar) uniform pushConstants { + u8buf bytestream; /* one slot per interval */ + uint slot_size; + uint nb_blocks; + uint nb_segs; + /* Per block of an MCU, a byte each: Huffman table class, and the + * distance to the previous block of the same component, which predicts + * its DC. 32-bit words, as in mjpeg_encode_dct.comp.glsl. */ + uint mcu_blocks[10]; +}; + +/* The 4 tables (DC luma, DC chroma, AC luma, AC chroma), 256 entries each */ +#define TAB_DC 0 +#define TAB_AC 2 +#define NB_TAB_ENTRIES (4*256) + +void put_sym(inout Coder c, const uint mode, uint tab, uint sym, + uint nb_extra, uint extra); + +/* Category of a value, and the extra bits that code it within it */ +uint category(int v) +{ + return uint(findMSB(abs(v)) + 1); +} + +uint extra_bits(int v, uint cat) +{ + return uint(v - int(v < 0)) & ((1u << cat) - 1u); +} + +/* + * A block is coded several times over, so all it takes is loaded once, up + * front, with independent loads, including its first nonzero AC + * coefficients, two per word. + */ +struct Block { + uint idx; + uint cls; /* Huffman table class */ + int diff; /* DC difference */ + u32vec2 mask; + u32vec4 ac_lo; + u32vec4 ac_hi; +}; + +/* Planes k and k + 1 of a block, as one word */ +uint load_pair(uint k, uint blk) +{ + return uint(uint16_t(coeffs[k*nb_blocks + blk])) | + uint(uint16_t(coeffs[(k + 1)*nb_blocks + blk])) << 16; +} + +/* Block blk is the b-th block of its restart interval; the DC prediction + * resets at the start of each interval. */ +Block load_block(uint blk, uint b) +{ + Block bl; + u8vec4 mb = unpack8(mcu_blocks[b % blocks_per_mcu]); + bl.idx = blk; + bl.cls = mb.x; + int dc = coeffs[blk]; + bl.diff = b >= mb.y ? dc - coeffs[blk - mb.y] : dc; + bl.mask = u32vec2(load_pair(64, blk), load_pair(66, blk)); + + /* Every block has room for 63 values, so these can be loaded whether + * or not they are used */ + uint w[8]; + [[unroll]] + for (uint i = 0; i < 8; i++) + w[i] = load_pair(1 + 2*i, blk); + bl.ac_lo = u32vec4(w[0], w[1], w[2], w[3]); + bl.ac_hi = u32vec4(w[4], w[5], w[6], w[7]); + return bl; +} + +/* The k-th nonzero AC coefficient */ +int ac_coeff(in Block bl, uint k) +{ + if (k >= 16) + return coeffs[(1 + k)*nb_blocks + bl.idx]; + uint w = k < 8 ? bl.ac_lo[(k >> 1) & 3] : bl.ac_hi[(k >> 1) & 3]; + return int(int16_t(uint16_t(w >> ((k & 1) << 4)))); +} + +void code_block(inout Coder c, const uint mode, in Block bl) +{ + const uint cls = bl.cls; + uint cat = category(bl.diff); + put_sym(c, mode, TAB_DC + cls, cat, cat, extra_bits(bl.diff, cat)); + + u32vec2 m = bl.mask; + uint k = 0; + uint last = 0; + while ((m.x | m.y) != 0) { + uint z; + if (m.x != 0) { + z = findLSB(m.x); + m.x &= m.x - 1; + } else { + z = 32 + findLSB(m.y); + m.y &= m.y - 1; + } + + uint run = z - last - 1; + last = z; + while (run > 15) { + put_sym(c, mode, TAB_AC + cls, 0xF0, 0, 0); /* ZRL */ + run -= 16; + } + + int v = ac_coeff(bl, k++); + cat = category(v); + put_sym(c, mode, TAB_AC + cls, (run << 4) | cat, cat, extra_bits(v, cat)); + } + + if (last != 63) + put_sym(c, mode, TAB_AC + cls, 0x00, 0, 0); /* EOB */ +} + +#endif /* VULKAN_MJPEG_ENCODE_BLOCK_H */ diff --git a/libavcodec/vulkan/mjpeg_encode_dct.comp.glsl b/libavcodec/vulkan/mjpeg_encode_dct.comp.glsl new file mode 100644 index 0000000000..028ab0e94c --- /dev/null +++ b/libavcodec/vulkan/mjpeg_encode_dct.comp.glsl @@ -0,0 +1,157 @@ +/* + * Copyright (c) 2026 Lynne <[email protected]> + * + * 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_clustered : require + +#include "common.glsl" +#include "dct.glsl" + +/* + * Forward DCT and quantisation of one component, dispatched over its blocks. + * Each workgroup transforms gl_WorkGroupSize.y neighbouring blocks of a row, + * with one invocation per row of a block. Blocks are stored in coding order, + * as described in mjpeg_encode_block.glsl. + */ + +/* Inverse zig-zag: zz_inv[raster index] = scan position */ +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), +}; + +layout (set = 0, binding = 0, scalar) writeonly buffer coeffs_buf { + int16_t coeffs[]; +}; + +/* Reciprocals of the quantisers, raster order */ +layout (set = 0, binding = 1, scalar) readonly buffer quant_buf { + float quant[2][64]; +}; + +layout (set = 0, binding = 2) uniform readonly uimage2D src[]; + +/* The arrays are indexed dynamically, and are made of 32-bit words: RADV + * miscompiles dynamically indexed arrays of 8-bit vectors in push + * constants. */ +layout (push_constant, scalar) uniform pushConstants { + ivec2 plane_size[3]; /* per component, in samples */ + /* Per component, a byte each: the plane and channel it is read from, + * and the log2 of its horizontal and vertical sampling factors */ + uint comp_info[3]; + uint comp_offset[3]; /* of the first block of each component in an MCU */ + uint mcu_width; + uint blocks_per_mcu; + uint total_blocks; + uint comp; /* the component dispatched */ +}; + +void main(void) +{ + const uint row = gl_LocalInvocationID.x; + const uint lb = gl_LocalInvocationID.y; + const uint bx = gl_WorkGroupID.x*gl_WorkGroupSize.y + lb; + const uint by = gl_WorkGroupID.y; + const u8vec4 ci = unpack8(comp_info[comp]); + const uint plane = ci.x, chan = ci.y, log2_h = ci.z, log2_v = ci.w; + + /* Blocks past the right edge take part in the barriers, and store + * nothing */ + const bool in_plane = bx < (mcu_width << log2_h); + const uint mcu = (by >> log2_v)*mcu_width + (bx >> log2_h); + const uint blk = mcu*blocks_per_mcu + comp_offset[comp] + + (bitfieldExtract(by, 0, int(log2_v)) << log2_h) + + bitfieldExtract(bx, 0, int(log2_h)); + + const ivec2 origin = ivec2(bx, by) << 3; + const ivec2 lim = plane_size[comp] - 1; + + /* Samples past the edges of the plane replicate the last row/column */ + const int y = min(origin.y + int(row), lim.y); + [[unroll]] + for (int i = 0; i < 8; i++) { + uint s = imageLoad(src[plane], ivec2(min(origin.x + i, lim.x), y))[chan]; + blocks[lb][row*9 + i] = float(s) - 128.0f; + } + barrier(); + + /* Columns, then rows. fdct8() is the orthonormal DCT, which is exactly + * the JPEG FDCT. The row pass leaves each invocation with its own row. */ + fdct8(lb, row, 9); + barrier(); + fdct8(lb, row*9, 1); + + /* Quantise, rounding the DC to nearest and the AC coefficients with the + * same deadzone as the software encoder. Values are clipped to what a + * baseline 8-bit stream can code: 11-bit DC differences and 10-bit AC + * coefficients. */ + int lvl[8]; + u32vec2 m = u32vec2(0); + [[unroll]] + for (uint i = 0; i < 8; i++) { + const uint r = row*8 + i; + float v = blocks[lb][row*9 + i]*quant[min(comp, 1)][r]; + int l = min(int(abs(v) + (r == 0 ? 0.5f : 0.375f)), 1023); + lvl[i] = v < 0.0f ? -l : l; + + uint z = zz_inv[r]; + if (z != 0 && l != 0) + m[z >> 5] |= 1u << (z & 31); + } + + /* The 8 invocations of a block are consecutive within the subgroup */ + m = subgroupClusteredOr(m, 8); + if (!in_plane) + return; + + /* The k-th nonzero coefficient in zig-zag order goes in plane 1 + k */ + [[unroll]] + for (uint i = 0; i < 8; i++) { + uint z = zz_inv[row*8 + i]; + if (z != 0 && lvl[i] != 0) { + u32vec2 below = z < 32 ? u32vec2((1u << z) - 1u, 0u) : + u32vec2(~0u, (1u << (z - 32)) - 1u); + uint rank = bitCount(m.x & below.x) + bitCount(m.y & below.y); + coeffs[(1 + rank)*total_blocks + blk] = int16_t(lvl[i]); + } + } + + /* The DC, and the mask as 4 halfwords from plane 64 */ + if (row == 0) + coeffs[blk] = int16_t(lvl[0]); + if (row < 4) + coeffs[(64 + row)*total_blocks + blk] = int16_t(m[row >> 1] >> (16*(row & 1))); +} diff --git a/libavcodec/vulkan/mjpeg_encode_histogram.comp.glsl b/libavcodec/vulkan/mjpeg_encode_histogram.comp.glsl new file mode 100644 index 0000000000..8d99733250 --- /dev/null +++ b/libavcodec/vulkan/mjpeg_encode_histogram.comp.glsl @@ -0,0 +1,66 @@ +/* + * Copyright (c) 2026 Lynne <[email protected]> + * + * 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" + +/* Counts the Huffman symbols of the frame, from which optimal tables are + * built. One invocation per block. */ + +struct Coder { + uint unused; +}; + +#include "mjpeg_encode_block.glsl" + +/* Zeroed before the dispatch */ +layout (set = 0, binding = 1, scalar) buffer hist_buf { + uint32_t hist[NB_TAB_ENTRIES]; +}; + +shared uint32_t local_hist[NB_TAB_ENTRIES]; + +void put_sym(inout Coder c, const uint mode, uint tab, uint sym, + uint nb_extra, uint extra) +{ + atomicAdd(local_hist[tab*256 + sym], 1u); +} + +void main(void) +{ + const uint blk = gl_GlobalInvocationID.x; + + for (uint i = gl_LocalInvocationIndex; i < NB_TAB_ENTRIES; i += gl_WorkGroupSize.x) + local_hist[i] = 0; + barrier(); + + if (blk < nb_blocks) { + Coder c; + code_block(c, 0, load_block(blk, blk % seg_blocks)); + } + barrier(); + + for (uint i = gl_LocalInvocationIndex; i < NB_TAB_ENTRIES; i += gl_WorkGroupSize.x) + if (local_hist[i] != 0) + atomicAdd(hist[i], local_hist[i]); +} diff --git a/libavcodec/vulkan/mjpeg_encode_huffman.comp.glsl b/libavcodec/vulkan/mjpeg_encode_huffman.comp.glsl new file mode 100644 index 0000000000..cf6370db99 --- /dev/null +++ b/libavcodec/vulkan/mjpeg_encode_huffman.comp.glsl @@ -0,0 +1,212 @@ +/* + * Copyright (c) 2026 Lynne <[email protected]> + * + * 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 + +#include "common.glsl" + +/* + * Huffman coding. One workgroup codes one restart interval into its own + * slot, with one invocation per block. + * + * Every block is coded three times over, by the same code: + * - to count its length in bits, from which a scan gives each block its + * start bit; + * - to count the 0xFF bytes it completes, from which a second scan gives + * it the position of its first byte after byte stuffing; + * - to write it out, stuffed. + * A block writes every byte that it completes, prepending the bits of the + * preceding blocks that the byte starts with, so blocks write disjoint + * bytes, and the interval is written without a separate stuffing pass. + */ + +layout (constant_id = 0) const uint max_subgroups = 32; /* at the smallest subgroup size */ + +#define MODE_BITS 0 +#define MODE_FF 1 +#define MODE_WRITE 2 + +struct Coder { + uint bits; /* MODE_BITS: bits coded, and */ + uint tail; /* the last 32 of them */ + uint64_t acc; /* bits not yet formed into a byte */ + uint nb_acc; + uint nb_ff; /* MODE_FF: 0xFF bytes formed */ + u8buf dst; /* MODE_WRITE */ + uint pos; +}; + +#include "mjpeg_encode_block.glsl" + +/* Per table entry: code length << 16 | code */ +layout (set = 0, binding = 1, scalar) readonly buffer huff_buf { + uint32_t huff_codes[NB_TAB_ENTRIES]; +}; + +/* Per interval, then the packed size, written by the gather pass, then the + * size of the largest interval that overflowed its slot, zeroed beforehand */ +layout (set = 0, binding = 2, scalar) buffer sizes_buf { + uint32_t seg_sizes[]; +}; + +shared uint32_t codes[NB_TAB_ENTRIES]; +shared uint blk_end[seg_blocks]; /* bits coded << 8 | the last 8 of them */ +shared uint sg_sum[2][max_subgroups]; + +void put_bits_stuffed(inout Coder c, const uint mode, uint n, uint v) +{ + if (mode == MODE_BITS) { + c.bits += n; + c.tail = (c.tail << n) | v; + return; + } + + c.acc = (c.acc << n) | v; + c.nb_acc += n; + while (c.nb_acc >= 8) { + c.nb_acc -= 8; + uint8_t byte = uint8_t(c.acc >> c.nb_acc); + if (mode == MODE_FF) { + c.nb_ff += uint(byte == uint8_t(0xFF)); + } else { + c.dst[c.pos++].v = byte; + if (byte == uint8_t(0xFF)) + c.dst[c.pos++].v = uint8_t(0x00); + } + } +} + +/* Codes and extra bits never take more than 16 + 11 bits */ +void put_sym(inout Coder c, const uint mode, uint tab, uint sym, + uint nb_extra, uint extra) +{ + uint e = codes[tab*256 + sym]; + put_bits_stuffed(c, mode, (e >> 16) + nb_extra, + ((e & 0xFFFF) << nb_extra) | extra); +} + +/* The last block of an interval pads its final byte with 1 bits */ +void pad_byte(inout Coder c, const uint mode) +{ + uint n = (8 - c.nb_acc) & 7; + put_bits_stuffed(c, mode, n, (1u << n) - 1u); +} + +/* Exclusive prefix sum over the workgroup */ +uint scan_add(uint v, const uint pass) +{ + uint excl = subgroupExclusiveAdd(v); + uint total = subgroupAdd(v); + if (subgroupElect()) + sg_sum[pass][gl_SubgroupID] = total; + barrier(); + + /* The totals of the subgroups before this one, a lane each */ + uint before = 0; + for (uint i = gl_SubgroupInvocationID; i < gl_SubgroupID; i += gl_SubgroupSize) + before += sg_sum[pass][i]; + return excl + subgroupAdd(before); +} + +void main(void) +{ + const uint seg = gl_WorkGroupID.x; + const uint b = gl_LocalInvocationID.x; + const uint first = seg*seg_blocks; + const uint nb = min(seg_blocks, nb_blocks - first); + const bool in_seg = b < nb; + const bool last = b == nb - 1; + + /* Both loads are in flight at once */ + Block bl; + if (in_seg) + bl = load_block(first + b, b); + + for (uint i = b; i < NB_TAB_ENTRIES; i += seg_blocks) + codes[i] = huff_codes[i]; + barrier(); + + Coder c; + c.bits = 0; + c.tail = 0; + if (in_seg) + code_block(c, MODE_BITS, bl); + + blk_end[b] = (c.bits << 8) | (c.tail & 0xFF); + const uint off = scan_add(c.bits, 0); + + /* The bits of the byte this block starts in that precede it come from + * the blocks before it: from at most 4 of them, as a block has at least + * 2 bits, a DC and an AC code */ + const uint lead = off & 7; + uint lead_bits = 0; + [[unroll]] + for (uint k = 1, got = 0; k <= 4; k++) { + uint e = b >= k ? blk_end[b - k] : 0; + uint n = min(lead - got, e >> 8); + lead_bits |= (e & ((1u << n) - 1u)) << got; + got += n; + } + + c.acc = lead_bits; + c.nb_acc = lead; + c.nb_ff = 0; + if (in_seg) { + code_block(c, MODE_FF, bl); + if (last) + pad_byte(c, MODE_FF); + } + + const uint ff_off = scan_add(c.nb_ff, 1); + if (!in_seg) + return; + + /* Slots are sized for typical intervals, not the worst case. A block + * that would write past its slot writes nothing, and the last block + * reports the size needed, for the frame to be coded again with larger + * slots. Ends only grow from block to block; every interval ends with + * a 2-byte marker. */ + const uint end = ((off + c.bits + (last ? 7 : 0)) >> 3) + ff_off + c.nb_ff + 2; + if (end > slot_size) { + if (last) { + seg_sizes[seg] = 0; + atomicMax(seg_sizes[nb_segs + 1], end); + } + return; + } + + c.acc = lead_bits; + c.nb_acc = lead; + c.dst = OFFBUF(u8buf, bytestream, uint64_t(seg)*slot_size); + c.pos = (off >> 3) + ff_off; + code_block(c, MODE_WRITE, bl); + + if (last) { + pad_byte(c, MODE_WRITE); + /* RSTn, or EOI after the last interval */ + c.dst[c.pos++].v = uint8_t(0xFF); + c.dst[c.pos++].v = uint8_t(seg == nb_segs - 1 ? 0xD9 : 0xD0 + (seg & 7)); + seg_sizes[seg] = c.pos; + } +} diff --git a/libavcodec/vulkan/mjpeg_encode_tables.comp.glsl b/libavcodec/vulkan/mjpeg_encode_tables.comp.glsl new file mode 100644 index 0000000000..579e99f118 --- /dev/null +++ b/libavcodec/vulkan/mjpeg_encode_tables.comp.glsl @@ -0,0 +1,227 @@ +/* + * Copyright (c) 2026 Lynne <[email protected]> + * + * 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" + +/* + * Builds an optimal Huffman table from the symbol histogram. One workgroup + * builds each table, DC luma, DC chroma, AC luma and AC chroma, with one + * invocation per symbol. + * + * As in libjpeg (and Annex K.2), every table gets one extra symbol with the + * lowest frequency, 256, which takes the longest code; dropping it once + * lengths are limited to 16 bits keeps the all-ones code unused. + * + * Only merging the two lightest nodes into a tree is sequential. + */ + +#define NB_SYMS 256 +#define MAX_LEN 64 /* code lengths of symbol counts below 2^32 stay under 48 */ + +layout (set = 0, binding = 0, scalar) buffer tables_buf { + uint32_t huff_codes[4*NB_SYMS]; /* per entry: code length << 16 | code */ + uint32_t hist[4*NB_SYMS]; +}; + +/* The tables as they are written in the DHT segment, for the host */ +layout (set = 0, binding = 1, scalar) writeonly buffer dht_buf { + uint8_t dht_bits[4][17]; + uint8_t dht_vals[4][NB_SYMS]; +}; + +shared uint nb_used; +shared uint64_t used_key[NB_SYMS]; /* of the used symbols, in any order */ +shared uint weight[NB_SYMS + 1]; /* of the used symbols, lightest first */ +shared uint parent[NB_SYMS + 1]; +/* Internal nodes, in order of creation: weight, then depth; parent, then + * the ancestor that pointer jumping has reached */ +shared uint node_val[NB_SYMS]; +shared uint node_up[NB_SYMS]; +shared uint bits[MAX_LEN]; /* codes per length */ +shared uint first_pos[17]; /* of each length, among the symbols in order of length */ +shared uint first_code[17]; + +/* + * Merges the two lightest nodes until one is left. Leaves and internal + * nodes each come in order of weight, so both are queues; the two first + * nodes of each are kept at hand, so that refilling them is not in the way + * of the next merges. + */ +void build_tree(uint n) +{ + const uint inf = 0xFFFFFFFF; + uint leaf = 0, node = 0; + uint l0 = weight[0], l1 = weight[1]; + uint n0 = inf, n1 = inf; + + for (uint i = 0; i < n - 1; i++) { + uint w = 0; + for (int k = 0; k < 2; k++) { + if (l0 <= n0) { + w += l0; + parent[leaf] = i; + l0 = l1; + l1 = leaf + 2 < n ? weight[leaf + 2] : inf; + leaf++; + } else { + w += n0; + node_up[node] = i; + n0 = n1; + n1 = node + 2 < i ? node_val[node + 2] : inf; + node++; + } + } + node_val[i] = w; + if (node == i) + n0 = w; + else if (node + 1 == i) + n1 = w; + } +} + +void main(void) +{ + const uint t = gl_WorkGroupID.x; + const uint s = gl_LocalInvocationIndex; + const uint f = hist[t*NB_SYMS + s]; + + if (s < MAX_LEN) + bits[s] = 0; + if (s == 0) + nb_used = 0; + barrier(); + + /* Symbols rank in ascending order of frequency, breaking ties by + * descending symbol */ + const uint64_t key = (uint64_t(f) << 8) | (255 - s); + if (f != 0) + used_key[atomicAdd(nb_used, 1)] = key; + barrier(); + + /* The extra symbol, with a frequency of 1, comes first */ + const uint n = nb_used + 1; + uint rank = 1; + if (f != 0) { + for (uint j = 0; j < n - 1; j++) + rank += uint(used_key[j] < key); + weight[rank] = f; + } + if (s == 0) + weight[0] = 1; + barrier(); + + /* A table with only the extra symbol is unused */ + if (n < 2) { + if (s < 17) + dht_bits[t][s] = uint8_t(0); + return; + } + + if (s == 0) + build_tree(n); + barrier(); + + /* Depths of the internal nodes, by pointer jumping from each to the + * root, the last node: depth += depth of ancestor, ancestor = its own + * ancestor, doubling the distance covered every round */ + const uint root = n - 2; + const bool is_node = s <= root; + if (is_node) { + node_val[s] = uint(s != root); + if (s == root) + node_up[s] = root; + } + barrier(); + + for (uint d = 1; d < n - 1; d <<= 1) { + uint val, up; + if (is_node) { + uint a = node_up[s]; + val = node_val[s] + node_val[a]; + up = node_up[a]; + } + barrier(); + if (is_node) { + node_val[s] = val; + node_up[s] = up; + } + barrier(); + } + + /* Code lengths, counted per length; the extra symbol's is the longest */ + const uint max_len = node_val[parent[0]] + 1; + if (f != 0) + atomicAdd(bits[node_val[parent[rank]] + 1], 1); + if (s == 0) + atomicAdd(bits[max_len], 1); + barrier(); + + if (s == 0) { + /* Limit lengths to 16 bits (Annex K.3); the symbols keep their + * order, and take the lengths in it */ + for (uint i = max_len; i > 16; i--) { + while (bits[i] > 0) { + uint j = i - 2; + while (bits[j] == 0) + j--; + bits[i] -= 2; + bits[i - 1]++; + bits[j + 1] += 2; + bits[j]--; + } + } + + /* Drop the extra symbol */ + uint i = 16; + while (bits[i] == 0) + i--; + bits[i]--; + + /* Canonical codes */ + uint code = 0; + uint p = 0; + for (uint l = 1; l <= 16; l++) { + first_pos[l] = p; + first_code[l] = code; + dht_bits[t][l] = uint8_t(bits[l]); + p += bits[l]; + code = (code + bits[l]) << 1; + } + } + barrier(); + + if (f == 0) + return; + + /* Merging in order of weight gives the heavier symbols the shorter + * codes, so in order of length the symbols come by descending rank, + * and a symbol's length is the one assigned at its place */ + const uint p = n - 1 - rank; + uint l = 16; + while (first_pos[l] > p) + l--; + + huff_codes[t*NB_SYMS + s] = (l << 16) | (first_code[l] + p - first_pos[l]); + dht_vals[t][p] = uint8_t(s); +} -- 2.54.0 _______________________________________________ ffmpeg-devel mailing list -- [email protected] To unsubscribe send an email to [email protected]
