git.sesse.net Git - nageru/blob - nageru/av1_encoder.cpp

   1 #include "av1_encoder.h"
   2
   3 #include <assert.h>
   4 #include <dlfcn.h>
   5 #include <stdio.h>
   6 #include <stdlib.h>
   7 #include <string.h>
   8 #include <unistd.h>
   9 #include <atomic>
  10 #include <cstdint>
  11 #include <functional>
  12 #include <mutex>
  13
  14 #include <EbSvtAv1.h>
  15 #include <EbSvtAv1Enc.h>
  16
  17 #include "defs.h"
  18 #include "flags.h"
  19 #include "shared/metrics.h"
  20 #include "shared/mux.h"
  21 #include "print_latency.h"
  22 #include "shared/timebase.h"
  23 #include "shared/memcpy_interleaved.h"
  24
  25 extern "C" {
  26 #include <libavcodec/avcodec.h>
  27 #include <libavformat/avformat.h>
  28 }
  29
  30 using namespace movit;
  31 using namespace std;
  32 using namespace std::chrono;
  33 using namespace std::placeholders;
  34
  35 namespace {
  36
  37 // AV1Encoder can be restarted if --record-av1-video is set, so make these
  38 // metrics global.
  39 atomic<int64_t> metric_av1_queued_frames{0};
  40 atomic<int64_t> metric_av1_max_queued_frames{AV1_QUEUE_LENGTH};
  41 atomic<int64_t> metric_av1_dropped_frames{0};
  42 atomic<int64_t> metric_av1_output_frames_i{0};
  43 atomic<int64_t> metric_av1_output_frames_p{0};
  44 Histogram metric_av1_qp;
  45 LatencyHistogram av1_latency_histogram;
  46
  47 once_flag av1_metrics_inited;
  48
  49 }  // namespace
  50
  51 AV1Encoder::AV1Encoder(const AVOutputFormat *oformat)
  52         : wants_global_headers(oformat->flags & AVFMT_GLOBALHEADER)
  53 {
  54                 call_once(av1_metrics_inited, []{
  55                         global_metrics.add("av1_queued_frames",  {}, &metric_av1_queued_frames, Metrics::TYPE_GAUGE);
  56                         global_metrics.add("av1_max_queued_frames", {},  &metric_av1_max_queued_frames, Metrics::TYPE_GAUGE);
  57                         global_metrics.add("av1_dropped_frames", {},  &metric_av1_dropped_frames);
  58                         global_metrics.add("av1_output_frames", {{ "type", "i" }}, &metric_av1_output_frames_i);
  59                         global_metrics.add("av1_output_frames", {{ "type", "p" }}, &metric_av1_output_frames_p);
  60
  61                         metric_av1_qp.init_uniform(50);
  62                         global_metrics.add("av1_qp", {}, &metric_av1_qp);
  63                         av1_latency_histogram.init("av1");
  64                 });
  65
  66         const size_t bytes_per_pixel = global_flags.bit_depth > 8 ? 2 : 1;
  67         frame_pool.reset(new uint8_t[global_flags.width * global_flags.height * 2 * bytes_per_pixel * AV1_QUEUE_LENGTH]);
  68         for (unsigned i = 0; i < AV1_QUEUE_LENGTH; ++i) {
  69                 free_frames.push(frame_pool.get() + i * (global_flags.width * global_flags.height * 2 * bytes_per_pixel));
  70         }
  71         encoder_thread = thread(&AV1Encoder::encoder_thread_func, this);
  72 }
  73
  74 AV1Encoder::~AV1Encoder()
  75 {
  76         should_quit = true;
  77         queued_frames_nonempty.notify_all();
  78         encoder_thread.join();
  79 }
  80
  81 void AV1Encoder::add_frame(int64_t pts, int64_t duration, YCbCrLumaCoefficients ycbcr_coefficients, const uint8_t *data, const ReceivedTimestamps &received_ts)
  82 {
  83         assert(!should_quit);
  84
  85         QueuedFrame qf;
  86         qf.pts = pts;
  87         qf.duration = duration;
  88         qf.ycbcr_coefficients = ycbcr_coefficients;
  89         qf.received_ts = received_ts;
  90
  91         {
  92                 lock_guard<mutex> lock(mu);
  93                 if (free_frames.empty()) {
  94                         fprintf(stderr, "WARNING: AV1 queue full, dropping frame with pts %" PRId64 "\n", pts);
  95                         ++metric_av1_dropped_frames;
  96                         return;
  97                 }
  98
  99                 qf.data = free_frames.front();
 100                 free_frames.pop();
 101         }
 102
 103         // Since we're copying anyway, we can unpack from NV12 to fully planar on the fly.
 104         // SVT-AV1 makes its own copy, though, and it would have been nice to avoid the
 105         // double-copy (and also perhaps let the GPU do the 10-bit compression SVT-AV1
 106         // wants, instead of doing it on the CPU).
 107         const size_t bytes_per_pixel = global_flags.bit_depth > 8 ? 2 : 1;
 108         size_t frame_size = global_flags.width * global_flags.height * bytes_per_pixel;
 109         assert(global_flags.width % 2 == 0);
 110         assert(global_flags.height % 2 == 0);
 111         uint8_t *y = qf.data;
 112         uint8_t *cb = y + frame_size;
 113         uint8_t *cr = cb + frame_size / 4;
 114         memcpy(y, data, frame_size);
 115         if (global_flags.bit_depth == 8) {
 116                 memcpy_interleaved(cb, cr, data + frame_size, frame_size / 2);
 117         } else {
 118                 const uint16_t *src = reinterpret_cast<const uint16_t *>(data + frame_size);
 119                 uint16_t *cb16 = reinterpret_cast<uint16_t *>(cb);
 120                 uint16_t *cr16 = reinterpret_cast<uint16_t *>(cr);
 121                 memcpy_interleaved_word(cb16, cr16, src, frame_size / 4);
 122         }
 123
 124         {
 125                 lock_guard<mutex> lock(mu);
 126                 queued_frames.push(qf);
 127                 queued_frames_nonempty.notify_all();
 128                 metric_av1_queued_frames = queued_frames.size();
 129         }
 130 }
 131
 132 void AV1Encoder::init_av1()
 133 {
 134         EbSvtAv1EncConfiguration config;
 135
 136         // svt_av1_enc_init_handle() is defined to fill config with the defaults;
 137         // yet, seemingly, not everything is written, and some of it can cause
 138         // Valgrind warnings and/or crashes. It should never hurt to put it
 139         // into a known state beforehand, and it seems to fix the crashes,
 140         // so we do that.
 141         memset(&config, 0, sizeof(config));
 142
 143         EbErrorType ret = svt_av1_enc_init_handle(&encoder, nullptr, &config);
 144         if (ret != EB_ErrorNone) {
 145                 fprintf(stderr, "Error initializing SVT-AV1 handle (error %08x)\n", ret);
 146                 exit(EXIT_FAILURE);
 147         }
 148
 149         // NOTE: We don't set CBR, as it requires low-delay mode, which is
 150         // generally problematic wrt. quality and performance.
 151         config.enc_mode = global_flags.av1_preset;
 152         config.intra_period_length = 63;  // Approx. one second, conforms to the (n % 8) - 1 == 0 rule.
 153         config.source_width = global_flags.width;
 154         config.source_height = global_flags.height;
 155         config.frame_rate_numerator = global_flags.av1_fps_num;
 156         config.frame_rate_denominator = global_flags.av1_fps_den;
 157         config.encoder_bit_depth = global_flags.bit_depth;
 158         config.rate_control_mode = 1;  // VBR.
 159         config.target_bit_rate = global_flags.av1_bitrate * 1000;
 160
 161         // NOTE: These should be in sync with the ones in quicksync_encoder.cpp (sps_rbsp()).
 162         config.color_primaries = EB_CICP_CP_BT_709;
 163         config.transfer_characteristics = EB_CICP_TC_SRGB;
 164         if (global_flags.ycbcr_rec709_coefficients) {
 165                 config.matrix_coefficients = EB_CICP_MC_BT_709;
 166         } else {
 167                 config.matrix_coefficients = EB_CICP_MC_BT_601;
 168         }
 169         config.color_range = EB_CR_STUDIO_RANGE;
 170         config.chroma_sample_position = EB_CSP_VERTICAL;
 171
 172         const vector<string> &extra_param = global_flags.av1_extra_param;
 173         for (const string &str : extra_param) {
 174                 const size_t pos = str.find(',');
 175                 if (pos == string::npos) {
 176                         if (svt_av1_enc_parse_parameter(&config, str.c_str(), nullptr) != EB_ErrorNone) {
 177                                 fprintf(stderr, "ERROR: SVT-AV1 rejected parameter '%s' with no value\n", str.c_str());
 178                                 exit(EXIT_FAILURE);
 179                         }
 180                 } else {
 181                         const string key = str.substr(0, pos);
 182                         const string value = str.substr(pos + 1);
 183                         if (svt_av1_enc_parse_parameter(&config, key.c_str(), value.c_str()) != EB_ErrorNone) {
 184                                 fprintf(stderr, "ERROR: SVT-AV1 rejected parameter '%s' set to '%s'\n",
 185                                         key.c_str(), value.c_str());
 186                                 exit(EXIT_FAILURE);
 187                         }
 188                 }
 189         }
 190
 191         ret = svt_av1_enc_set_parameter(encoder, &config);
 192         if (ret != EB_ErrorNone) {
 193                 fprintf(stderr, "Error configuring SVT-AV1 (error %08x)\n", ret);
 194                 exit(EXIT_FAILURE);
 195         }
 196
 197         ret = svt_av1_enc_init(encoder);
 198         if (ret != EB_ErrorNone) {
 199                 fprintf(stderr, "Error initializing SVT-AV1 (error %08x)\n", ret);
 200                 exit(EXIT_FAILURE);
 201         }
 202
 203         if (wants_global_headers) {
 204                 EbBufferHeaderType *header = NULL;
 205
 206                 ret = svt_av1_enc_stream_header(encoder, &header);
 207                 if (ret != EB_ErrorNone) {
 208                         fprintf(stderr, "Error building SVT-AV1 header (error %08x)\n", ret);
 209                         exit(EXIT_FAILURE);
 210                 }
 211
 212                 global_headers = string(reinterpret_cast<const char *>(header->p_buffer), header->n_filled_len);
 213
 214                 svt_av1_enc_stream_header_release(header);  // Don't care about errors.
 215           }
 216 }
 217
 218 void AV1Encoder::encoder_thread_func()
 219 {
 220         if (nice(5) == -1) {
 221                 perror("nice()");
 222                 // No exit; it's not fatal.
 223         }
 224         pthread_setname_np(pthread_self(), "AV1_encode");
 225         init_av1();
 226         av1_init_done = true;
 227
 228         bool frames_left;
 229
 230         do {
 231                 QueuedFrame qf;
 232
 233                 // Wait for a queued frame, then dequeue it.
 234                 {
 235                         unique_lock<mutex> lock(mu);
 236                         queued_frames_nonempty.wait(lock, [this]() { return !queued_frames.empty() || should_quit; });
 237                         if (!queued_frames.empty()) {
 238                                 qf = queued_frames.front();
 239                                 queued_frames.pop();
 240                         } else {
 241                                 qf.pts = -1;
 242                                 qf.duration = -1;
 243                                 qf.data = nullptr;
 244                         }
 245
 246                         metric_av1_queued_frames = queued_frames.size();
 247                         frames_left = !queued_frames.empty();
 248                 }
 249
 250                 encode_frame(qf);
 251
 252                 {
 253                         lock_guard<mutex> lock(mu);
 254                         free_frames.push(qf.data);
 255                 }
 256
 257                 // We should quit only if the should_quit flag is set _and_ we have nothing
 258                 // in our queue.
 259         } while (!should_quit || frames_left);
 260
 261         // Signal end of stream.
 262         EbBufferHeaderType hdr;
 263         hdr.n_alloc_len   = 0;
 264         hdr.n_filled_len  = 0;
 265         hdr.n_tick_count  = 0;
 266         hdr.p_app_private = nullptr;
 267         hdr.pic_type      = EB_AV1_INVALID_PICTURE;
 268         hdr.p_buffer      = nullptr;
 269         hdr.metadata      = nullptr;
 270         hdr.flags         = EB_BUFFERFLAG_EOS;
 271         svt_av1_enc_send_picture(encoder, &hdr);
 272
 273         bool seen_eof = false;
 274         do {
 275                 EbBufferHeaderType *buf;
 276                 EbErrorType ret = svt_av1_enc_get_packet(encoder, &buf, /*pic_send_done=*/true);
 277                 if (ret == EB_NoErrorEmptyQueue) {
 278                         assert(false);
 279                 }
 280                 seen_eof = (buf->flags & EB_BUFFERFLAG_EOS);
 281                 process_packet(buf);
 282         } while (!seen_eof);
 283
 284         svt_av1_enc_deinit(encoder);
 285         svt_av1_enc_deinit_handle(encoder);
 286 }
 287
 288 void AV1Encoder::encode_frame(AV1Encoder::QueuedFrame qf)
 289 {
 290         if (qf.data) {
 291                 const size_t bytes_per_pixel = global_flags.bit_depth > 8 ? 2 : 1;
 292
 293                 EbSvtIOFormat pic;
 294                 pic.luma = qf.data;
 295                 pic.cb = pic.luma + global_flags.width * global_flags.height * bytes_per_pixel;
 296                 pic.cr = pic.cb + (global_flags.width * global_flags.height / 4) * bytes_per_pixel;
 297                 pic.y_stride = global_flags.width;  // In pixels, so no bytes_per_pixel.
 298                 pic.cb_stride = global_flags.width / 2;  // Likewise.
 299                 pic.cr_stride = global_flags.width / 2;  // Likewise.
 300                 pic.width = global_flags.width;
 301                 pic.height = global_flags.height;
 302                 pic.org_x = 0;
 303                 pic.org_y = 0;
 304                 pic.color_fmt = EB_YUV420;
 305                 pic.bit_depth = global_flags.bit_depth > 8 ? EB_TEN_BIT : EB_EIGHT_BIT;
 306
 307                 EbBufferHeaderType hdr;
 308                 hdr.p_buffer      = reinterpret_cast<uint8_t *>(&pic);
 309                 hdr.n_alloc_len   = (global_flags.width * global_flags.height * 3 / 2) * bytes_per_pixel;
 310                 hdr.n_filled_len  = hdr.n_alloc_len;
 311                 hdr.n_tick_count  = 0;
 312                 hdr.p_app_private = nullptr;
 313                 hdr.pic_type      = EB_AV1_INVALID_PICTURE;  // Actually means auto, according to FFmpeg.
 314                 hdr.metadata      = nullptr;
 315                 hdr.flags         = 0;
 316                 hdr.pts           = av_rescale_q(qf.pts, AVRational{ 1, TIMEBASE }, AVRational{ global_flags.av1_fps_den, global_flags.av1_fps_num });
 317                 if (hdr.pts <= last_pts) {
 318                         fprintf(stderr, "WARNING: Receiving frames faster than given --av1-fps value (%d/%d); dropping frame.\n",
 319                                 global_flags.av1_fps_num, global_flags.av1_fps_den);
 320                 } else {
 321                         svt_av1_enc_send_picture(encoder, &hdr);
 322                         frames_being_encoded[hdr.pts] = qf.received_ts;
 323                         last_pts = hdr.pts;
 324                 }
 325         }
 326
 327         for ( ;; ) {
 328                 EbBufferHeaderType *buf;
 329                 EbErrorType ret = svt_av1_enc_get_packet(encoder, &buf, /*pic_send_done=*/false);
 330                 if (ret == EB_NoErrorEmptyQueue) {
 331                         return;
 332                 }
 333                 process_packet(buf);
 334         }
 335 }
 336
 337 void AV1Encoder::process_packet(EbBufferHeaderType *buf)
 338 {
 339         if (buf->n_filled_len == 0) {
 340                 // TODO: Can this ever happen?
 341                 svt_av1_enc_release_out_buffer(&buf);
 342                 return;
 343         }
 344
 345         switch (buf->pic_type) {
 346                 case EB_AV1_KEY_PICTURE:
 347                 case EB_AV1_INTRA_ONLY_PICTURE:
 348                         ++metric_av1_output_frames_i;
 349                         break;
 350                 case EB_AV1_INTER_PICTURE:  // We don't really know whether it's P or B.
 351                         ++metric_av1_output_frames_p;
 352                         break;
 353                 default:
 354                         break;
 355         }
 356         metric_av1_qp.count_event(buf->qp);
 357
 358         if (frames_being_encoded.count(buf->pts)) {
 359                 ReceivedTimestamps received_ts = frames_being_encoded[buf->pts];
 360                 frames_being_encoded.erase(buf->pts);
 361
 362                 static int frameno = 0;
 363                 print_latency("Current AV1 latency (video inputs → network mux):",
 364                                 received_ts, /*b_frame=*/false, &frameno, &av1_latency_histogram);
 365         } else {
 366                 assert(false);
 367         }
 368
 369         AVPacket pkt;
 370         memset(&pkt, 0, sizeof(pkt));
 371         pkt.buf = nullptr;
 372         pkt.data = buf->p_buffer;
 373         pkt.size = buf->n_filled_len;
 374         pkt.stream_index = 0;
 375         if (buf->pic_type == EB_AV1_KEY_PICTURE) {
 376                 pkt.flags = AV_PKT_FLAG_KEY;
 377         } else if (buf->pic_type == EB_AV1_NON_REF_PICTURE) {
 378                 // I have no idea if this does anything in practice,
 379                 // but the libavcodec plugin does it.
 380                 pkt.flags = AV_PKT_FLAG_DISPOSABLE;
 381         } else {
 382                 pkt.flags = 0;
 383         }
 384         pkt.pts = av_rescale_q(buf->pts, AVRational{ global_flags.av1_fps_den, global_flags.av1_fps_num }, AVRational{ 1, TIMEBASE });
 385         pkt.dts = av_rescale_q(buf->dts, AVRational{ global_flags.av1_fps_den, global_flags.av1_fps_num }, AVRational{ 1, TIMEBASE });
 386
 387         for (Mux *mux : muxes) {
 388                 mux->add_packet(pkt, pkt.pts, pkt.dts);
 389         }
 390
 391         svt_av1_enc_release_out_buffer(&buf);
 392 }