From patchwork Wed Oct 23 23:03:27 2024 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: IndecisiveTurtle X-Patchwork-Id: 52458 Delivered-To: ffmpegpatchwork2@gmail.com Received: by 2002:a05:612c:1c41:b0:48e:c0f8:d0de with SMTP id im1csp30231vqb; Wed, 23 Oct 2024 19:09:04 -0700 (PDT) X-Forwarded-Encrypted: i=2; AJvYcCVwcSuJBeozSJm+hqguwQ3uZHpOUEBCW3+hqn5329iu7gR4sMa5r3EMWqBKMNsJAsLVB1LJSjtUdFC/MmkXBPoa@gmail.com X-Google-Smtp-Source: AGHT+IFBi3tve7U4de8/+hxVql71IVqPp6XmtfHbLlKqkYFQt/XdzQmKYmF0Dt8V8A1eUfrkyvEU X-Received: by 2002:a17:907:3187:b0:a99:f945:8776 with SMTP id a640c23a62f3a-a9abf88a58bmr363992866b.24.1729735743775; Wed, 23 Oct 2024 19:09:03 -0700 (PDT) ARC-Seal: i=1; a=rsa-sha256; t=1729735743; cv=none; d=google.com; s=arc-20240605; b=McK5YwVuZx11rTRZcgnN2PmqOcfNEv3nOEiW6NNlNl8jggUj3uhojlci/Aif1b7i64 0lFHm07Jm84DXNCnQnu3Th2KsoZWKZzK2ptM0lpsEmf8Cz09xD+4CuIi5qBg6v577lQk oBShZkYbF9vZx+N5n5vBa4yXMyC2KBpcPzJt5UEqBR7XMrS7dvXQS4mdySXfFgvAEu+3 dcbDja+ZJZwcLv7ZeD7cQZPs+Nmytjtu3ro4wO4msftAhrX7UZfGTXwfplSO77/aBriV 9gNBPOP+yoNRDNDWuLakXpzUfQQtdyyd25xlNWbWvQzmmzKqFHw76kCS6C7MPan6E4J1 ecTg== ARC-Message-Signature: i=1; a=rsa-sha256; c=relaxed/relaxed; d=google.com; s=arc-20240605; h=sender:errors-to:content-transfer-encoding:cc:reply-to :list-subscribe:list-help:list-post:list-archive:list-unsubscribe :list-id:precedence:subject:mime-version:message-id:date:to:from :dkim-signature:delivered-to; bh=J0Bkt5FUKdyC1PA8dklewf/V6GbP/VH+XUg9BJwv+98=; fh=ifgEp66h78QTxewnunWI44SZGWqBDDDWcajsEi/qBSk=; b=OLmC3ABdFfvEVpEurADHcqkCM37i6iuHavFIlvFkXzTlMDOfXCpuUggS/J2iSxxYYo pnY0bLeoRdM60STSwfrh/wfsv9G29Fe6hUrfKoCPVLOfRN2BYQDGAjsvM9u0gr3lodtp NppUbhP/bfSReB9tGnN8jC6xdVhBR7uAPhjVwSuFdqCdUBCOpbDcxlV0xKdez+foVzA/ AU1uoKqfFEOOQCFovBT1/CKHO9do4UvfJkEJ2r7YUSv/hmC9nluTvf9tHC+71N0Ys/ZD z95eL01dBUp9Wi8UjG+4Cv+vD1QpGVvpt6TkzfIQlU3ZBQd5HrOS0yNYKnqLC1iaTpjB ubQQ==; dara=google.com ARC-Authentication-Results: i=1; mx.google.com; dkim=neutral (body hash did not verify) header.i=@gmail.com header.s=20230601 header.b=Mo9z1LBk; spf=pass (google.com: domain of ffmpeg-devel-bounces@ffmpeg.org designates 79.124.17.100 as permitted sender) smtp.mailfrom=ffmpeg-devel-bounces@ffmpeg.org; dmarc=fail (p=NONE sp=QUARANTINE dis=NONE) header.from=gmail.com; dara=fail header.i=@gmail.com Return-Path: Received: from ffbox0-bg.mplayerhq.hu (ffbox0-bg.ffmpeg.org. [79.124.17.100]) by mx.google.com with ESMTP id a640c23a62f3a-a9a913705f9si657280866b.759.2024.10.23.19.09.02; Wed, 23 Oct 2024 19:09:03 -0700 (PDT) Received-SPF: pass (google.com: domain of ffmpeg-devel-bounces@ffmpeg.org designates 79.124.17.100 as permitted sender) client-ip=79.124.17.100; Authentication-Results: mx.google.com; dkim=neutral (body hash did not verify) header.i=@gmail.com header.s=20230601 header.b=Mo9z1LBk; spf=pass (google.com: domain of ffmpeg-devel-bounces@ffmpeg.org designates 79.124.17.100 as permitted sender) smtp.mailfrom=ffmpeg-devel-bounces@ffmpeg.org; dmarc=fail (p=NONE sp=QUARANTINE dis=NONE) header.from=gmail.com; dara=fail header.i=@gmail.com Received: from [127.0.1.1] (localhost [127.0.0.1]) by ffbox0-bg.mplayerhq.hu (Postfix) with ESMTP id 9389D68DE55; Thu, 24 Oct 2024 02:03:44 +0300 (EEST) X-Original-To: ffmpeg-devel@ffmpeg.org Delivered-To: ffmpeg-devel@ffmpeg.org Received: from mail-ed1-f41.google.com (mail-ed1-f41.google.com [209.85.208.41]) by ffbox0-bg.mplayerhq.hu (Postfix) with ESMTPS id 1CEC468DE44 for ; Thu, 24 Oct 2024 02:03:37 +0300 (EEST) Received: by mail-ed1-f41.google.com with SMTP id 4fb4d7f45d1cf-5c9454f3bfaso262088a12.2 for ; Wed, 23 Oct 2024 16:03:37 -0700 (PDT) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=gmail.com; s=20230601; t=1729724616; x=1730329416; darn=ffmpeg.org; h=content-transfer-encoding:mime-version:message-id:date:subject:cc :to:from:from:to:cc:subject:date:message-id:reply-to; bh=UZw4B/LzQu8JyQtmjcMyOu5vATvqQBL9bwR/R0BDK2U=; b=Mo9z1LBkzxssauHuC4irstrCY20Gwjtw92C3Wi/PLiC7rXq7TrZ9Kh6gYI4+vVdiOn 8MDkNU9fiO1OO3+0ua5C9MRx8PI3n+lZLbzn7DMnc1igSLN8XzokrW0mvm7vdtZh/r4B wRXQZ4cIS080M9awVFZ+QgzU2I499WEOhz+/pAIV61B6o0LdypfyZ+JHmLnXmiZxXpUF dqygeGbNXNoYTp5HOnTi3f82x18iTkKuj5T4W1h3P90NBLgcPghBMZdLX9J/JKOWBhkc sB9pfd8Hv+kx/de3WRi1gPsj8sDrg0ZgBdX4cemmfU3QxTGJtkfNUq+27wN4rY4Ut60A MYUA== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1729724616; x=1730329416; h=content-transfer-encoding:mime-version:message-id:date:subject:cc :to:from:x-gm-message-state:from:to:cc:subject:date:message-id :reply-to; bh=UZw4B/LzQu8JyQtmjcMyOu5vATvqQBL9bwR/R0BDK2U=; b=hgrlb2G2uj1bXUaIhLakb0ElSHMRghSJoOnfsi8Ibr/dmwLA17mN7NXNUbc62YdF/4 PWmdnu4AnWiWyOOp3iQNPb3psDgfH8PSNIyLFdnNYS/s2y2+QyGCz4JQLyWIUuFhrH/U yg8A1bC2vy0IfsmVD95fHSi/cw/xt8jzq5G6KqCdtHJE0bNcKq7hYb3tdst9eGqZ4Eit qUExbmJaxa55hKpN0kkMG7LKxyCspIkzEsliAL0WS1U0DB9Pzy1U86pQ5lVzQGjcoFR1 SjPzL7HZz25GIAubrTv+XnTjZnbSMSINC9BYPSRip354qhiU08URyz0lBHYefuUtdgya zlwA== X-Gm-Message-State: AOJu0YzDIrBlKo/bpfIe+Cx0rCRyx1wFsUgKWpjDtZ6UJPYxI/aJhdH/ 13FYDAL/8tZPE3oLF6heNb/uRLn6q1qEEu54c6PP29/GJ7zNyC03XLx7eg== X-Received: by 2002:a05:6402:50ce:b0:5cb:7318:aab7 with SMTP id 4fb4d7f45d1cf-5cb8b263bfbmr2786298a12.27.1729724614602; Wed, 23 Oct 2024 16:03:34 -0700 (PDT) Received: from localhost.localdomain ([2a02:586:492f:c100:4f65:9001:e562:8dd4]) by smtp.gmail.com with ESMTPSA id 4fb4d7f45d1cf-5cb66c6b7d5sm4997904a12.75.2024.10.23.16.03.31 (version=TLS1_3 cipher=TLS_AES_256_GCM_SHA384 bits=256/256); Wed, 23 Oct 2024 16:03:33 -0700 (PDT) From: IndecisiveTurtle To: ffmpeg-devel@ffmpeg.org Date: Thu, 24 Oct 2024 02:03:27 +0300 Message-ID: <20241023230327.381622-1-geoster3d@gmail.com> X-Mailer: git-send-email 2.47.0 MIME-Version: 1.0 Subject: [FFmpeg-devel] [PATCH] avcodec/vc2enc: Initial vulkan VC2 encoder X-BeenThere: ffmpeg-devel@ffmpeg.org X-Mailman-Version: 2.1.29 Precedence: list List-Id: FFmpeg development discussions and patches List-Unsubscribe: , List-Archive: List-Post: List-Help: List-Subscribe: , Reply-To: FFmpeg development discussions and patches Cc: IndecisiveTurtle Errors-To: ffmpeg-devel-bounces@ffmpeg.org Sender: "ffmpeg-devel" X-TUID: djIQ0lXS3BKn Implements a Vulkan based dirac encoder. Supports Haar and Legall wavelets and should work with all wavelet depths. Performance wise, encoding a 1080p 1-minute video is performed in about 2.5 minutes with the cpu encoder running on my Ryzen 5 4600H, while it takes about 30 seconds on my NVIDIA GTX 1650 Haar shader has a subgroup optimized variant that applies when configured wavelet depth allows it lavapipe seems to be bugged for some reason, after a bunch of debugging I'm not quite sure if it's a bug here or in lavapipe. But people probably dont want to use this with a software implementation anyway. --- libavcodec/Makefile | 5 +- libavcodec/allcodecs.c | 1 + libavcodec/vc2enc.c | 501 +-------------- libavcodec/vc2enc_common.c | 368 +++++++++++ libavcodec/vc2enc_common.h | 279 ++++++++ libavcodec/vc2enc_vulkan.c | 781 +++++++++++++++++++++++ libavcodec/vulkan/dwt_haar.comp | 76 +++ libavcodec/vulkan/dwt_haar_subgroup.comp | 94 +++ libavcodec/vulkan/dwt_hor_legall.comp | 61 ++ libavcodec/vulkan/dwt_legall.comp | 74 +++ libavcodec/vulkan/dwt_upload.comp | 45 ++ libavcodec/vulkan/dwt_ver_legall.comp | 55 ++ libavcodec/vulkan/encode.comp | 256 ++++++++ libavcodec/vulkan/slice_sizes.comp | 184 ++++++ 14 files changed, 2279 insertions(+), 501 deletions(-) create mode 100644 libavcodec/vc2enc_common.c create mode 100644 libavcodec/vc2enc_common.h create mode 100644 libavcodec/vc2enc_vulkan.c create mode 100644 libavcodec/vulkan/dwt_haar.comp create mode 100644 libavcodec/vulkan/dwt_haar_subgroup.comp create mode 100644 libavcodec/vulkan/dwt_hor_legall.comp create mode 100644 libavcodec/vulkan/dwt_legall.comp create mode 100644 libavcodec/vulkan/dwt_upload.comp create mode 100644 libavcodec/vulkan/dwt_ver_legall.comp create mode 100644 libavcodec/vulkan/encode.comp create mode 100644 libavcodec/vulkan/slice_sizes.comp diff --git a/libavcodec/Makefile b/libavcodec/Makefile index dd5d0de898..2b197ad386 100644 --- a/libavcodec/Makefile +++ b/libavcodec/Makefile @@ -765,7 +765,10 @@ OBJS-$(CONFIG_VC1_CUVID_DECODER) += cuviddec.o OBJS-$(CONFIG_VC1_MMAL_DECODER) += mmaldec.o OBJS-$(CONFIG_VC1_QSV_DECODER) += qsvdec.o OBJS-$(CONFIG_VC1_V4L2M2M_DECODER) += v4l2_m2m_dec.o -OBJS-$(CONFIG_VC2_ENCODER) += vc2enc.o vc2enc_dwt.o diractab.o +OBJS-$(CONFIG_VC2_ENCODER) += vc2enc.o vc2enc_dwt.o vc2enc_common.o vc2enc_vulkan.o diractab.o \ + vulkan/encode.o vulkan/slice_sizes.o \ + vulkan/dwt_hor_legall.o vulkan/dwt_ver_legall.o \ + vulkan/dwt_upload.o vulkan/dwt_haar.o vulkan/dwt_haar_subgroup.o OBJS-$(CONFIG_VCR1_DECODER) += vcr1.o OBJS-$(CONFIG_VMDAUDIO_DECODER) += vmdaudio.o OBJS-$(CONFIG_VMDVIDEO_DECODER) += vmdvideo.o diff --git a/libavcodec/allcodecs.c b/libavcodec/allcodecs.c index c7e5f9910c..1db44b220f 100644 --- a/libavcodec/allcodecs.c +++ b/libavcodec/allcodecs.c @@ -362,6 +362,7 @@ extern const FFCodec ff_vc1image_decoder; extern const FFCodec ff_vc1_mmal_decoder; extern const FFCodec ff_vc1_qsv_decoder; extern const FFCodec ff_vc1_v4l2m2m_decoder; +extern const FFCodec ff_vc2_vulkan_encoder; extern const FFCodec ff_vc2_encoder; extern const FFCodec ff_vcr1_decoder; extern const FFCodec ff_vmdvideo_decoder; diff --git a/libavcodec/vc2enc.c b/libavcodec/vc2enc.c index b82370a753..54fcbe5f83 100644 --- a/libavcodec/vc2enc.c +++ b/libavcodec/vc2enc.c @@ -29,506 +29,7 @@ #include "put_bits.h" #include "version.h" -#include "vc2enc_dwt.h" -#include "diractab.h" - -/* The limited size resolution of each slice forces us to do this */ -#define SSIZE_ROUND(b) (FFALIGN((b), s->size_scaler) + 4 + s->prefix_bytes) - -/* Decides the cutoff point in # of slices to distribute the leftover bytes */ -#define SLICE_REDIST_TOTAL 150 - -typedef struct VC2BaseVideoFormat { - enum AVPixelFormat pix_fmt; - AVRational time_base; - int width, height; - uint8_t interlaced, level; - char name[13]; -} VC2BaseVideoFormat; - -static const VC2BaseVideoFormat base_video_fmts[] = { - { 0 }, /* Custom format, here just to make indexing equal to base_vf */ - { AV_PIX_FMT_YUV420P, { 1001, 15000 }, 176, 120, 0, 1, "QSIF525" }, - { AV_PIX_FMT_YUV420P, { 2, 25 }, 176, 144, 0, 1, "QCIF" }, - { AV_PIX_FMT_YUV420P, { 1001, 15000 }, 352, 240, 0, 1, "SIF525" }, - { AV_PIX_FMT_YUV420P, { 2, 25 }, 352, 288, 0, 1, "CIF" }, - { AV_PIX_FMT_YUV420P, { 1001, 15000 }, 704, 480, 0, 1, "4SIF525" }, - { AV_PIX_FMT_YUV420P, { 2, 25 }, 704, 576, 0, 1, "4CIF" }, - - { AV_PIX_FMT_YUV422P10, { 1001, 30000 }, 720, 480, 1, 2, "SD480I-60" }, - { AV_PIX_FMT_YUV422P10, { 1, 25 }, 720, 576, 1, 2, "SD576I-50" }, - - { AV_PIX_FMT_YUV422P10, { 1001, 60000 }, 1280, 720, 0, 3, "HD720P-60" }, - { AV_PIX_FMT_YUV422P10, { 1, 50 }, 1280, 720, 0, 3, "HD720P-50" }, - { AV_PIX_FMT_YUV422P10, { 1001, 30000 }, 1920, 1080, 1, 3, "HD1080I-60" }, - { AV_PIX_FMT_YUV422P10, { 1, 25 }, 1920, 1080, 1, 3, "HD1080I-50" }, - { AV_PIX_FMT_YUV422P10, { 1001, 60000 }, 1920, 1080, 0, 3, "HD1080P-60" }, - { AV_PIX_FMT_YUV422P10, { 1, 50 }, 1920, 1080, 0, 3, "HD1080P-50" }, - - { AV_PIX_FMT_YUV444P12, { 1, 24 }, 2048, 1080, 0, 4, "DC2K" }, - { AV_PIX_FMT_YUV444P12, { 1, 24 }, 4096, 2160, 0, 5, "DC4K" }, - - { AV_PIX_FMT_YUV422P10, { 1001, 60000 }, 3840, 2160, 0, 6, "UHDTV 4K-60" }, - { AV_PIX_FMT_YUV422P10, { 1, 50 }, 3840, 2160, 0, 6, "UHDTV 4K-50" }, - - { AV_PIX_FMT_YUV422P10, { 1001, 60000 }, 7680, 4320, 0, 7, "UHDTV 8K-60" }, - { AV_PIX_FMT_YUV422P10, { 1, 50 }, 7680, 4320, 0, 7, "UHDTV 8K-50" }, - - { AV_PIX_FMT_YUV422P10, { 1001, 24000 }, 1920, 1080, 0, 3, "HD1080P-24" }, - { AV_PIX_FMT_YUV422P10, { 1001, 30000 }, 720, 486, 1, 2, "SD Pro486" }, -}; -static const int base_video_fmts_len = FF_ARRAY_ELEMS(base_video_fmts); - -enum VC2_QM { - VC2_QM_DEF = 0, - VC2_QM_COL, - VC2_QM_FLAT, - - VC2_QM_NB -}; - -typedef struct SubBand { - dwtcoef *buf; - ptrdiff_t stride; - int width; - int height; -} SubBand; - -typedef struct Plane { - SubBand band[MAX_DWT_LEVELS][4]; - dwtcoef *coef_buf; - int width; - int height; - int dwt_width; - int dwt_height; - ptrdiff_t coef_stride; -} Plane; - -typedef struct SliceArgs { - const struct VC2EncContext *ctx; - union { - int cache[DIRAC_MAX_QUANT_INDEX]; - uint8_t *buf; - }; - int x; - int y; - int quant_idx; - int bits_ceil; - int bits_floor; - int bytes; -} SliceArgs; - -typedef struct TransformArgs { - const struct VC2EncContext *ctx; - Plane *plane; - const void *idata; - ptrdiff_t istride; - int field; - VC2TransformContext t; -} TransformArgs; - -typedef struct VC2EncContext { - AVClass *av_class; - PutBitContext pb; - Plane plane[3]; - AVCodecContext *avctx; - DiracVersionInfo ver; - - SliceArgs *slice_args; - TransformArgs transform_args[3]; - - /* For conversion from unsigned pixel values to signed */ - int diff_offset; - int bpp; - int bpp_idx; - - /* Picture number */ - uint32_t picture_number; - - /* Base video format */ - int base_vf; - int level; - int profile; - - /* Quantization matrix */ - uint8_t quant[MAX_DWT_LEVELS][4]; - int custom_quant_matrix; - - /* Division LUT */ - uint32_t qmagic_lut[116][2]; - - int num_x; /* #slices horizontally */ - int num_y; /* #slices vertically */ - int prefix_bytes; - int size_scaler; - int chroma_x_shift; - int chroma_y_shift; - - /* Rate control stuff */ - int frame_max_bytes; - int slice_max_bytes; - int slice_min_bytes; - int q_ceil; - int q_avg; - - /* Options */ - double tolerance; - int wavelet_idx; - int wavelet_depth; - int strict_compliance; - int slice_height; - int slice_width; - int interlaced; - enum VC2_QM quant_matrix; - - /* Parse code state */ - uint32_t next_parse_offset; - enum DiracParseCodes last_parse_code; -} VC2EncContext; - -static av_always_inline void put_vc2_ue_uint(PutBitContext *pb, uint32_t val) -{ - int i; - int bits = 0; - unsigned topbit = 1, maxval = 1; - uint64_t pbits = 0; - - if (!val++) { - put_bits(pb, 1, 1); - return; - } - - while (val > maxval) { - topbit <<= 1; - maxval <<= 1; - maxval |= 1; - } - - bits = ff_log2(topbit); - - for (i = 0; i < bits; i++) { - topbit >>= 1; - av_assert2(pbits <= UINT64_MAX>>3); - pbits <<= 2; - if (val & topbit) - pbits |= 0x1; - } - - put_bits64(pb, bits*2 + 1, (pbits << 1) | 1); -} - -static av_always_inline int count_vc2_ue_uint(uint32_t val) -{ - int topbit = 1, maxval = 1; - - if (!val++) - return 1; - - while (val > maxval) { - topbit <<= 1; - maxval <<= 1; - maxval |= 1; - } - - return ff_log2(topbit)*2 + 1; -} - -/* VC-2 10.4 - parse_info() */ -static void encode_parse_info(VC2EncContext *s, enum DiracParseCodes pcode) -{ - uint32_t cur_pos, dist; - - align_put_bits(&s->pb); - - cur_pos = put_bytes_count(&s->pb, 0); - - /* Magic string */ - ff_put_string(&s->pb, "BBCD", 0); - - /* Parse code */ - put_bits(&s->pb, 8, pcode); - - /* Next parse offset */ - dist = cur_pos - s->next_parse_offset; - AV_WB32(s->pb.buf + s->next_parse_offset + 5, dist); - s->next_parse_offset = cur_pos; - put_bits32(&s->pb, pcode == DIRAC_PCODE_END_SEQ ? 13 : 0); - - /* Last parse offset */ - put_bits32(&s->pb, s->last_parse_code == DIRAC_PCODE_END_SEQ ? 13 : dist); - - s->last_parse_code = pcode; -} - -/* VC-2 11.1 - parse_parameters() - * The level dictates what the decoder should expect in terms of resolution - * and allows it to quickly reject whatever it can't support. Remember, - * this codec kinda targets cheapo FPGAs without much memory. Unfortunately - * it also limits us greatly in our choice of formats, hence the flag to disable - * strict_compliance */ -static void encode_parse_params(VC2EncContext *s) -{ - put_vc2_ue_uint(&s->pb, s->ver.major); /* VC-2 demands this to be 2 */ - put_vc2_ue_uint(&s->pb, s->ver.minor); /* ^^ and this to be 0 */ - put_vc2_ue_uint(&s->pb, s->profile); /* 3 to signal HQ profile */ - put_vc2_ue_uint(&s->pb, s->level); /* 3 - 1080/720, 6 - 4K */ -} - -/* VC-2 11.3 - frame_size() */ -static void encode_frame_size(VC2EncContext *s) -{ - put_bits(&s->pb, 1, !s->strict_compliance); - if (!s->strict_compliance) { - AVCodecContext *avctx = s->avctx; - put_vc2_ue_uint(&s->pb, avctx->width); - put_vc2_ue_uint(&s->pb, avctx->height); - } -} - -/* VC-2 11.3.3 - color_diff_sampling_format() */ -static void encode_sample_fmt(VC2EncContext *s) -{ - put_bits(&s->pb, 1, !s->strict_compliance); - if (!s->strict_compliance) { - int idx; - if (s->chroma_x_shift == 1 && s->chroma_y_shift == 0) - idx = 1; /* 422 */ - else if (s->chroma_x_shift == 1 && s->chroma_y_shift == 1) - idx = 2; /* 420 */ - else - idx = 0; /* 444 */ - put_vc2_ue_uint(&s->pb, idx); - } -} - -/* VC-2 11.3.4 - scan_format() */ -static void encode_scan_format(VC2EncContext *s) -{ - put_bits(&s->pb, 1, !s->strict_compliance); - if (!s->strict_compliance) - put_vc2_ue_uint(&s->pb, s->interlaced); -} - -/* VC-2 11.3.5 - frame_rate() */ -static void encode_frame_rate(VC2EncContext *s) -{ - put_bits(&s->pb, 1, !s->strict_compliance); - if (!s->strict_compliance) { - AVCodecContext *avctx = s->avctx; - put_vc2_ue_uint(&s->pb, 0); - put_vc2_ue_uint(&s->pb, avctx->time_base.den); - put_vc2_ue_uint(&s->pb, avctx->time_base.num); - } -} - -/* VC-2 11.3.6 - aspect_ratio() */ -static void encode_aspect_ratio(VC2EncContext *s) -{ - put_bits(&s->pb, 1, !s->strict_compliance); - if (!s->strict_compliance) { - AVCodecContext *avctx = s->avctx; - put_vc2_ue_uint(&s->pb, 0); - put_vc2_ue_uint(&s->pb, avctx->sample_aspect_ratio.num); - put_vc2_ue_uint(&s->pb, avctx->sample_aspect_ratio.den); - } -} - -/* VC-2 11.3.7 - clean_area() */ -static void encode_clean_area(VC2EncContext *s) -{ - put_bits(&s->pb, 1, 0); -} - -/* VC-2 11.3.8 - signal_range() */ -static void encode_signal_range(VC2EncContext *s) -{ - put_bits(&s->pb, 1, !s->strict_compliance); - if (!s->strict_compliance) - put_vc2_ue_uint(&s->pb, s->bpp_idx); -} - -/* VC-2 11.3.9 - color_spec() */ -static void encode_color_spec(VC2EncContext *s) -{ - AVCodecContext *avctx = s->avctx; - put_bits(&s->pb, 1, !s->strict_compliance); - if (!s->strict_compliance) { - int val; - put_vc2_ue_uint(&s->pb, 0); - - /* primaries */ - put_bits(&s->pb, 1, 1); - if (avctx->color_primaries == AVCOL_PRI_BT470BG) - val = 2; - else if (avctx->color_primaries == AVCOL_PRI_SMPTE170M) - val = 1; - else if (avctx->color_primaries == AVCOL_PRI_SMPTE240M) - val = 1; - else - val = 0; - put_vc2_ue_uint(&s->pb, val); - - /* color matrix */ - put_bits(&s->pb, 1, 1); - if (avctx->colorspace == AVCOL_SPC_RGB) - val = 3; - else if (avctx->colorspace == AVCOL_SPC_YCOCG) - val = 2; - else if (avctx->colorspace == AVCOL_SPC_BT470BG) - val = 1; - else - val = 0; - put_vc2_ue_uint(&s->pb, val); - - /* transfer function */ - put_bits(&s->pb, 1, 1); - if (avctx->color_trc == AVCOL_TRC_LINEAR) - val = 2; - else if (avctx->color_trc == AVCOL_TRC_BT1361_ECG) - val = 1; - else - val = 0; - put_vc2_ue_uint(&s->pb, val); - } -} - -/* VC-2 11.3 - source_parameters() */ -static void encode_source_params(VC2EncContext *s) -{ - encode_frame_size(s); - encode_sample_fmt(s); - encode_scan_format(s); - encode_frame_rate(s); - encode_aspect_ratio(s); - encode_clean_area(s); - encode_signal_range(s); - encode_color_spec(s); -} - -/* VC-2 11 - sequence_header() */ -static void encode_seq_header(VC2EncContext *s) -{ - align_put_bits(&s->pb); - encode_parse_params(s); - put_vc2_ue_uint(&s->pb, s->base_vf); - encode_source_params(s); - put_vc2_ue_uint(&s->pb, s->interlaced); /* Frames or fields coding */ -} - -/* VC-2 12.1 - picture_header() */ -static void encode_picture_header(VC2EncContext *s) -{ - align_put_bits(&s->pb); - put_bits32(&s->pb, s->picture_number++); -} - -/* VC-2 12.3.4.1 - slice_parameters() */ -static void encode_slice_params(VC2EncContext *s) -{ - put_vc2_ue_uint(&s->pb, s->num_x); - put_vc2_ue_uint(&s->pb, s->num_y); - put_vc2_ue_uint(&s->pb, s->prefix_bytes); - put_vc2_ue_uint(&s->pb, s->size_scaler); -} - -/* 1st idx = LL, second - vertical, third - horizontal, fourth - total */ -static const uint8_t vc2_qm_col_tab[][4] = { - {20, 9, 15, 4}, - { 0, 6, 6, 4}, - { 0, 3, 3, 5}, - { 0, 3, 5, 1}, - { 0, 11, 10, 11} -}; - -static const uint8_t vc2_qm_flat_tab[][4] = { - { 0, 0, 0, 0}, - { 0, 0, 0, 0}, - { 0, 0, 0, 0}, - { 0, 0, 0, 0}, - { 0, 0, 0, 0} -}; - -static void init_quant_matrix(VC2EncContext *s) -{ - int level, orientation; - - if (s->wavelet_depth <= 4 && s->quant_matrix == VC2_QM_DEF) { - s->custom_quant_matrix = 0; - for (level = 0; level < s->wavelet_depth; level++) { - s->quant[level][0] = ff_dirac_default_qmat[s->wavelet_idx][level][0]; - s->quant[level][1] = ff_dirac_default_qmat[s->wavelet_idx][level][1]; - s->quant[level][2] = ff_dirac_default_qmat[s->wavelet_idx][level][2]; - s->quant[level][3] = ff_dirac_default_qmat[s->wavelet_idx][level][3]; - } - return; - } - - s->custom_quant_matrix = 1; - - if (s->quant_matrix == VC2_QM_DEF) { - for (level = 0; level < s->wavelet_depth; level++) { - for (orientation = 0; orientation < 4; orientation++) { - if (level <= 3) - s->quant[level][orientation] = ff_dirac_default_qmat[s->wavelet_idx][level][orientation]; - else - s->quant[level][orientation] = vc2_qm_col_tab[level][orientation]; - } - } - } else if (s->quant_matrix == VC2_QM_COL) { - for (level = 0; level < s->wavelet_depth; level++) { - for (orientation = 0; orientation < 4; orientation++) { - s->quant[level][orientation] = vc2_qm_col_tab[level][orientation]; - } - } - } else { - for (level = 0; level < s->wavelet_depth; level++) { - for (orientation = 0; orientation < 4; orientation++) { - s->quant[level][orientation] = vc2_qm_flat_tab[level][orientation]; - } - } - } -} - -/* VC-2 12.3.4.2 - quant_matrix() */ -static void encode_quant_matrix(VC2EncContext *s) -{ - int level; - put_bits(&s->pb, 1, s->custom_quant_matrix); - if (s->custom_quant_matrix) { - put_vc2_ue_uint(&s->pb, s->quant[0][0]); - for (level = 0; level < s->wavelet_depth; level++) { - put_vc2_ue_uint(&s->pb, s->quant[level][1]); - put_vc2_ue_uint(&s->pb, s->quant[level][2]); - put_vc2_ue_uint(&s->pb, s->quant[level][3]); - } - } -} - -/* VC-2 12.3 - transform_parameters() */ -static void encode_transform_params(VC2EncContext *s) -{ - put_vc2_ue_uint(&s->pb, s->wavelet_idx); - put_vc2_ue_uint(&s->pb, s->wavelet_depth); - - encode_slice_params(s); - encode_quant_matrix(s); -} - -/* VC-2 12.2 - wavelet_transform() */ -static void encode_wavelet_transform(VC2EncContext *s) -{ - encode_transform_params(s); - align_put_bits(&s->pb); -} - -/* VC-2 12 - picture_parse() */ -static void encode_picture_start(VC2EncContext *s) -{ - align_put_bits(&s->pb); - encode_picture_header(s); - align_put_bits(&s->pb); - encode_wavelet_transform(s); -} +#include "vc2enc_common.h" #define QUANT(c, mul, add, shift) (((mul) * (c) + (add)) >> (shift)) diff --git a/libavcodec/vc2enc_common.c b/libavcodec/vc2enc_common.c new file mode 100644 index 0000000000..3cc59c4b62 --- /dev/null +++ b/libavcodec/vc2enc_common.c @@ -0,0 +1,368 @@ +/* + * Copyright (C) 2016 Open Broadcast Systems Ltd. + * Author 2016 Rostislav Pehlivanov + * + * 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 "vc2enc_common.h" + +void put_vc2_ue_uint(PutBitContext *pb, uint32_t val) +{ + int i; + int bits = 0; + unsigned topbit = 1, maxval = 1; + uint64_t pbits = 0; + + if (!val++) { + put_bits(pb, 1, 1); + return; + } + + while (val > maxval) { + topbit <<= 1; + maxval <<= 1; + maxval |= 1; + } + + bits = ff_log2(topbit); + + for (i = 0; i < bits; i++) { + topbit >>= 1; + av_assert2(pbits <= UINT64_MAX>>3); + pbits <<= 2; + if (val & topbit) + pbits |= 0x1; + } + + put_bits64(pb, bits*2 + 1, (pbits << 1) | 1); +} + +int count_vc2_ue_uint(uint32_t val) +{ + int topbit = 1, maxval = 1; + + if (!val++) + return 1; + + while (val > maxval) { + topbit <<= 1; + maxval <<= 1; + maxval |= 1; + } + + return ff_log2(topbit)*2 + 1; +} + +/* VC-2 10.4 - parse_info() */ +void encode_parse_info(VC2EncContext *s, enum DiracParseCodes pcode) +{ + uint32_t cur_pos, dist; + + align_put_bits(&s->pb); + + cur_pos = put_bytes_count(&s->pb, 0); + + /* Magic string */ + ff_put_string(&s->pb, "BBCD", 0); + + /* Parse code */ + put_bits(&s->pb, 8, pcode); + + /* Next parse offset */ + dist = cur_pos - s->next_parse_offset; + AV_WB32(s->pb.buf + s->next_parse_offset + 5, dist); + s->next_parse_offset = cur_pos; + put_bits32(&s->pb, pcode == DIRAC_PCODE_END_SEQ ? 13 : 0); + + cur_pos = put_bytes_count(&s->pb, 0); + + /* Last parse offset */ + put_bits32(&s->pb, s->last_parse_code == DIRAC_PCODE_END_SEQ ? 13 : dist); + + s->last_parse_code = pcode; +} + +/* VC-2 11.1 - parse_parameters() + * The level dictates what the decoder should expect in terms of resolution + * and allows it to quickly reject whatever it can't support. Remember, + * this codec kinda targets cheapo FPGAs without much memory. Unfortunately + * it also limits us greatly in our choice of formats, hence the flag to disable + * strict_compliance */ +static void encode_parse_params(VC2EncContext *s) +{ + put_vc2_ue_uint(&s->pb, s->ver.major); /* VC-2 demands this to be 2 */ + put_vc2_ue_uint(&s->pb, s->ver.minor); /* ^^ and this to be 0 */ + put_vc2_ue_uint(&s->pb, s->profile); /* 3 to signal HQ profile */ + put_vc2_ue_uint(&s->pb, s->level); /* 3 - 1080/720, 6 - 4K */ +} + +/* VC-2 11.3 - frame_size() */ +static void encode_frame_size(VC2EncContext *s) +{ + put_bits(&s->pb, 1, !s->strict_compliance); + if (!s->strict_compliance) { + AVCodecContext *avctx = s->avctx; + put_vc2_ue_uint(&s->pb, avctx->width); + put_vc2_ue_uint(&s->pb, avctx->height); + } +} + +/* VC-2 11.3.3 - color_diff_sampling_format() */ +static void encode_sample_fmt(VC2EncContext *s) +{ + put_bits(&s->pb, 1, !s->strict_compliance); + if (!s->strict_compliance) { + int idx; + if (s->chroma_x_shift == 1 && s->chroma_y_shift == 0) + idx = 1; /* 422 */ + else if (s->chroma_x_shift == 1 && s->chroma_y_shift == 1) + idx = 2; /* 420 */ + else + idx = 0; /* 444 */ + put_vc2_ue_uint(&s->pb, idx); + } +} + +/* VC-2 11.3.4 - scan_format() */ +static void encode_scan_format(VC2EncContext *s) +{ + put_bits(&s->pb, 1, !s->strict_compliance); + if (!s->strict_compliance) + put_vc2_ue_uint(&s->pb, s->interlaced); +} + +/* VC-2 11.3.5 - frame_rate() */ +static void encode_frame_rate(VC2EncContext *s) +{ + put_bits(&s->pb, 1, !s->strict_compliance); + if (!s->strict_compliance) { + AVCodecContext *avctx = s->avctx; + put_vc2_ue_uint(&s->pb, 0); + put_vc2_ue_uint(&s->pb, avctx->time_base.den); + put_vc2_ue_uint(&s->pb, avctx->time_base.num); + } +} + +/* VC-2 11.3.6 - aspect_ratio() */ +static void encode_aspect_ratio(VC2EncContext *s) +{ + put_bits(&s->pb, 1, !s->strict_compliance); + if (!s->strict_compliance) { + AVCodecContext *avctx = s->avctx; + put_vc2_ue_uint(&s->pb, 0); + put_vc2_ue_uint(&s->pb, avctx->sample_aspect_ratio.num); + put_vc2_ue_uint(&s->pb, avctx->sample_aspect_ratio.den); + } +} + +/* VC-2 11.3.7 - clean_area() */ +static void encode_clean_area(VC2EncContext *s) +{ + put_bits(&s->pb, 1, 0); +} + +/* VC-2 11.3.8 - signal_range() */ +static void encode_signal_range(VC2EncContext *s) +{ + put_bits(&s->pb, 1, !s->strict_compliance); + if (!s->strict_compliance) + put_vc2_ue_uint(&s->pb, s->bpp_idx); +} + +/* VC-2 11.3.9 - color_spec() */ +static void encode_color_spec(VC2EncContext *s) +{ + AVCodecContext *avctx = s->avctx; + put_bits(&s->pb, 1, !s->strict_compliance); + if (!s->strict_compliance) { + int val; + put_vc2_ue_uint(&s->pb, 0); + + /* primaries */ + put_bits(&s->pb, 1, 1); + if (avctx->color_primaries == AVCOL_PRI_BT470BG) + val = 2; + else if (avctx->color_primaries == AVCOL_PRI_SMPTE170M) + val = 1; + else if (avctx->color_primaries == AVCOL_PRI_SMPTE240M) + val = 1; + else + val = 0; + put_vc2_ue_uint(&s->pb, val); + + /* color matrix */ + put_bits(&s->pb, 1, 1); + if (avctx->colorspace == AVCOL_SPC_RGB) + val = 3; + else if (avctx->colorspace == AVCOL_SPC_YCOCG) + val = 2; + else if (avctx->colorspace == AVCOL_SPC_BT470BG) + val = 1; + else + val = 0; + put_vc2_ue_uint(&s->pb, val); + + /* transfer function */ + put_bits(&s->pb, 1, 1); + if (avctx->color_trc == AVCOL_TRC_LINEAR) + val = 2; + else if (avctx->color_trc == AVCOL_TRC_BT1361_ECG) + val = 1; + else + val = 0; + put_vc2_ue_uint(&s->pb, val); + } +} + +/* VC-2 11.3 - source_parameters() */ +static void encode_source_params(VC2EncContext *s) +{ + encode_frame_size(s); + encode_sample_fmt(s); + encode_scan_format(s); + encode_frame_rate(s); + encode_aspect_ratio(s); + encode_clean_area(s); + encode_signal_range(s); + encode_color_spec(s); +} + +/* VC-2 11 - sequence_header() */ +void encode_seq_header(VC2EncContext *s) +{ + align_put_bits(&s->pb); + encode_parse_params(s); + put_vc2_ue_uint(&s->pb, s->base_vf); + encode_source_params(s); + put_vc2_ue_uint(&s->pb, s->interlaced); /* Frames or fields coding */ +} + +/* VC-2 12.1 - picture_header() */ +static void encode_picture_header(VC2EncContext *s) +{ + align_put_bits(&s->pb); + put_bits32(&s->pb, s->picture_number++); +} + +/* VC-2 12.3.4.1 - slice_parameters() */ +static void encode_slice_params(VC2EncContext *s) +{ + put_vc2_ue_uint(&s->pb, s->num_x); + put_vc2_ue_uint(&s->pb, s->num_y); + put_vc2_ue_uint(&s->pb, s->prefix_bytes); + put_vc2_ue_uint(&s->pb, s->size_scaler); +} + +/* 1st idx = LL, second - vertical, third - horizontal, fourth - total */ +static const uint8_t vc2_qm_col_tab[][4] = { + {20, 9, 15, 4}, + { 0, 6, 6, 4}, + { 0, 3, 3, 5}, + { 0, 3, 5, 1}, + { 0, 11, 10, 11} +}; + +static const uint8_t vc2_qm_flat_tab[][4] = { + { 0, 0, 0, 0}, + { 0, 0, 0, 0}, + { 0, 0, 0, 0}, + { 0, 0, 0, 0}, + { 0, 0, 0, 0} +}; + +void init_quant_matrix(VC2EncContext *s) +{ + int level, orientation; + + if (s->wavelet_depth <= 4 && s->quant_matrix == VC2_QM_DEF) { + s->custom_quant_matrix = 0; + for (level = 0; level < s->wavelet_depth; level++) { + s->quant[level][0] = ff_dirac_default_qmat[s->wavelet_idx][level][0]; + s->quant[level][1] = ff_dirac_default_qmat[s->wavelet_idx][level][1]; + s->quant[level][2] = ff_dirac_default_qmat[s->wavelet_idx][level][2]; + s->quant[level][3] = ff_dirac_default_qmat[s->wavelet_idx][level][3]; + } + return; + } + + s->custom_quant_matrix = 1; + + if (s->quant_matrix == VC2_QM_DEF) { + for (level = 0; level < s->wavelet_depth; level++) { + for (orientation = 0; orientation < 4; orientation++) { + if (level <= 3) + s->quant[level][orientation] = ff_dirac_default_qmat[s->wavelet_idx][level][orientation]; + else + s->quant[level][orientation] = vc2_qm_col_tab[level][orientation]; + } + } + } else if (s->quant_matrix == VC2_QM_COL) { + for (level = 0; level < s->wavelet_depth; level++) { + for (orientation = 0; orientation < 4; orientation++) { + s->quant[level][orientation] = vc2_qm_col_tab[level][orientation]; + } + } + } else { + for (level = 0; level < s->wavelet_depth; level++) { + for (orientation = 0; orientation < 4; orientation++) { + s->quant[level][orientation] = vc2_qm_flat_tab[level][orientation]; + } + } + } +} + +/* VC-2 12.3.4.2 - quant_matrix() */ +static void encode_quant_matrix(VC2EncContext *s) +{ + int level; + put_bits(&s->pb, 1, s->custom_quant_matrix); + if (s->custom_quant_matrix) { + put_vc2_ue_uint(&s->pb, s->quant[0][0]); + for (level = 0; level < s->wavelet_depth; level++) { + put_vc2_ue_uint(&s->pb, s->quant[level][1]); + put_vc2_ue_uint(&s->pb, s->quant[level][2]); + put_vc2_ue_uint(&s->pb, s->quant[level][3]); + } + } +} + +/* VC-2 12.3 - transform_parameters() */ +static void encode_transform_params(VC2EncContext *s) +{ + put_vc2_ue_uint(&s->pb, s->wavelet_idx); + put_vc2_ue_uint(&s->pb, s->wavelet_depth); + + encode_slice_params(s); + encode_quant_matrix(s); +} + +/* VC-2 12.2 - wavelet_transform() */ +static void encode_wavelet_transform(VC2EncContext *s) +{ + encode_transform_params(s); + align_put_bits(&s->pb); +} + +/* VC-2 12 - picture_parse() */ +void encode_picture_start(VC2EncContext *s) +{ + align_put_bits(&s->pb); + encode_picture_header(s); + align_put_bits(&s->pb); + encode_wavelet_transform(s); +} diff --git a/libavcodec/vc2enc_common.h b/libavcodec/vc2enc_common.h new file mode 100644 index 0000000000..dfcd77752c --- /dev/null +++ b/libavcodec/vc2enc_common.h @@ -0,0 +1,279 @@ +/* + * Copyright (C) 2016 Open Broadcast Systems Ltd. + * Author 2016 Rostislav Pehlivanov + * + * 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 AVCODEC_VC2ENC_COMMON_H +#define AVCODEC_VC2ENC_COMMON_H + +#include "avcodec.h" +#include "dirac.h" +#include "put_bits.h" + +#include "vc2enc_dwt.h" +#include "diractab.h" +#include "libavutil/vulkan.h" + +/* The limited size resolution of each slice forces us to do this */ +#define SSIZE_ROUND(b) (FFALIGN((b), s->size_scaler) + 4 + s->prefix_bytes) + +/* Decides the cutoff point in # of slices to distribute the leftover bytes */ +#define SLICE_REDIST_TOTAL 150 + +typedef struct VC2BaseVideoFormat { + enum AVPixelFormat pix_fmt; + AVRational time_base; + int width, height; + uint8_t interlaced, level; + char name[13]; +} VC2BaseVideoFormat; + +static const VC2BaseVideoFormat base_video_fmts[] = { + { 0 }, /* Custom format, here just to make indexing equal to base_vf */ + { AV_PIX_FMT_YUV420P, { 1001, 15000 }, 176, 120, 0, 1, "QSIF525" }, + { AV_PIX_FMT_YUV420P, { 2, 25 }, 176, 144, 0, 1, "QCIF" }, + { AV_PIX_FMT_YUV420P, { 1001, 15000 }, 352, 240, 0, 1, "SIF525" }, + { AV_PIX_FMT_YUV420P, { 2, 25 }, 352, 288, 0, 1, "CIF" }, + { AV_PIX_FMT_YUV420P, { 1001, 15000 }, 704, 480, 0, 1, "4SIF525" }, + { AV_PIX_FMT_YUV420P, { 2, 25 }, 704, 576, 0, 1, "4CIF" }, + + { AV_PIX_FMT_YUV422P10, { 1001, 30000 }, 720, 480, 1, 2, "SD480I-60" }, + { AV_PIX_FMT_YUV422P10, { 1, 25 }, 720, 576, 1, 2, "SD576I-50" }, + + { AV_PIX_FMT_YUV422P10, { 1001, 60000 }, 1280, 720, 0, 3, "HD720P-60" }, + { AV_PIX_FMT_YUV422P10, { 1, 50 }, 1280, 720, 0, 3, "HD720P-50" }, + { AV_PIX_FMT_YUV422P10, { 1001, 30000 }, 1920, 1080, 1, 3, "HD1080I-60" }, + { AV_PIX_FMT_YUV422P10, { 1, 25 }, 1920, 1080, 1, 3, "HD1080I-50" }, + { AV_PIX_FMT_YUV422P10, { 1001, 60000 }, 1920, 1080, 0, 3, "HD1080P-60" }, + { AV_PIX_FMT_YUV422P10, { 1, 50 }, 1920, 1080, 0, 3, "HD1080P-50" }, + + { AV_PIX_FMT_YUV444P12, { 1, 24 }, 2048, 1080, 0, 4, "DC2K" }, + { AV_PIX_FMT_YUV444P12, { 1, 24 }, 4096, 2160, 0, 5, "DC4K" }, + + { AV_PIX_FMT_YUV422P10, { 1001, 60000 }, 3840, 2160, 0, 6, "UHDTV 4K-60" }, + { AV_PIX_FMT_YUV422P10, { 1, 50 }, 3840, 2160, 0, 6, "UHDTV 4K-50" }, + + { AV_PIX_FMT_YUV422P10, { 1001, 60000 }, 7680, 4320, 0, 7, "UHDTV 8K-60" }, + { AV_PIX_FMT_YUV422P10, { 1, 50 }, 7680, 4320, 0, 7, "UHDTV 8K-50" }, + + { AV_PIX_FMT_YUV422P10, { 1001, 24000 }, 1920, 1080, 0, 3, "HD1080P-24" }, + { AV_PIX_FMT_YUV422P10, { 1001, 30000 }, 720, 486, 1, 2, "SD Pro486" }, +}; +static const int base_video_fmts_len = FF_ARRAY_ELEMS(base_video_fmts); + +enum VC2_QM { + VC2_QM_DEF = 0, + VC2_QM_COL, + VC2_QM_FLAT, + + VC2_QM_NB +}; + +typedef struct SubBand { + dwtcoef *buf; + ptrdiff_t stride; + int width; + int height; + int shift; +} SubBand; + +typedef struct Plane { + SubBand band[MAX_DWT_LEVELS][4]; + dwtcoef *coef_buf; + int width; + int height; + int dwt_width; + int dwt_height; + ptrdiff_t coef_stride; +} Plane; + +typedef struct SliceArgs { + const struct VC2EncContext *ctx; + union { + int cache[DIRAC_MAX_QUANT_INDEX]; + uint8_t *buf; + }; + int x; + int y; + int quant_idx; + int bits_ceil; + int bits_floor; + int bytes; +} SliceArgs; + +typedef struct TransformArgs { + struct VC2EncContext *ctx; + Plane *plane; + const void *idata; + ptrdiff_t istride; + int field; + VC2TransformContext t; +} TransformArgs; + +typedef struct VC2DwtPlane { + int width; + int height; + int dwt_width; + int dwt_height; +} VC2DwtPlane; + +typedef struct VC2DwtPushData { + int s; + union { + int diff_offset; + int plane_idx; + }; + int level; + VC2DwtPlane planes[3]; + VkDeviceAddress pbuf[3]; +} VC2DwtPushData; + +typedef struct VC2EncAuxData { + uint32_t quant[MAX_DWT_LEVELS][4]; + int ff_dirac_qscale_tab[116]; +} VC2EncAuxData; + +typedef struct VC2EncPushData { + VkDeviceAddress p[3]; + VkDeviceAddress pb; + VkDeviceAddress luts; + VkDeviceAddress slice; + int num_x; + int num_y; + VC2DwtPlane planes[3]; + int wavelet_depth; + int size_scaler; + int prefix_bytes; +} VC2EncPushData; + +typedef struct VC2EncSliceArgs { + int quant_idx; + int bytes; + int pb_start; + int pad; +} VC2EncSliceArgs; + +typedef struct VC2EncSliceCalcPushData { + VkDeviceAddress p[3]; + VkDeviceAddress luts; + VkDeviceAddress slice; + int num_x; + int num_y; + VC2DwtPlane planes[3]; + int wavelet_depth; + int size_scaler; + int prefix_bytes; + int bits_ceil; + int bits_floor; +} VC2EncSliceCalcPushData; + +typedef struct VC2EncContext { + AVClass *av_class; + PutBitContext pb; + Plane plane[3]; + AVCodecContext *avctx; + DiracVersionInfo ver; + + SliceArgs *slice_args; + VC2EncSliceArgs* vk_slice_args; + TransformArgs transform_args[3]; + + /* For conversion from unsigned pixel values to signed */ + int diff_offset; + int bpp; + int bpp_idx; + + /* Picture number */ + uint32_t picture_number; + + /* Base video format */ + int base_vf; + int level; + int profile; + + /* Quantization matrix */ + uint8_t quant[MAX_DWT_LEVELS][4]; + int custom_quant_matrix; + + /* Division LUT */ + uint32_t qmagic_lut[116][2]; + + int num_x; /* #slices horizontally */ + int num_y; /* #slices vertically */ + int group_x; + int group_y; + int prefix_bytes; + int size_scaler; + int chroma_x_shift; + int chroma_y_shift; + + /* Rate control stuff */ + int frame_max_bytes; + int slice_max_bytes; + int slice_min_bytes; + int q_ceil; + int q_avg; + + /* Options */ + double tolerance; + int wavelet_idx; + int wavelet_depth; + int strict_compliance; + int slice_height; + int slice_width; + int interlaced; + enum VC2_QM quant_matrix; + + /* Parse code state */ + uint32_t next_parse_offset; + enum DiracParseCodes last_parse_code; + + /* Vulkan state */ + FFVulkanContext vkctx; + FFVkQueueFamilyCtx qf; + FFVkExecPool e; + + FFVulkanShader dwt_haar_shd; + FFVulkanShader dwt_upload_shd; + FFVulkanShader dwt_hor_shd, dwt_ver_shd; + FFVulkanShader slice_shd; + FFVulkanShader enc_shd; + AVBufferPool* dwt_buf_pool; + int haar_subgroup; + + VkBuffer plane_buf, slice_buf; + uint32_t buf_plane_size; + VC2EncPushData enc_consts; + VC2DwtPushData dwt_consts; + VC2EncSliceCalcPushData calc_consts; +} VC2EncContext; + +void put_vc2_ue_uint(PutBitContext *pb, uint32_t val); + +int count_vc2_ue_uint(uint32_t val); + +void init_quant_matrix(VC2EncContext *s); + +void encode_parse_info(VC2EncContext *s, enum DiracParseCodes pcode); + +void encode_seq_header(VC2EncContext *s); + +void encode_picture_start(VC2EncContext *s); + +#endif diff --git a/libavcodec/vc2enc_vulkan.c b/libavcodec/vc2enc_vulkan.c new file mode 100644 index 0000000000..09f2cf8fa6 --- /dev/null +++ b/libavcodec/vc2enc_vulkan.c @@ -0,0 +1,781 @@ +/* + * Copyright (C) 2016 Open Broadcast Systems Ltd. + * Author 2016 Rostislav Pehlivanov + * + * 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/avassert.h" +#include "libavutil/mem.h" +#include "libavutil/pixdesc.h" +#include "libavutil/opt.h" +#include "libavutil/version.h" +#include "libavutil/vulkan_spirv.h" +#include "libavutil/hwcontext_vulkan.h" +#include "libavutil/vulkan_loader.h" +#include "libavutil/vulkan.h" +#include "codec_internal.h" +#include "internal.h" +#include "encode.h" +#include "version.h" +#include "vc2enc_common.h" +#include "hwconfig.h" + +#define LEGALL_WORKGROUP_X 64 +#define SLICE_WORKGROUP_X 128 + +extern const char *ff_source_encode_comp; +extern const char *ff_source_dwt_hor_legall_comp; +extern const char *ff_source_dwt_ver_legall_comp; +extern const char *ff_source_slice_sizes_comp; +extern const char *ff_source_dwt_upload_comp; +extern const char *ff_source_dwt_haar_comp; +extern const char *ff_source_dwt_haar_subgroup_comp; + +static int init_vulkan_pipeline(VC2EncContext* s, FFVkSPIRVCompiler *spv, + FFVulkanShader* shd, int push_size, + int lg_x, int lg_y, int lg_z, + const char* pl_name, const char* pl_source, + int plane_img) +{ + uint8_t *spv_data; + size_t spv_len; + void *spv_opaque = NULL; + FFVulkanContext *vkctx = &s->vkctx; + FFVulkanDescriptorSetBinding *desc; + int err = 0; + + ff_vk_shader_init(vkctx, shd, pl_name, VK_SHADER_STAGE_COMPUTE_BIT, + NULL, 0, lg_x, lg_y, lg_z, 0); + + if (plane_img) { + desc = (FFVulkanDescriptorSetBinding []) { + { + .name = "plane_imgs", + .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, + .mem_layout = ff_vk_shader_rep_fmt(vkctx->frames->sw_format, FF_VK_REP_NATIVE), + .dimensions = 2, + .elems = 3, + .stages = VK_SHADER_STAGE_COMPUTE_BIT, + }, + }; + RET(ff_vk_shader_add_descriptor_set(vkctx, shd, desc, 1, 0, 0)); + } + + ff_vk_shader_add_push_const(shd, 0, push_size, VK_SHADER_STAGE_COMPUTE_BIT); + GLSLD(pl_source); + + /* Compile Haar shader */ + RET(spv->compile_shader(vkctx, spv, shd, &spv_data, &spv_len, "main", &spv_opaque)); + RET(ff_vk_shader_link(vkctx, shd, spv_data, spv_len, "main")); + RET(ff_vk_shader_register_exec(vkctx, &s->e, shd)); + +fail: + return err; +} + +static int init_vulkan(AVCodecContext *avctx) +{ + VC2EncContext *s = avctx->priv_data; + FFVulkanContext *vkctx = &s->vkctx; + FFVkSPIRVCompiler *spv; + AVBufferRef* dwt_buf = NULL; + AVBufferRef* coef_buf = NULL; + FFVkBuffer* vk_buf = NULL; + VC2EncAuxData* ad = NULL; + Plane *p; + VC2DwtPlane vk_plane; + int i, level, ret; + uint32_t subgroup_size = vkctx->subgroup_props.maxSubgroupSize; + + /* Initialize spirv compiler */ + spv = ff_vk_spirv_init(); + if (!spv) { + av_log(avctx, AV_LOG_ERROR, "Unable to initialize SPIR-V compiler!\n"); + return -1; + } + + ff_vk_qf_init(vkctx, &s->qf, VK_QUEUE_COMPUTE_BIT); + ff_vk_exec_pool_init(vkctx, &s->qf, &s->e, 1, 0, 0, 0, NULL); + + /* Allocate coefficient buffer for each plane */ + p = &s->plane[0]; + s->buf_plane_size = p->coef_stride*p->dwt_height*sizeof(dwtcoef); + ret = ff_vk_get_pooled_buffer(vkctx, &s->dwt_buf_pool, &coef_buf, + VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | + VK_BUFFER_USAGE_TRANSFER_DST_BIT, NULL, + s->buf_plane_size * 3, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + vk_buf = (FFVkBuffer*)coef_buf->data; + s->plane_buf = vk_buf->buf; + + for (i = 0; i < 3; i++) { + p = &s->plane[i]; + vk_plane.dwt_width = p->dwt_width; + vk_plane.dwt_height = p->dwt_height; + vk_plane.width = p->width; + vk_plane.height = p->height; + memcpy(&s->calc_consts.planes[i], &vk_plane, sizeof(vk_plane)); + memcpy(&s->dwt_consts.planes[i], &vk_plane, sizeof(vk_plane)); + memcpy(&s->enc_consts.planes[i], &vk_plane, sizeof(vk_plane)); + s->enc_consts.p[i] = vk_buf->address + s->buf_plane_size * i; + s->calc_consts.p[i] = vk_buf->address + s->buf_plane_size * i; + s->dwt_consts.pbuf[i] = vk_buf->address + s->buf_plane_size * i; + } + + /* Initialize Haar push data */ + s->dwt_consts.diff_offset = s->diff_offset; + s->dwt_consts.s = s->wavelet_idx == VC2_TRANSFORM_HAAR_S ? 1 : 0; + s->dwt_consts.level = 0; + + /* Initializer slice calc push data */ + s->calc_consts.num_x = s->num_x; + s->calc_consts.num_y = s->num_y; + s->calc_consts.wavelet_depth = s->wavelet_depth; + s->calc_consts.prefix_bytes = s->prefix_bytes; + + /* Initialize encoder push data */ + s->enc_consts.wavelet_depth = s->wavelet_depth; + s->enc_consts.num_x = s->num_x; + s->enc_consts.num_y = s->num_y; + + /* Create buffer for encoder auxilary data. */ + ret = ff_vk_get_pooled_buffer(vkctx, &s->dwt_buf_pool, &dwt_buf, + VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, NULL, + sizeof(VC2EncAuxData), + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT); + vk_buf = (FFVkBuffer*)dwt_buf->data; + s->calc_consts.luts = vk_buf->address; + s->enc_consts.luts = vk_buf->address; + ad = (VC2EncAuxData*)vk_buf->mapped_mem; + if (s->wavelet_depth <= 4 && s->quant_matrix == VC2_QM_DEF) { + s->custom_quant_matrix = 0; + for (level = 0; level < s->wavelet_depth; level++) { + ad->quant[level][0] = ff_dirac_default_qmat[s->wavelet_idx][level][0]; + ad->quant[level][1] = ff_dirac_default_qmat[s->wavelet_idx][level][1]; + ad->quant[level][2] = ff_dirac_default_qmat[s->wavelet_idx][level][2]; + ad->quant[level][3] = ff_dirac_default_qmat[s->wavelet_idx][level][3]; + } + } + memcpy(ad->ff_dirac_qscale_tab, ff_dirac_qscale_tab, sizeof(ff_dirac_qscale_tab)); + + /* Create buffer for slice arguments */ + ret = ff_vk_get_pooled_buffer(vkctx, &s->dwt_buf_pool, &dwt_buf, + VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, NULL, + sizeof(VC2EncSliceArgs) * s->num_x * s->num_y, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT); + vk_buf = (FFVkBuffer*)dwt_buf->data; + s->slice_buf = vk_buf->buf; + s->vk_slice_args = (VC2EncSliceArgs*)vk_buf->mapped_mem; + s->calc_consts.slice = vk_buf->address; + s->enc_consts.slice = vk_buf->address; + + s->haar_subgroup = 0; + + /* Initialize encoding pipelines */ + init_vulkan_pipeline(s, spv, &s->dwt_upload_shd, sizeof(VC2DwtPushData), + 8, 8, 1, "dwt_upload_pl", ff_source_dwt_upload_comp, 1); + init_vulkan_pipeline(s, spv, &s->slice_shd, sizeof(VC2EncPushData), + 128, 1, 1, "slice_pl", ff_source_slice_sizes_comp, 0); + init_vulkan_pipeline(s, spv, &s->enc_shd, sizeof(VC2EncPushData), + 128, 1, 1, "enc_pl", ff_source_encode_comp, 0); + + if (s->wavelet_idx == VC2_TRANSFORM_HAAR || s->wavelet_idx == VC2_TRANSFORM_HAAR_S) { + if (subgroup_size == 32 && s->wavelet_depth < 3) { + init_vulkan_pipeline(s, spv, &s->dwt_haar_shd, sizeof(VC2DwtPushData), + 64, 1, 1, "dwt_haar_pl", ff_source_dwt_haar_subgroup_comp, 0); + s->haar_subgroup = 1; + } else if (subgroup_size == 64 && s->wavelet_depth < 4) { + init_vulkan_pipeline(s, spv, &s->dwt_haar_shd, sizeof(VC2DwtPushData), + 64, 1, 1, "dwt_haar_pl", ff_source_dwt_haar_subgroup_comp, 0); + s->haar_subgroup = 1; + } else { + init_vulkan_pipeline(s, spv, &s->dwt_haar_shd, sizeof(VC2DwtPushData), + 16, 16, 1, "dwt_haar_pl", ff_source_dwt_haar_comp, 0); + } + } else if (s->wavelet_idx == VC2_TRANSFORM_5_3) { + init_vulkan_pipeline(s, spv, &s->dwt_hor_shd, sizeof(VC2DwtPushData), + 64, 1, 1, "dwt_hor_pl", ff_source_dwt_hor_legall_comp, 0); + init_vulkan_pipeline(s, spv, &s->dwt_ver_shd, sizeof(VC2DwtPushData), + 64, 1, 1, "dwt_ver_pl", ff_source_dwt_ver_legall_comp, 0); + } + + s->group_x = s->plane[0].dwt_width >> 3; + s->group_y = s->plane[0].dwt_height >> 3; + return ret; +} + +static void dwt_plane_haar(VC2EncContext *s, FFVkExecContext *exec, VkBufferMemoryBarrier2* buf_bar) +{ + int p, group_x, group_y; + FFVulkanContext *vkctx = &s->vkctx; + FFVulkanFunctions *vk = &vkctx->vkfn; + + s->dwt_consts.level = s->wavelet_depth; + ff_vk_exec_bind_shader(vkctx, exec, &s->dwt_haar_shd); + + /* Haar pass */ + for (p = 0; p < 3; p++) { + s->dwt_consts.plane_idx = p; + if (s->haar_subgroup) { + group_x = FFALIGN(s->plane[p].dwt_width, 8) >> 3; + group_y = FFALIGN(s->plane[p].dwt_height, 8) >> 3; + } else { + group_x = FFALIGN(s->plane[p].dwt_width, 16) >> 4; + group_y = FFALIGN(s->plane[p].dwt_height, 16) >> 4; + } + + ff_vk_shader_update_push_const(vkctx, exec, &s->dwt_haar_shd, VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(VC2DwtPushData), &s->dwt_consts); + vk->CmdDispatch(exec->buf, group_x, group_y, 1); + } + + /* Wait for Haar dispatches to complete */ + vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .pBufferMemoryBarriers = buf_bar, + .bufferMemoryBarrierCount = 1U, + }); +} + +static void dwt_plane_legall(VC2EncContext *s, FFVkExecContext *exec, VkBufferMemoryBarrier2* buf_bar) +{ + int i; + int legall_group_x = (s->plane[0].dwt_height + LEGALL_WORKGROUP_X - 1) >> 6; + int legall_group_y = (s->plane[0].dwt_width + LEGALL_WORKGROUP_X - 1) >> 6; + FFVulkanContext *vkctx = &s->vkctx; + FFVulkanFunctions *vk = &vkctx->vkfn; + + /* Perform Haar wavelet trasform */ + for (i = 0; i < s->wavelet_depth; i++) { + s->dwt_consts.level = i; + + /* Horizontal Haar pass */ + ff_vk_exec_bind_shader(vkctx, exec, &s->dwt_hor_shd); + ff_vk_shader_update_push_const(vkctx, exec, &s->dwt_hor_shd, VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(VC2DwtPushData), &s->dwt_consts); + vk->CmdDispatch(exec->buf, legall_group_x, 1, 3); + vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .pBufferMemoryBarriers = buf_bar, + .bufferMemoryBarrierCount = 1U, + }); + + /* Vertical Haar pass */ + ff_vk_exec_bind_shader(vkctx, exec, &s->dwt_ver_shd); + ff_vk_shader_update_push_const(vkctx, exec, &s->dwt_ver_shd, VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(VC2DwtPushData), &s->dwt_consts); + vk->CmdDispatch(exec->buf, legall_group_y, 1, 3); + vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .pBufferMemoryBarriers = buf_bar, + .bufferMemoryBarrierCount = 1U, + }); + } +} + +static void dwt_plane(VC2EncContext *s, FFVkExecContext *exec, AVFrame *frame) +{ + int i, group_x = s->group_x, group_y = s->group_y; + FFVulkanContext *vkctx = &s->vkctx; + FFVulkanFunctions *vk = &vkctx->vkfn; + uint32_t num_slice_groups = (s->num_x*s->num_y + SLICE_WORKGROUP_X - 1) >> 7; + VkBufferMemoryBarrier2 buf_bar; + VkBufferMemoryBarrier2 slice_buf_bar; + VkImageView views[AV_NUM_DATA_POINTERS]; + VkImageMemoryBarrier2 img_bar[AV_NUM_DATA_POINTERS]; + int nb_img_bar = 0; + + buf_bar = (VkBufferMemoryBarrier2) { + .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2, + .srcStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + .srcAccessMask = VK_ACCESS_2_SHADER_WRITE_BIT | VK_ACCESS_2_SHADER_READ_BIT, + .dstAccessMask = VK_ACCESS_2_SHADER_WRITE_BIT | VK_ACCESS_2_SHADER_READ_BIT, + .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED, + .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED, + .buffer = s->plane_buf, + .size = s->buf_plane_size * 3, + .offset = 0, + }; + + ff_vk_exec_add_dep_frame(vkctx, exec, frame, + VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT, + VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT); + ff_vk_create_imageviews(vkctx, exec, views, frame, FF_VK_REP_UINT); + ff_vk_frame_barrier(vkctx, 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, + }); + + ff_vk_shader_update_img_array(vkctx, exec, &s->dwt_upload_shd, frame, views, 0, 0, + VK_IMAGE_LAYOUT_GENERAL, + VK_NULL_HANDLE); + + /* Upload coefficients from planes to the buffer. */ + s->dwt_consts.diff_offset = s->diff_offset; + ff_vk_exec_bind_shader(vkctx, exec, &s->dwt_upload_shd); + ff_vk_shader_update_push_const(vkctx, exec, &s->dwt_upload_shd, VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(VC2DwtPushData), &s->dwt_consts); + vk->CmdDispatch(exec->buf, group_x, group_y, 3); + vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .pBufferMemoryBarriers = &buf_bar, + .bufferMemoryBarrierCount = 1U, + }); + + /* Perform Haar wavelet trasform */ + if (s->wavelet_idx == VC2_TRANSFORM_HAAR || s->wavelet_idx == VC2_TRANSFORM_HAAR_S) { + dwt_plane_haar(s, exec, &buf_bar); + } else if (s->wavelet_idx == VC2_TRANSFORM_5_3) { + dwt_plane_legall(s, exec, &buf_bar); + } + + /* Calculate slice sizes. */ + ff_vk_exec_bind_shader(vkctx, exec, &s->slice_shd); + ff_vk_shader_update_push_const(vkctx, exec, &s->slice_shd, VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(VC2EncSliceCalcPushData), &s->calc_consts); + vk->CmdDispatch(exec->buf, num_slice_groups, 1, 1); + + slice_buf_bar = (VkBufferMemoryBarrier2) { + .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2, + .srcStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + .srcAccessMask = VK_ACCESS_2_SHADER_READ_BIT | VK_ACCESS_2_SHADER_WRITE_BIT, + .dstAccessMask = VK_ACCESS_2_SHADER_READ_BIT, + .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED, + .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED, + .buffer = s->slice_buf, + .size = sizeof(VC2EncSliceArgs) * s->num_x * s->num_y, + .offset = 0, + }; + vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .pBufferMemoryBarriers = &slice_buf_bar, + .bufferMemoryBarrierCount = 1U, + }); +} + +static void vulkan_encode_slices(VC2EncContext *s, FFVkExecContext *exec) +{ + uint32_t num_slice_groups = (s->num_x*s->num_y + SLICE_WORKGROUP_X - 1) >> 7; + FFVulkanContext *vkctx = &s->vkctx; + FFVulkanFunctions *vk = &vkctx->vkfn; + int skip = 0; + + flush_put_bits(&s->pb); + s->enc_consts.pb += put_bytes_output(&s->pb); + + ff_vk_exec_bind_shader(vkctx, exec, &s->enc_shd); + ff_vk_shader_update_push_const(vkctx, exec, &s->enc_shd, VK_SHADER_STAGE_COMPUTE_BIT, + 0, sizeof(VC2EncPushData), &s->enc_consts); + + vk->CmdDispatch(exec->buf, num_slice_groups, 1, 1); + + ff_vk_exec_submit(vkctx, exec); + ff_vk_exec_wait(vkctx, exec); + + for (int slice_y = 0; slice_y < s->num_y; slice_y++) { + for (int slice_x = 0; slice_x < s->num_x; slice_x++) { + VC2EncSliceArgs *args = &s->vk_slice_args[s->num_x*slice_y + slice_x]; + skip += args->bytes; + } + } + + /* Skip forward to write end header */ + skip_put_bytes(&s->pb, skip); +} + +static int encode_frame(VC2EncContext *s, AVPacket *avpkt, const AVFrame *frame, + const char *aux_data, const int header_size, int field) +{ + int ret; + int64_t max_frame_bytes; + AVBufferRef *avpkt_buf = NULL; + FFVkBuffer* buf_vk = NULL; + FFVulkanContext *vkctx = &s->vkctx; + FFVkExecContext *exec = ff_vk_exec_get(vkctx, &s->e); + + ff_vk_exec_start(vkctx, exec); + + /* Perform Haar DWT pass on the inpute frame. */ + dwt_plane(s, exec, (AVFrame*)frame); + + /* Allocate a buffer that can fit at all all 3 planes of data */ + max_frame_bytes = header_size + s->avctx->width * s->avctx->height * sizeof(dwtcoef); + s->custom_quant_matrix = 0; + + /* Get a pooled device local host visible buffer for writing output data */ + if (field < 2) { + ret = ff_vk_get_pooled_buffer(vkctx, &s->dwt_buf_pool, &avpkt_buf, + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | + VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, NULL, + max_frame_bytes << s->interlaced, + VK_MEMORY_PROPERTY_HOST_CACHED_BIT | + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | + VK_MEMORY_PROPERTY_HOST_COHERENT_BIT); + avpkt->buf = avpkt_buf; + buf_vk = (FFVkBuffer *)avpkt_buf->data; + avpkt->data = buf_vk->mapped_mem; + avpkt->size = max_frame_bytes << s->interlaced; + s->enc_consts.pb = buf_vk->address; + + if (ret < 0) { + return ret; + } + init_put_bits(&s->pb, avpkt->data, avpkt->size); + } + + /* Sequence header */ + encode_parse_info(s, DIRAC_PCODE_SEQ_HEADER); + encode_seq_header(s); + + /* Encoder version */ + if (aux_data) { + encode_parse_info(s, DIRAC_PCODE_AUX); + ff_put_string(&s->pb, aux_data, 1); + } + + /* Picture header */ + encode_parse_info(s, DIRAC_PCODE_PICTURE_HQ); + encode_picture_start(s); + + /* Encode slices */ + vulkan_encode_slices(s, exec); + + /* End sequence */ + encode_parse_info(s, DIRAC_PCODE_END_SEQ); + + return 0; +} + +static av_cold int vc2_encode_frame(AVCodecContext *avctx, AVPacket *avpkt, + const AVFrame *frame, int *got_packet) +{ + int ret = 0; + int slice_ceil, sig_size = 256; + VC2EncContext *s = avctx->priv_data; + const int bitexact = avctx->flags & AV_CODEC_FLAG_BITEXACT; + const char *aux_data = bitexact ? "Lavc" : LIBAVCODEC_IDENT; + const int aux_data_size = bitexact ? sizeof("Lavc") : sizeof(LIBAVCODEC_IDENT); + const int header_size = 100 + aux_data_size; + int64_t r_bitrate = avctx->bit_rate >> (s->interlaced); + + s->avctx = avctx; + s->size_scaler = 2; + s->prefix_bytes = 0; + s->last_parse_code = 0; + s->next_parse_offset = 0; + + /* Rate control */ + s->frame_max_bytes = (av_rescale(r_bitrate, s->avctx->time_base.num, + s->avctx->time_base.den) >> 3) - header_size; + s->slice_max_bytes = slice_ceil = av_rescale(s->frame_max_bytes, 1, s->num_x*s->num_y); + + /* Find an appropriate size scaler */ + while (sig_size > 255) { + int r_size = SSIZE_ROUND(s->slice_max_bytes); + if (r_size > slice_ceil) { + s->slice_max_bytes -= r_size - slice_ceil; + r_size = SSIZE_ROUND(s->slice_max_bytes); + } + sig_size = r_size/s->size_scaler; /* Signalled slize size */ + s->size_scaler <<= 1; + } + + s->slice_min_bytes = s->slice_max_bytes - s->slice_max_bytes*(s->tolerance/100.0f); + if (s->slice_min_bytes < 0) + return AVERROR(EINVAL); + + /* Update slice calc push data */ + s->calc_consts.size_scaler = s->size_scaler; + s->calc_consts.bits_ceil = s->slice_max_bytes << 3; + s->calc_consts.bits_floor = s->slice_min_bytes << 3; + s->enc_consts.prefix_bytes = 0; + s->enc_consts.size_scaler = s->size_scaler; + + ret = encode_frame(s, avpkt, frame, aux_data, header_size, s->interlaced); + if (ret) + return ret; + if (s->interlaced) { + ret = encode_frame(s, avpkt, frame, aux_data, header_size, 2); + if (ret) + return ret; + } + + flush_put_bits(&s->pb); + av_shrink_packet(avpkt, put_bytes_output(&s->pb)); + avpkt->flags |= AV_PKT_FLAG_KEY; + *got_packet = 1; + + return 0; +} + +static av_cold int vc2_encode_end(AVCodecContext *avctx) +{ + int i; + VC2EncContext *s = avctx->priv_data; + + av_log(avctx, AV_LOG_INFO, "Qavg: %i\n", s->q_avg); + + for (i = 0; i < 3; i++) { + ff_vc2enc_free_transforms(&s->transform_args[i].t); + av_freep(&s->plane[i].coef_buf); + } + + return 0; +} + +static av_cold int vc2_encode_init(AVCodecContext *avctx) +{ + Plane *p; + SubBand *b; + int i, level, o, ret, depth; + const AVPixFmtDescriptor *fmt; + VC2EncContext *s = avctx->priv_data; + FFVulkanContext *vkctx = &s->vkctx; + + vkctx->frames_ref = av_buffer_ref(avctx->hw_frames_ctx); + vkctx->frames = (AVHWFramesContext *)vkctx->frames_ref->data; + vkctx->hwfc = vkctx->frames->hwctx; + vkctx->device = (AVHWDeviceContext *)vkctx->frames->device_ref->data; + vkctx->hwctx = vkctx->device->hwctx; + vkctx->extensions = ff_vk_extensions_to_mask(vkctx->hwctx->enabled_dev_extensions, + vkctx->hwctx->nb_enabled_dev_extensions); + ff_vk_load_functions(vkctx->device, &vkctx->vkfn, vkctx->extensions, 1, 1); + ff_vk_load_props(vkctx); + + s->picture_number = 0; + + /* Total allowed quantization range */ + s->q_ceil = DIRAC_MAX_QUANT_INDEX; + + s->ver.major = 2; + s->ver.minor = 0; + s->profile = 3; + s->level = 3; + + s->base_vf = -1; + s->strict_compliance = 1; + + s->q_avg = 0; + s->slice_max_bytes = 0; + s->slice_min_bytes = 0; + + /* Mark unknown as progressive */ + s->interlaced = !((avctx->field_order == AV_FIELD_UNKNOWN) || + (avctx->field_order == AV_FIELD_PROGRESSIVE)); + + for (i = 0; i < base_video_fmts_len; i++) { + const VC2BaseVideoFormat *fmt = &base_video_fmts[i]; + if (avctx->pix_fmt != fmt->pix_fmt) + continue; + if (avctx->time_base.num != fmt->time_base.num) + continue; + if (avctx->time_base.den != fmt->time_base.den) + continue; + if (avctx->width != fmt->width) + continue; + if (avctx->height != fmt->height) + continue; + if (s->interlaced != fmt->interlaced) + continue; + s->base_vf = i; + s->level = base_video_fmts[i].level; + break; + } + + if (s->interlaced) + av_log(avctx, AV_LOG_WARNING, "Interlacing enabled!\n"); + + if ((s->slice_width & (s->slice_width - 1)) || + (s->slice_height & (s->slice_height - 1))) { + av_log(avctx, AV_LOG_ERROR, "Slice size is not a power of two!\n"); + return AVERROR_UNKNOWN; + } + + if ((s->slice_width > avctx->width) || + (s->slice_height > avctx->height)) { + av_log(avctx, AV_LOG_ERROR, "Slice size is bigger than the image!\n"); + return AVERROR_UNKNOWN; + } + + if (s->base_vf <= 0) { + if (avctx->strict_std_compliance < FF_COMPLIANCE_STRICT) { + s->strict_compliance = s->base_vf = 0; + av_log(avctx, AV_LOG_WARNING, "Format does not strictly comply with VC2 specs\n"); + } else { + av_log(avctx, AV_LOG_WARNING, "Given format does not strictly comply with " + "the specifications, decrease strictness to use it.\n"); + return AVERROR_UNKNOWN; + } + } else { + av_log(avctx, AV_LOG_INFO, "Selected base video format = %i (%s)\n", + s->base_vf, base_video_fmts[s->base_vf].name); + } + + /* Chroma subsampling */ + ret = av_pix_fmt_get_chroma_sub_sample(vkctx->frames->sw_format, &s->chroma_x_shift, &s->chroma_y_shift); + if (ret) + return ret; + + /* Bit depth and color range index */ + fmt = av_pix_fmt_desc_get(vkctx->frames->sw_format); + depth = fmt->comp[0].depth; + if (depth == 8 && avctx->color_range == AVCOL_RANGE_JPEG) { + s->bpp = 1; + s->bpp_idx = 1; + s->diff_offset = 128; + } else if (depth == 8 && (avctx->color_range == AVCOL_RANGE_MPEG || + avctx->color_range == AVCOL_RANGE_UNSPECIFIED)) { + s->bpp = 1; + s->bpp_idx = 2; + s->diff_offset = 128; + } else if (depth == 10) { + s->bpp = 2; + s->bpp_idx = 3; + s->diff_offset = 512; + } else { + s->bpp = 2; + s->bpp_idx = 4; + s->diff_offset = 2048; + } + + /* Planes initialization */ + for (i = 0; i < 3; i++) { + int w, h; + p = &s->plane[i]; + p->width = avctx->width >> (i ? s->chroma_x_shift : 0); + p->height = avctx->height >> (i ? s->chroma_y_shift : 0); + if (s->interlaced) + p->height >>= 1; + p->dwt_width = w = FFALIGN(p->width, (1 << s->wavelet_depth)); + p->dwt_height = h = FFALIGN(p->height, (1 << s->wavelet_depth)); + p->coef_stride = FFALIGN(p->dwt_width, 32); + for (level = s->wavelet_depth-1; level >= 0; level--) { + w = w >> 1; + h = h >> 1; + for (o = 0; o < 4; o++) { + b = &p->band[level][o]; + b->width = w; + b->height = h; + b->stride = p->coef_stride; + b->shift = (o > 1)*b->height*b->stride + (o & 1)*b->width; + } + } + + /* DWT init */ + if (ff_vc2enc_init_transforms(&s->transform_args[i].t, + s->plane[i].coef_stride, + s->plane[i].dwt_height, + s->slice_width, s->slice_height)) + return AVERROR(ENOMEM); + } + + /* Slices */ + s->num_x = s->plane[0].dwt_width/s->slice_width; + s->num_y = s->plane[0].dwt_height/s->slice_height; + + s->slice_args = av_calloc(s->num_x*s->num_y, sizeof(SliceArgs)); + if (!s->slice_args) + return AVERROR(ENOMEM); + + for (i = 0; i < 116; i++) { + const uint64_t qf = ff_dirac_qscale_tab[i]; + const uint32_t m = av_log2(qf); + const uint32_t t = (1ULL << (m + 32)) / qf; + const uint32_t r = (t*qf + qf) & UINT32_MAX; + if (!(qf & (qf - 1))) { + s->qmagic_lut[i][0] = 0xFFFFFFFF; + s->qmagic_lut[i][1] = 0xFFFFFFFF; + } else if (r <= 1 << m) { + s->qmagic_lut[i][0] = t + 1; + s->qmagic_lut[i][1] = 0; + } else { + s->qmagic_lut[i][0] = t; + s->qmagic_lut[i][1] = t; + } + } + init_vulkan(avctx); + + return 0; +} + +#define VC2ENC_FLAGS (AV_OPT_FLAG_ENCODING_PARAM | AV_OPT_FLAG_VIDEO_PARAM) +static const AVOption vc2enc_options[] = { + {"tolerance", "Max undershoot in percent", offsetof(VC2EncContext, tolerance), AV_OPT_TYPE_DOUBLE, {.dbl = 5.0f}, 0.0f, 45.0f, VC2ENC_FLAGS, .unit = "tolerance"}, + {"slice_width", "Slice width", offsetof(VC2EncContext, slice_width), AV_OPT_TYPE_INT, {.i64 = 32}, 32, 1024, VC2ENC_FLAGS, .unit = "slice_width"}, + {"slice_height", "Slice height", offsetof(VC2EncContext, slice_height), AV_OPT_TYPE_INT, {.i64 = 16}, 8, 1024, VC2ENC_FLAGS, .unit = "slice_height"}, + {"wavelet_depth", "Transform depth", offsetof(VC2EncContext, wavelet_depth), AV_OPT_TYPE_INT, {.i64 = 4}, 1, 5, VC2ENC_FLAGS, .unit = "wavelet_depth"}, + {"wavelet_type", "Transform type", offsetof(VC2EncContext, wavelet_idx), AV_OPT_TYPE_INT, {.i64 = VC2_TRANSFORM_HAAR_S}, 0, VC2_TRANSFORMS_NB, VC2ENC_FLAGS, .unit = "wavelet_idx"}, + {"5_3", "LeGall (5,3)", 0, AV_OPT_TYPE_CONST, {.i64 = VC2_TRANSFORM_5_3}, INT_MIN, INT_MAX, VC2ENC_FLAGS, .unit = "wavelet_idx"}, + {"haar", "Haar (with shift)", 0, AV_OPT_TYPE_CONST, {.i64 = VC2_TRANSFORM_HAAR_S}, INT_MIN, INT_MAX, VC2ENC_FLAGS, .unit = "wavelet_idx"}, + {"haar_noshift", "Haar (without shift)", 0, AV_OPT_TYPE_CONST, {.i64 = VC2_TRANSFORM_HAAR}, INT_MIN, INT_MAX, VC2ENC_FLAGS, .unit = "wavelet_idx"}, + {"qm", "Custom quantization matrix", offsetof(VC2EncContext, quant_matrix), AV_OPT_TYPE_INT, {.i64 = VC2_QM_DEF}, 0, VC2_QM_NB, VC2ENC_FLAGS, .unit = "quant_matrix"}, + {"default", "Default from the specifications", 0, AV_OPT_TYPE_CONST, {.i64 = VC2_QM_DEF}, INT_MIN, INT_MAX, VC2ENC_FLAGS, .unit = "quant_matrix"}, + {"color", "Prevents low bitrate discoloration", 0, AV_OPT_TYPE_CONST, {.i64 = VC2_QM_COL}, INT_MIN, INT_MAX, VC2ENC_FLAGS, .unit = "quant_matrix"}, + {"flat", "Optimize for PSNR", 0, AV_OPT_TYPE_CONST, {.i64 = VC2_QM_FLAT}, INT_MIN, INT_MAX, VC2ENC_FLAGS, .unit = "quant_matrix"}, + {NULL} +}; + +static const AVClass vc2enc_class = { + .class_name = "vc2_vulkan_encoder", + .category = AV_CLASS_CATEGORY_ENCODER, + .option = vc2enc_options, + .item_name = av_default_item_name, + .version = LIBAVUTIL_VERSION_INT +}; + +static const FFCodecDefault vc2enc_defaults[] = { + { "b", "600000000" }, + { NULL }, +}; + +const AVCodecHWConfigInternal *const ff_vc2_hw_configs[] = { + HW_CONFIG_ENCODER_FRAMES(VULKAN, VULKAN), + HW_CONFIG_ENCODER_DEVICE(NONE, VULKAN), + NULL, +}; + +const FFCodec ff_vc2_vulkan_encoder = { + .p.name = "vc2_vulkan", + CODEC_LONG_NAME("SMPTE VC-2"), + .p.type = AVMEDIA_TYPE_VIDEO, + .p.id = AV_CODEC_ID_DIRAC, + .p.capabilities = AV_CODEC_CAP_DR1 | AV_CODEC_CAP_HARDWARE, + .caps_internal = FF_CODEC_CAP_INIT_CLEANUP, + .priv_data_size = sizeof(VC2EncContext), + .init = vc2_encode_init, + .close = vc2_encode_end, + FF_CODEC_ENCODE_CB(vc2_encode_frame), + .p.priv_class = &vc2enc_class, + .defaults = vc2enc_defaults, + .p.pix_fmts = (const enum AVPixelFormat[]) { + AV_PIX_FMT_VULKAN, + AV_PIX_FMT_NONE, + }, + .hw_configs = ff_vc2_hw_configs, +}; diff --git a/libavcodec/vulkan/dwt_haar.comp b/libavcodec/vulkan/dwt_haar.comp new file mode 100644 index 0000000000..69073cb17f --- /dev/null +++ b/libavcodec/vulkan/dwt_haar.comp @@ -0,0 +1,76 @@ +#extension GL_EXT_scalar_block_layout : require +#extension GL_EXT_buffer_reference : require + +#define LOCAL_X 256 + +layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in; + +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer DwtCoef { + int coef_buf[]; +}; + +struct Plane { + ivec2 dim; + ivec2 dwt_dim; +}; + +layout(push_constant, scalar) uniform ComputeInfo { + int s; + int plane_idx; + int wavelet_depth; + Plane planes[3]; + DwtCoef pbuf[3]; +}; + +shared int local_coef[LOCAL_X]; + +void main() { + ivec2 coord = ivec2(gl_GlobalInvocationID.xy); + ivec2 dwt_dim = planes[plane_idx].dwt_dim; + if (any(greaterThanEqual(coord, dwt_dim))) { + return; + } + int index = dwt_dim.x * coord.y + coord.x; + int value = pbuf[plane_idx].coef_buf[index]; + + /* Perform Haar wavelet on the 16x16 local workgroup with shared memory */ + for (int i = 0; i < wavelet_depth; i++) { + ivec2 mask = ivec2((1 << i) - 1); + if (any(notEqual(coord & mask, ivec2(0)))) { + break; + } + + /* Offset between valid hor pixels for each level, +1, +2, +4 etc */ + int dist = (1 << i); + + local_coef[gl_LocalInvocationIndex] = value; + barrier(); + + /* Horizontal haar wavelet */ + uint other_id = gl_LocalInvocationIndex ^ dist; + int other = local_coef[other_id]; + int a = gl_LocalInvocationIndex < other_id ? value : other; + int b = gl_LocalInvocationIndex < other_id ? other : value; + int dst_b = (b - a) * (1 << s); + int dst_a = a * (1 << s) + ((dst_b + 1) >> 1); + value = gl_LocalInvocationIndex < other_id ? dst_a : dst_b; + + /* Offset between valid ver pixels for each level, +1, +2, +4 etc */ + dist <<= 4; + + local_coef[gl_LocalInvocationIndex] = value; + barrier(); + + /* Vertical haar wavelet */ + other_id = gl_LocalInvocationIndex ^ dist; + other = local_coef[other_id]; + a = gl_LocalInvocationIndex < other_id ? value : other; + b = gl_LocalInvocationIndex < other_id ? other : value; + dst_b = b - a; + dst_a = a + ((dst_b + 1) >> 1); + value = gl_LocalInvocationIndex < other_id ? dst_a : dst_b; + } + + /* Store value */ + pbuf[plane_idx].coef_buf[index] = value; +} \ No newline at end of file diff --git a/libavcodec/vulkan/dwt_haar_subgroup.comp b/libavcodec/vulkan/dwt_haar_subgroup.comp new file mode 100644 index 0000000000..32fb04561a --- /dev/null +++ b/libavcodec/vulkan/dwt_haar_subgroup.comp @@ -0,0 +1,94 @@ +#extension GL_EXT_scalar_block_layout : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_buffer_reference : require + +#define TILE_DIM 8 +#define LOCAL_X 64 + +layout(local_size_x = LOCAL_X, local_size_y = 1, local_size_z = 1) in; + +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer DwtCoef { + int coef_buf[]; +}; + +struct Plane { + ivec2 dim; + ivec2 dwt_dim; +}; + +layout(push_constant, scalar) uniform ComputeInfo { + int s; + int plane_idx; + int wavelet_depth; + Plane planes[3]; + DwtCoef pbuf[3]; +}; + +int dwt_haar_subgroup(int value, int i) { + /* Offset between valid hor pixels for each level, +1, +2, +4 etc */ + int dist = (1 << i); + + /* Horizontal haar wavelet */ + uint other_sub_id = gl_SubgroupInvocationID ^ dist; + int other = subgroupShuffle(value, other_sub_id); + int a = gl_SubgroupInvocationID < other_sub_id ? value : other; + int b = gl_SubgroupInvocationID < other_sub_id ? other : value; + int dst_b = (b - a) * (1 << s); + int dst_a = a * (1 << s) + ((dst_b + 1) >> 1); + value = gl_SubgroupInvocationID < other_sub_id ? dst_a : dst_b; + + /* Offset between valid ver pixels for each level, +1, +2, +4 etc */ + dist <<= 3; + + /* Vertical haar wavelet */ + other_sub_id = gl_SubgroupInvocationID ^ dist; + other = subgroupShuffle(value, other_sub_id); + a = gl_SubgroupInvocationID < other_sub_id ? value : other; + b = gl_SubgroupInvocationID < other_sub_id ? other : value; + dst_b = b - a; + dst_a = a + ((dst_b + 1) >> 1); + return gl_SubgroupInvocationID < other_sub_id ? dst_a : dst_b; +} + +bool is_thread_active(int i, ivec2 coord) { + if (i > wavelet_depth - 1) { + return false; + } + ivec2 mask = ivec2((1 << i) - 1); + if (any(notEqual(coord & mask, ivec2(0)))) { + return false; + } + return true; +} + +void main() { + ivec2 tile_coord = ivec2(gl_WorkGroupID.xy); + ivec2 local_coord = ivec2(gl_LocalInvocationIndex & 7, gl_LocalInvocationIndex >> 3); + ivec2 coord = tile_coord * ivec2(TILE_DIM) + local_coord; + ivec2 dwt_dim = planes[plane_idx].dwt_dim; + if (any(greaterThanEqual(coord, dwt_dim))) { + return; + } + int index = dwt_dim.x * coord.y + coord.x; + int value = pbuf[plane_idx].coef_buf[index]; + + if (gl_SubgroupSize == 64) { + for (int i = 0; i < 3; i++) { + if (!is_thread_active(i, local_coord)) { + break; + } + value = dwt_haar_subgroup(value, i); + } + } else { + for (int i = 0; i < 2; i++) { + if (!is_thread_active(i, local_coord)) { + break; + } + value = dwt_haar_subgroup(value, i); + } + } + + // Store value + pbuf[plane_idx].coef_buf[index] = value; +} \ No newline at end of file diff --git a/libavcodec/vulkan/dwt_hor_legall.comp b/libavcodec/vulkan/dwt_hor_legall.comp new file mode 100644 index 0000000000..9c3945825e --- /dev/null +++ b/libavcodec/vulkan/dwt_hor_legall.comp @@ -0,0 +1,61 @@ +#extension GL_EXT_scalar_block_layout : require +#extension GL_EXT_buffer_reference : require + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer DwtCoef { + int coef_buf[]; +}; + +struct Plane { + ivec2 dim; + ivec2 dwt_dim; +}; + +layout(push_constant, scalar) uniform ComputeInfo { + int s; + int diff_offset; + int level; + Plane planes[3]; + DwtCoef pbuf[3]; +}; + +int align(int x, int a) { + return (x + a - 1) & ~(a - 1); +} + +void main() { + int coord_y = int(gl_GlobalInvocationID.x); + uint plane_idx = gl_GlobalInvocationID.z; + ivec2 work_area = planes[plane_idx].dwt_dim; + int dist = 1 << level; + if (coord_y >= work_area.y || (coord_y & (dist - 1)) != 0) { + return; + } + + DwtCoef buf = pbuf[plane_idx]; + ivec2 dwt_area = work_area >> 1; + int stride = align(planes[plane_idx].dwt_dim.x, 32); + int start = stride * coord_y; + + // Shift in one bit that is used for additional precision + for (int x = 0; x < work_area.x; x += dist) { + buf.coef_buf[start + x] = buf.coef_buf[start + x] << 1; + } + + // Lifting stage 2 + for (int x = 0; x < work_area.x - 2 * dist; x += 2 * dist) { + buf.coef_buf[start + x + dist] -= (buf.coef_buf[start + x] + + buf.coef_buf[start + x + 2 * dist] + 1) >> 1; + } + buf.coef_buf[start + work_area.x - dist] -= (2 * buf.coef_buf[start + work_area.x - 2 * dist] + 1) >> 1; + + // Lifting stage 1 + buf.coef_buf[start] += (2 * buf.coef_buf[start + dist] + 2) >> 2; + for (int x = 2 * dist; x < work_area.x - 2 * dist; x += 2 * dist) { + buf.coef_buf[start + x] += (buf.coef_buf[start + x - dist] + + buf.coef_buf[start + x + dist] + 2) >> 2; + } + buf.coef_buf[start + work_area.x - 2 * dist] += (buf.coef_buf[start + work_area.x - 3 * dist] + + buf.coef_buf[start + work_area.x - dist] + 2) >> 2; +} \ No newline at end of file diff --git a/libavcodec/vulkan/dwt_legall.comp b/libavcodec/vulkan/dwt_legall.comp new file mode 100644 index 0000000000..1c640022f5 --- /dev/null +++ b/libavcodec/vulkan/dwt_legall.comp @@ -0,0 +1,74 @@ +#extension GL_EXT_scalar_block_layout : require +#extension GL_EXT_debug_printf : require + +#define TILE_DIM 16 +#define LOCAL_X 256 + +layout(local_size_x = TILE_DIM, local_size_y = TILE_DIM, local_size_z = 1) in; + +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer DwtCoef { + int coef_buf[]; +}; + +struct Plane { + ivec2 dim; + ivec2 dwt_dim; +}; + +layout(push_constant, scalar) uniform ComputeInfo { + int s; + int plane_idx; + int wavelet_depth; + Plane planes[3]; + DwtCoef dst_buf[3]; + DwtCoef src_buf[3]; +}; + +int align(int x, int a) { + return (x + a - 1) & ~(a - 1); +} + +shared uint local_coef[LOCAL_X]; + +void main() { + ivec2 coord = ivec2(gl_GlobalInvocationID.xy); + ivec2 work_area = planes[plane_idx].dwt_dim; + ivec2 dwt_area = work_area >> 1; + if (any(greaterThanEqual(coord, work_area))) { + return; + } + + DwtCoef src = src_buf[plane_idx]; + DwtCoef dst = dst_buf[plane_idx]; + int stride = align(planes[plane_idx].dwt_dim.x, 32); + int start = stride * coord.y; + + for (int i = 0; i < wavelet_depth; i++) { + ivec2 mask = ivec2((1 << i) - 1); + if (any(notEqual(coord & mask, ivec2(0)))) { + break; + } + + mask <<= 1; + mask |= 1; + + // Shift in one bit that is used for additional precision + for (int x = 0; x < work_area.x; x++) { + dst.coef_buf[start + x] = src.coef_buf[start + x] << 1; + } + + // Lifting stage 2 + for (int x = 0; x < dwt_area.x - 1; x++) { + dst.coef_buf[start + 2 * x + 1] -= (dst.coef_buf[start + 2 * x] + + dst.coef_buf[start + 2 * x + 2] + 1) >> 1; + } + dst.coef_buf[start + work_area.x - 1] -= (2 * dst.coef_buf[start + work_area.x - 2] + 1) >> 1; + + // Lifting stage 1 + dst.coef_buf[start] += (2 * dst.coef_buf[start + 1] + 2) >> 2; + for (int x = 1; x <= dwt_area.x - 1; x++) { + dst.coef_buf[start + 2 * x] += (dst.coef_buf[start + 2 * x - 1] + + dst.coef_buf[start + 2 * x + 1] + 2) >> 2; + } + } +} \ No newline at end of file diff --git a/libavcodec/vulkan/dwt_upload.comp b/libavcodec/vulkan/dwt_upload.comp new file mode 100644 index 0000000000..943ebf23d7 --- /dev/null +++ b/libavcodec/vulkan/dwt_upload.comp @@ -0,0 +1,45 @@ +#extension GL_EXT_scalar_block_layout : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_buffer_reference : require + +layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in; + +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer DwtCoef { + int coef_buf[]; +}; + +layout(scalar, buffer_reference, buffer_reference_align = 1) buffer PlaneBuf { + uint8_t data[]; +}; + +struct Plane { + ivec2 dim; + ivec2 dwt_dim; +}; + +layout(push_constant, scalar) uniform ComputeInfo { + int s; + int diff_offset; + int level; + Plane planes[3]; + DwtCoef pbuf[3]; +}; + +int align(int x, int a) { + return (x + a - 1) & ~(a - 1); +} + +void main() { + ivec2 coord = ivec2(gl_GlobalInvocationID.xy); + uint plane_idx = gl_GlobalInvocationID.z; + ivec2 work_area = planes[plane_idx].dwt_dim; + if (any(greaterThanEqual(coord, work_area))) { + return; + } + int stride = align(planes[plane_idx].dwt_dim.x, 32); + uint coef_idx = coord.y * stride + coord.x; + ivec2 coord_i = clamp(coord, ivec2(0), planes[plane_idx].dim); + uint texel = imageLoad(plane_imgs[plane_idx], coord_i).x; + int result = int(texel - diff_offset); + pbuf[plane_idx].coef_buf[coef_idx] = result; +} \ No newline at end of file diff --git a/libavcodec/vulkan/dwt_ver_legall.comp b/libavcodec/vulkan/dwt_ver_legall.comp new file mode 100644 index 0000000000..6662bd656f --- /dev/null +++ b/libavcodec/vulkan/dwt_ver_legall.comp @@ -0,0 +1,55 @@ +#extension GL_EXT_scalar_block_layout : require +#extension GL_EXT_buffer_reference : require + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer DwtCoef { + int coef_buf[]; +}; + +struct Plane { + ivec2 dim; + ivec2 dwt_dim; +}; + +layout(push_constant, scalar) uniform ComputeInfo { + int s; + int diff_offset; + int level; + Plane planes[3]; + DwtCoef pbuf[3]; +}; + +int align(int x, int a) { + return (x + a - 1) & ~(a - 1); +} + +void main() { + int coord_x = int(gl_GlobalInvocationID.x); + uint plane_idx = gl_GlobalInvocationID.z; + ivec2 work_area = planes[plane_idx].dwt_dim; + int dist = 1 << level; + if (coord_x >= work_area.x || (coord_x & (dist - 1)) != 0) { + return; + } + + DwtCoef buf = pbuf[plane_idx]; + ivec2 dwt_area = work_area >> 1; + int stride = align(planes[plane_idx].dwt_dim.x, 32); + + // Lifting stage 2 + for (int y = dist; y < work_area.y - 2 * dist; y += 2 * dist) { + buf.coef_buf[stride * y + coord_x] -= (buf.coef_buf[stride * (y - dist) + coord_x] + + buf.coef_buf[stride * (y + dist) + coord_x] + 1) >> 1; + } + buf.coef_buf[stride * (work_area.y - dist) + coord_x] -= (2 * buf.coef_buf[stride * (work_area.y - 2 * dist) + coord_x] + 1) >> 1; + + // Lifting stage 1 + buf.coef_buf[coord_x] += (2 * buf.coef_buf[stride * dist + coord_x] + 2) >> 2; + for (int y = 2 * dist; y < work_area.y - 2 * dist; y += 2 * dist) { + buf.coef_buf[stride * y + coord_x] += (buf.coef_buf[stride * (y + dist) + coord_x] + + buf.coef_buf[stride * (y - dist) + coord_x] + 2) >> 2; + } + buf.coef_buf[stride * (work_area.y - 2 * dist) + coord_x] += (buf.coef_buf[stride * (work_area.y - 3 * dist) + coord_x] + + buf.coef_buf[stride * (work_area.y - dist) + coord_x] + 2) >> 2; +} \ No newline at end of file diff --git a/libavcodec/vulkan/encode.comp b/libavcodec/vulkan/encode.comp new file mode 100644 index 0000000000..892674c6da --- /dev/null +++ b/libavcodec/vulkan/encode.comp @@ -0,0 +1,256 @@ +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_scalar_block_layout : require +#extension GL_EXT_buffer_reference : require +#extension GL_EXT_debug_printf : require + +#define WORKGROUP_X 128 +layout(local_size_x = WORKGROUP_X, local_size_y = 1, local_size_z = 1) in; + +#define MAX_DWT_LEVELS (5) + +struct SliceArgs { + int quant_idx; + int bytes; + int pb_start; + int pad; +}; + +struct Plane { + ivec2 dim; + ivec2 dwt_dim; +}; + +layout(std430, buffer_reference, buffer_reference_align = 16) buffer SliceArgBuf { + SliceArgs args[]; +}; +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer DwtCoef { + int coef_buf[]; +}; +layout(scalar, buffer_reference, buffer_reference_align = 1) buffer BitBuf { + uint data[]; +}; +layout(scalar, buffer_reference, buffer_reference_align = 1) buffer BitBufByte { + uint8_t data[]; +}; +layout(scalar, buffer_reference, buffer_reference_align = 1) buffer BitBufQword { + uint64_t data[]; +}; +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer QuantLuts { + int quant[5][4]; + int ff_dirac_qscale_tab[116]; +}; + +layout(push_constant, scalar) uniform ComputeInfo { + DwtCoef plane_dat[3]; + BitBuf pb; + QuantLuts luts; + SliceArgBuf slice; + ivec2 num_slices; + Plane planes[3]; + int wavelet_depth; + int size_scaler; + int prefix_bytes; +}; + +#define BUF_BITS 64 +#define BUF_BYTES 4 + +uint64_t bit_buf = 0; +int bit_left = BUF_BITS; +int write_start = 0; +int write_ptr = 0; + +uint64_t byteswap(uint64_t value) { + u32vec2 val = unpack32(value); + val.x = pack32(unpack8(val.x).wzyx); + val.y = pack32(unpack8(val.y).wzyx); + return pack64(val.yx); +} + +void put_bits(int n, int value) { + if (n < bit_left) { + bit_buf = (bit_buf << n) | value; + bit_left -= n; + } else { + bit_buf <<= bit_left; + bit_buf |= (value >> (n - bit_left)); + uint64_t pbb = uint64_t(pb) + write_ptr * 4; + BitBufQword(pbb).data[0] = byteswap(bit_buf); + write_ptr += 2; + bit_left += BUF_BITS - n; + bit_buf = value; + } +} + +void flush_put_bits() { + int bit_left_lo = (bit_left >> 3) << 3; + bit_buf <<= (bit_left - bit_left_lo); + if (bit_left_lo == 0) { + uint64_t pbb = uint64_t(pb) + write_ptr * 4; + BitBufQword(pbb).data[0] = byteswap(bit_buf); + write_ptr += 2; + bit_left = BUF_BITS; + bit_buf = 0; + } else { + bit_left = bit_left_lo; + } +} + +int put_bytes_count() { + return (write_ptr - write_start) * BUF_BYTES + ((BUF_BITS - bit_left) >> 3); +} + +/* Same as skip_put_bytes in put_bits.h but fills in 0xFF */ +void skip_put_bytes(int n) { + int bytes_left = bit_left >> 3; + if (n < bytes_left) { + int n_bits = n << 3; + int mask = (1 << n_bits) - 1; + bit_buf <<= n_bits; + bit_buf |= mask; + bit_left -= n_bits; + return; + } + if (bit_left < BUF_BITS) { + int mask = (1 << bit_left) - 1; + bit_buf <<= bit_left; + bit_buf |= mask; + uint64_t pbb = uint64_t(pb) + write_ptr * 4; + BitBufQword(pbb).data[0] = byteswap(bit_buf); + write_ptr += 2; + n -= bit_left >> 3; + } + int skip_dwords = n >> 2; + while (skip_dwords > 0) { + pb.data[write_ptr++] = 0xFFFFFFFF; + skip_dwords--; + } + int skip_bits = (n & 3) << 3; + bit_buf = (1 << skip_bits) - 1; + bit_left = BUF_BITS - skip_bits; +} + +void put_vc2_ue_uint(uint val) { + int pbits = 0, topbit = 1, maxval = 1, bits = 0; + if (val == 0) { + put_bits(1, 1); + return; + } + val++; + + while (val > maxval) { + topbit <<= 1; + bits++; + maxval <<= 1; + maxval |= 1; + } + + for (int i = 0; i < bits; i++) { + topbit >>= 1; + pbits <<= 2; + if ((val & topbit) != 0) { + pbits |= 1; + } + } + + put_bits(bits * 2 + 1, (pbits << 1) | 1); +} + +int align(int x, int a) { + return (x + a - 1) & ~(a - 1); +} + +int quants[MAX_DWT_LEVELS][4]; + +int subband_coord(int index, int h, int lvl) { + int coord = index; + coord <<= 1; + coord |= h; + coord <<= (wavelet_depth-lvl-1); + return coord; +} + +void main() { + int slice_index = int(gl_GlobalInvocationID.x); + int max_index = num_slices.x * num_slices.y; + if (slice_index >= max_index) { + return; + } + + /* Step 2. Quantize and encode */ + int pb_start = slice.args[slice_index].pb_start; + for (int i = 0, index = WORKGROUP_X - 1; i < gl_WorkGroupID.x; i++) { + pb_start += slice.args[index].pb_start + slice.args[index].bytes; + index += WORKGROUP_X; + } + ivec2 slice_coord = ivec2(slice_index % num_slices.x, slice_index / num_slices.x); + write_ptr = (pb_start >> 2); + write_start = write_ptr; + + int slice_bytes_max = slice.args[slice_index].bytes; + int quant_index = slice.args[slice_index].quant_idx; + + for (int level = 0; level < wavelet_depth; level++) + for (int orientation = int(level > 0); orientation < 4; orientation++) + quants[level][orientation] = max(quant_index - luts.quant[level][orientation], 0); + + /* Write quant index for this slice */ + put_bits(8, quant_index); + + /* Luma + 2 Chroma planes */ + for (int p = 0; p < 3; p++) { + int pad_s, pad_c; + int bytes_start = put_bytes_count(); + + /* Save current location and write a zero value */ + int write_ptr_start = write_ptr; + int bit_left_start = bit_left; + put_bits(8, 0); + for (int level = 0; level < wavelet_depth; level++) { + ivec2 band_size = planes[p].dwt_dim >> (wavelet_depth - level); + for (int o = int(level > 0); o < 4; o++) { + /* Encode subband */ + int left = band_size.x * (slice_coord.x) / num_slices.x; + int right = band_size.x * (slice_coord.x+1) / num_slices.x; + int top = band_size.y * (slice_coord.y) / num_slices.y; + int bottom = band_size.y * (slice_coord.y+1) / num_slices.y; + + const int q_idx = quants[level][o]; + const int qfactor = luts.ff_dirac_qscale_tab[q_idx]; + + const int yh = o >> 1; + const int xh = o & 1; + + int stride = align(planes[p].dwt_dim.x, 32); + for (int y = top; y < bottom; y++) { + for (int x = left; x < right; x++) { + int sx = subband_coord(x, xh, level); + int sy = subband_coord(y, yh, level); + int coef = plane_dat[p].coef_buf[sy * stride + sx]; + uint c_abs = uint(abs(coef)); + c_abs = (c_abs << 2) / qfactor; + put_vc2_ue_uint(c_abs); + if (c_abs != 0) { + put_bits(1, int(coef < 0)); + } + } + } + } + } + flush_put_bits(); + int bytes_len = put_bytes_count() - bytes_start - 1; + if (p == 2) { + int len_diff = slice_bytes_max - put_bytes_count(); + pad_s = align((bytes_len + len_diff), size_scaler)/size_scaler; + pad_c = (pad_s*size_scaler) - bytes_len; + } else { + pad_s = align(bytes_len, size_scaler)/size_scaler; + pad_c = (pad_s*size_scaler) - bytes_len; + } + BitBufByte pbb = BitBufByte(pb); + int start_ptr = (write_ptr_start << 2) + ((BUF_BITS - bit_left_start) >> 3); + pbb.data[start_ptr] = uint8_t(pad_s); + /* vc2-reference uses that padding that decodes to '0' coeffs */ + skip_put_bytes(pad_c); + } +} \ No newline at end of file diff --git a/libavcodec/vulkan/slice_sizes.comp b/libavcodec/vulkan/slice_sizes.comp new file mode 100644 index 0000000000..a965e911c7 --- /dev/null +++ b/libavcodec/vulkan/slice_sizes.comp @@ -0,0 +1,184 @@ +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_scalar_block_layout : require +#extension GL_EXT_buffer_reference : require + +#define WORKGROUP_X 128 +layout(local_size_x = WORKGROUP_X, local_size_y = 1, local_size_z = 1) in; + +#define DIRAC_MAX_QUANT_INDEX 116 +#define MAX_DWT_LEVELS 5 + +struct SliceArgs { + int quant_idx; + int bytes; + int pb_start; + int pad; +}; + +struct Plane { + ivec2 dim; + ivec2 dwt_dim; +}; + +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer DwtCoef { + int coef_buf[]; +}; + +layout(std430, buffer_reference) buffer SliceArgBuf { + SliceArgs args[]; +}; + +layout(scalar, buffer_reference, buffer_reference_align = 4) buffer QuantLuts { + int quant[5][4]; + int ff_dirac_qscale_tab[116]; +}; + +layout(push_constant, scalar) uniform ComputeInfo { + DwtCoef plane_dat[3]; + QuantLuts luts; + SliceArgBuf slice; + ivec2 num_slices; + Plane planes[3]; + int wavelet_depth; + int size_scaler; + int prefix_bytes; + int bits_ceil; + int bits_floor; +}; + +int count_vc2_ue_uint(uint val) { + uint topbit = 1, maxval = 1; + int bits = 0; + if (val == 0) { + return 1; + } + val++; + while (val > maxval) { + bits++; + topbit <<= 1; + maxval <<= 1; + maxval |= 1; + } + return bits * 2 + 1; +} + +int ffalign(int x, int a) { + return (x + a - 1) & ~(a - 1); +} + +int cache[DIRAC_MAX_QUANT_INDEX]; +int quants[MAX_DWT_LEVELS][4]; +shared int slice_sizes[WORKGROUP_X]; + +int subband_coord(int index, int h, int lvl) { + int coord = index; + coord <<= 1; + coord |= h; + coord <<= (wavelet_depth-lvl-1); + return coord; +} + +int count_hq_slice(int quant_index) { + int bits = 0; + if (cache[quant_index] != 0) { + return cache[quant_index]; + } + + bits += 8*prefix_bytes; + bits += 8; /* quant_idx */ + + for (int level = 0; level < wavelet_depth; level++) + for (int orientation = int(level > 0); orientation < 4; orientation++) + quants[level][orientation] = max(quant_index - luts.quant[level][orientation], 0); + + int slice_index = int(gl_GlobalInvocationID.x); + ivec2 slice_coord = ivec2(slice_index % num_slices.x, slice_index / num_slices.x); + for (int p = 0; p < 3; p++) { + int bytes_start = bits >> 3; + bits += 8; + for (int level = 0; level < wavelet_depth; level++) { + ivec2 band_dim = planes[p].dwt_dim >> (wavelet_depth - level); + for (int o = int(level > 0); o < 4; o++) { + const int left = band_dim.x * slice_coord.x / num_slices.x; + const int right = band_dim.x * (slice_coord.x+1) / num_slices.x; + const int top = band_dim.y * slice_coord.y / num_slices.y; + const int bottom = band_dim.y * (slice_coord.y+1) / num_slices.y; + + const int q_idx = quants[level][o]; + const int qfactor = luts.ff_dirac_qscale_tab[q_idx]; + + const int yh = o >> 1; + const int xh = o & 1; + + const int stride = ffalign(planes[p].dwt_dim.x, 32); + for (int y = top; y < bottom; y++) { + for (int x = left; x < right; x++) { + int sx = subband_coord(x, xh, level); + int sy = subband_coord(y, yh, level); + int coef = plane_dat[p].coef_buf[sy * stride + sx]; + uint c_abs = uint(abs(coef)); + c_abs = (c_abs << 2) / qfactor; + bits += count_vc2_ue_uint(c_abs); + bits += int(c_abs > 0); + } + } + } + } + bits += ffalign(bits, 8) - bits; + int bytes_len = (bits >> 3) - bytes_start - 1; + int pad_s = ffalign(bytes_len, size_scaler) / size_scaler; + int pad_c = (pad_s * size_scaler) - bytes_len; + bits += pad_c * 8; + } + + cache[quant_index] = bits; + return bits; +} + +int ssize_round(int b) { + return ffalign(b, size_scaler) + 4 + prefix_bytes; +} + +void main() { + int slice_index = int(gl_GlobalInvocationID.x); + int max_index = num_slices.x * num_slices.y; + if (slice_index >= max_index) { + return; + } + for (int i = 0; i < DIRAC_MAX_QUANT_INDEX; i++) { + cache[i] = 0; + } + const int q_ceil = DIRAC_MAX_QUANT_INDEX; + const int top = bits_ceil; + const int bottom = bits_floor; + int quant_buf[2] = int[2](-1, -1); + int quant = slice.args[slice_index].quant_idx; + int step = 1; + int bits_last = 0; + int bits = count_hq_slice(quant); + while ((bits > top) || (bits < bottom)) { + const int signed_step = bits > top ? +step : -step; + quant = clamp(quant + signed_step, 0, q_ceil-1); + bits = count_hq_slice(quant); + if (quant_buf[1] == quant) { + quant = max(quant_buf[0], quant); + bits = quant == quant_buf[0] ? bits_last : bits; + break; + } + step = clamp(step / 2, 1, (q_ceil - 1) / 2); + quant_buf[1] = quant_buf[0]; + quant_buf[0] = quant; + bits_last = bits; + } + int bytes = ssize_round(bits >> 3); + slice.args[slice_index].quant_idx = clamp(quant, 0, q_ceil-1); + slice.args[slice_index].bytes = bytes; + slice_sizes[gl_LocalInvocationIndex] = bytes; + barrier(); + /* Prefix sum for all slices in current workgroup */ + int total_bytes = 0; + for (int i = 0; i < gl_LocalInvocationIndex; i++) { + total_bytes += slice_sizes[i]; + } + slice.args[slice_index].pb_start = total_bytes; +} \ No newline at end of file