replay: shared decoder context (#32255)

share decoder context

cleanup includes

tools/replay/framereader.cc

@@ -1,5 +1,10 @@
 #include "tools/replay/framereader.h"
 
+#include <map>
+#include <memory>
+#include <tuple>
+#include <utility>
+
 #include "common/util.h"
 #include "third_party/libyuv/include/libyuv.h"
 #include "tools/replay/util.h"
@@ -12,7 +17,6 @@
 #define HW_PIX_FMT AV_PIX_FMT_CUDA
 #endif
 
-
 namespace {
 
 enum AVPixelFormat get_hw_format(AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts) {
@@ -21,11 +25,32 @@ enum AVPixelFormat get_hw_format(AVCodecContext *ctx, const enum AVPixelFormat *
     if (*p == *hw_pix_fmt) return *p;
   }
   rWarning("Please run replay with the --no-hw-decoder flag!");
-  // fallback to YUV420p
   *hw_pix_fmt = AV_PIX_FMT_NONE;
   return AV_PIX_FMT_YUV420P;
 }
 
+struct DecoderManager {
+  VideoDecoder *acquire(CameraType type, AVCodecParameters *codecpar, bool hw_decoder) {
+    auto key = std::tuple(type, codecpar->width, codecpar->height);
+    std::unique_lock lock(mutex_);
+    if (auto it = decoders_.find(key); it != decoders_.end()) {
+      return it->second.get();
+    }
+
+    auto decoder = std::make_unique<VideoDecoder>();
+    if (!decoder->open(codecpar, hw_decoder)) {
+      decoder.reset(nullptr);
+    }
+    decoders_[key] = std::move(decoder);
+    return decoders_[key].get();
+  }
+
+  std::mutex mutex_;
+  std::map<std::tuple<CameraType, int, int>, std::unique_ptr<VideoDecoder>> decoders_;
+};
+
+DecoderManager decoder_manager;
+
 }  // namespace
 
 FrameReader::FrameReader() {
@@ -33,12 +58,10 @@ FrameReader::FrameReader() {
 }
 
 FrameReader::~FrameReader() {
-  if (decoder_ctx) avcodec_free_context(&decoder_ctx);
   if (input_ctx) avformat_close_input(&input_ctx);
-  if (hw_device_ctx) av_buffer_unref(&hw_device_ctx);
 }
 
-bool FrameReader::load(const std::string &url, bool no_hw_decoder, std::atomic<bool> *abort, bool local_cache, int chunk_size, int retries) {
+bool FrameReader::load(CameraType type, const std::string &url, bool no_hw_decoder, std::atomic<bool> *abort, bool local_cache, int chunk_size, int retries) {
   auto local_file_path = url.find("https://") == 0 ? cacheFilePath(url) : url;
   if (!util::file_exists(local_file_path)) {
     FileReader f(local_cache, chunk_size, retries);
@@ -46,10 +69,10 @@ bool FrameReader::load(const std::string &url, bool no_hw_decoder, std::atomic<b
       return false;
     }
   }
-  return loadFromFile(local_file_path, no_hw_decoder, abort);
+  return loadFromFile(type, local_file_path, no_hw_decoder, abort);
 }
 
-bool FrameReader::loadFromFile(const std::string &file, bool no_hw_decoder, std::atomic<bool> *abort) {
+bool FrameReader::loadFromFile(CameraType type, const std::string &file, bool no_hw_decoder, std::atomic<bool> *abort) {
   if (avformat_open_input(&input_ctx, file.c_str(), nullptr, nullptr) != 0 ||
       avformat_find_stream_info(input_ctx, nullptr) < 0) {
     rError("Failed to open input file or find video stream");
@@ -57,20 +80,57 @@ bool FrameReader::loadFromFile(const std::string &file, bool no_hw_decoder, std:
   }
   input_ctx->probesize = 10 * 1024 * 1024;  // 10MB
 
-  AVStream *video = input_ctx->streams[0];
+  decoder_ = decoder_manager.acquire(type, input_ctx->streams[0]->codecpar, !no_hw_decoder);
-  const AVCodec *decoder = avcodec_find_decoder(video->codecpar->codec_id);
+  if (!decoder_) {
+    return false;
+  }
+  width = decoder_->width;
+  height = decoder_->height;
+
+  AVPacket pkt;
+  packets_info.reserve(60 * 20);  // 20fps, one minute
+  while (!(abort && *abort) && av_read_frame(input_ctx, &pkt) == 0) {
+    packets_info.emplace_back(PacketInfo{.flags = pkt.flags, .pos = pkt.pos});
+    av_packet_unref(&pkt);
+  }
+  avio_seek(input_ctx->pb, 0, SEEK_SET);
+  return !packets_info.empty();
+}
+
+bool FrameReader::get(int idx, VisionBuf *buf) {
+  if (!buf || idx < 0 || idx >= packets_info.size()) {
+    return false;
+  }
+  return decoder_->decode(this, idx, buf);
+}
+
+// class VideoDecoder
+
+VideoDecoder::VideoDecoder() {
+  av_frame_ = av_frame_alloc();
+  hw_frame_ = av_frame_alloc();
+}
+
+VideoDecoder::~VideoDecoder() {
+  if (hw_device_ctx) av_buffer_unref(&hw_device_ctx);
+  if (decoder_ctx) avcodec_free_context(&decoder_ctx);
+  av_frame_free(&av_frame_);
+  av_frame_free(&hw_frame_);
+}
+
+bool VideoDecoder::open(AVCodecParameters *codecpar, bool hw_decoder) {
+  const AVCodec *decoder = avcodec_find_decoder(codecpar->codec_id);
   if (!decoder) return false;
 
   decoder_ctx = avcodec_alloc_context3(decoder);
-  if (!decoder_ctx || avcodec_parameters_to_context(decoder_ctx, video->codecpar) != 0) {
+  if (!decoder_ctx || avcodec_parameters_to_context(decoder_ctx, codecpar) != 0) {
     rError("Failed to allocate or initialize codec context");
     return false;
   }
-
   width = (decoder_ctx->width + 3) & ~3;
   height = decoder_ctx->height;
 
-  if (has_hw_decoder && !no_hw_decoder && !initHardwareDecoder(HW_DEVICE_TYPE)) {
+  if (hw_decoder && !initHardwareDecoder(HW_DEVICE_TYPE)) {
     rWarning("No device with hardware decoder found. fallback to CPU decoding.");
   }
 
@@ -78,36 +138,25 @@ bool FrameReader::loadFromFile(const std::string &file, bool no_hw_decoder, std:
     rError("Failed to open codec");
     return false;
   }
-
+  return true;
-  AVPacket pkt;
-  packets_info.reserve(60 * 20);  // 20fps, one minute
-  while (!(abort && *abort) && av_read_frame(input_ctx, &pkt) == 0) {
-    packets_info.emplace_back(PacketInfo{.flags = pkt.flags, .pos = pkt.pos});
-    av_packet_unref(&pkt);
-  }
-
-  avio_seek(input_ctx->pb, 0, SEEK_SET);
-  return !packets_info.empty();
 }
 
-bool FrameReader::initHardwareDecoder(AVHWDeviceType hw_device_type) {
+bool VideoDecoder::initHardwareDecoder(AVHWDeviceType hw_device_type) {
-  for (int i = 0;; i++) {
+  const AVCodecHWConfig *config = nullptr;
-    const AVCodecHWConfig *config = avcodec_get_hw_config(decoder_ctx->codec, i);
+  for (int i = 0; (config = avcodec_get_hw_config(decoder_ctx->codec, i)) != nullptr; i++) {
-    if (!config) {
-      rWarning("decoder %s does not support hw device type %s.", decoder_ctx->codec->name,
-               av_hwdevice_get_type_name(hw_device_type));
-      return false;
-    }
     if (config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX && config->device_type == hw_device_type) {
       hw_pix_fmt = config->pix_fmt;
       break;
     }
   }
+  if (!config) {
+    rWarning("Hardware configuration not found");
+    return false;
+  }
 
   int ret = av_hwdevice_ctx_create(&hw_device_ctx, hw_device_type, nullptr, nullptr, 0);
   if (ret < 0) {
     hw_pix_fmt = AV_PIX_FMT_NONE;
-    has_hw_decoder = false;
     rWarning("Failed to create specified HW device %d.", ret);
     return false;
   }
@@ -118,34 +167,27 @@ bool FrameReader::initHardwareDecoder(AVHWDeviceType hw_device_type) {
   return true;
 }
 
-bool FrameReader::get(int idx, VisionBuf *buf) {
+bool VideoDecoder::decode(FrameReader *reader, int idx, VisionBuf *buf) {
-  if (!buf || idx < 0 || idx >= packets_info.size()) {
-    return false;
-  }
-  return decode(idx, buf);
-}
-
-bool FrameReader::decode(int idx, VisionBuf *buf) {
   int from_idx = idx;
-  if (idx != prev_idx + 1) {
+  if (idx != reader->prev_idx + 1) {
     // seeking to the nearest key frame
     for (int i = idx; i >= 0; --i) {
-      if (packets_info[i].flags & AV_PKT_FLAG_KEY) {
+      if (reader->packets_info[i].flags & AV_PKT_FLAG_KEY) {
         from_idx = i;
         break;
       }
     }
-    avio_seek(input_ctx->pb, packets_info[from_idx].pos, SEEK_SET);
+    avio_seek(reader->input_ctx->pb, reader->packets_info[from_idx].pos, SEEK_SET);
   }
-  prev_idx = idx;
+  reader->prev_idx = idx;
 
   bool result = false;
   AVPacket pkt;
   for (int i = from_idx; i <= idx; ++i) {
-    if (av_read_frame(input_ctx, &pkt) == 0) {
+    if (av_read_frame(reader->input_ctx, &pkt) == 0) {
       AVFrame *f = decodeFrame(&pkt);
       if (f && i == idx) {
-        result = copyBuffers(f, buf);
+        result = copyBuffer(f, buf);
       }
       av_packet_unref(&pkt);
     }
@@ -153,33 +195,27 @@ bool FrameReader::decode(int idx, VisionBuf *buf) {
   return result;
 }
 
-AVFrame *FrameReader::decodeFrame(AVPacket *pkt) {
+AVFrame *VideoDecoder::decodeFrame(AVPacket *pkt) {
   int ret = avcodec_send_packet(decoder_ctx, pkt);
   if (ret < 0) {
     rError("Error sending a packet for decoding: %d", ret);
     return nullptr;
   }
 
-  av_frame_.reset(av_frame_alloc());
+  ret = avcodec_receive_frame(decoder_ctx, av_frame_);
-  ret = avcodec_receive_frame(decoder_ctx, av_frame_.get());
   if (ret != 0) {
     rError("avcodec_receive_frame error: %d", ret);
     return nullptr;
   }
 
-  if (av_frame_->format == hw_pix_fmt) {
+  if (av_frame_->format == hw_pix_fmt && av_hwframe_transfer_data(hw_frame_, av_frame_, 0) < 0) {
-    hw_frame.reset(av_frame_alloc());
+    rError("error transferring frame data from GPU to CPU");
-    if ((ret = av_hwframe_transfer_data(hw_frame.get(), av_frame_.get(), 0)) < 0) {
-      rError("error transferring the data from GPU to CPU");
     return nullptr;
   }
-    return hw_frame.get();
+  return (av_frame_->format == hw_pix_fmt) ? hw_frame_ : av_frame_;
-  } else {
-    return av_frame_.get();
-  }
 }
 
-bool FrameReader::copyBuffers(AVFrame *f, VisionBuf *buf) {
+bool VideoDecoder::copyBuffer(AVFrame *f, VisionBuf *buf) {
   if (hw_pix_fmt == HW_PIX_FMT) {
     for (int i = 0; i < height/2; i++) {
       memcpy(buf->y + (i*2 + 0)*buf->stride, f->data[0] + (i*2 + 0)*f->linesize[0], width);

tools/replay/framereader.h

@@ -1,10 +1,10 @@
 #pragma once
 
-#include <memory>
 #include <string>
 #include <vector>
 
 #include "cereal/visionipc/visionbuf.h"
+#include "system/camerad/cameras/camera_common.h"
 #include "tools/replay/filereader.h"
 
 extern "C" {
@@ -12,39 +12,46 @@ extern "C" {
 #include <libavformat/avformat.h>
 }
 
-struct AVFrameDeleter {
+class VideoDecoder;
-  void operator()(AVFrame* frame) const { av_frame_free(&frame); }
-};
 
 class FrameReader {
 public:
   FrameReader();
   ~FrameReader();
-  bool load(const std::string &url, bool no_hw_decoder = false, std::atomic<bool> *abort = nullptr, bool local_cache = false,
+  bool load(CameraType type, const std::string &url, bool no_hw_decoder = false, std::atomic<bool> *abort = nullptr, bool local_cache = false,
             int chunk_size = -1, int retries = 0);
-  bool loadFromFile(const std::string &file, bool no_hw_decoder = false, std::atomic<bool> *abort = nullptr);
+  bool loadFromFile(CameraType type, const std::string &file, bool no_hw_decoder = false, std::atomic<bool> *abort = nullptr);
   bool get(int idx, VisionBuf *buf);
   size_t getFrameCount() const { return packets_info.size(); }
 
   int width = 0, height = 0;
 
-private:
+  VideoDecoder *decoder_ = nullptr;
-  bool initHardwareDecoder(AVHWDeviceType hw_device_type);
-  bool decode(int idx, VisionBuf *buf);
-  AVFrame * decodeFrame(AVPacket *pkt);
-  bool copyBuffers(AVFrame *f, VisionBuf *buf);
-
-  std::unique_ptr<AVFrame, AVFrameDeleter>av_frame_, hw_frame;
   AVFormatContext *input_ctx = nullptr;
-  AVCodecContext *decoder_ctx = nullptr;
-
-  AVPixelFormat hw_pix_fmt = AV_PIX_FMT_NONE;
-  AVBufferRef *hw_device_ctx = nullptr;
   int prev_idx = -1;
   struct PacketInfo {
     int flags;
     int64_t pos;
   };
   std::vector<PacketInfo> packets_info;
-  inline static std::atomic<bool> has_hw_decoder = true;
+};
+
+
+class VideoDecoder {
+public:
+  VideoDecoder();
+  ~VideoDecoder();
+  bool open(AVCodecParameters *codecpar, bool hw_decoder);
+  bool decode(FrameReader *reader, int idx, VisionBuf *buf);
+  int width = 0, height = 0;
+
+private:
+  bool initHardwareDecoder(AVHWDeviceType hw_device_type);
+  AVFrame *decodeFrame(AVPacket *pkt);
+  bool copyBuffer(AVFrame *f, VisionBuf *buf);
+
+  AVFrame *av_frame_, *hw_frame_;
+  AVCodecContext *decoder_ctx = nullptr;
+  AVPixelFormat hw_pix_fmt = AV_PIX_FMT_NONE;
+  AVBufferRef *hw_device_ctx = nullptr;
 };

tools/replay/route.cc

@@ -160,7 +160,7 @@ void Segment::loadFile(int id, const std::string file) {
   bool success = false;
   if (id < MAX_CAMERAS) {
     frames[id] = std::make_unique<FrameReader>();
-    success = frames[id]->load(file, flags & REPLAY_FLAG_NO_HW_DECODER, &abort_, local_cache, 20 * 1024 * 1024, 3);
+    success = frames[id]->load((CameraType)id, file, flags & REPLAY_FLAG_NO_HW_DECODER, &abort_, local_cache, 20 * 1024 * 1024, 3);
   } else {
     log = std::make_unique<LogReader>(filters_);
     success = log->load(file, &abort_, local_cache, 0, 3);