Refac, simplify

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -144,7 +144,7 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device)
 
 BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
   // TODONVDEC P0: we probably need to free the frames that have been decoded by
-  // NVDEC but not yet "mapped" - i.e. those that are still in frameBuffer_?
+  // NVDEC but not yet "mapped" - i.e. those that are still in readyFrames_?
 
   if (decoder_) {
     NVDECCache::getCache(device_.index())
@@ -328,40 +328,34 @@ int BetaCudaDeviceInterface::frameReadyForDecoding(CUVIDPICPARAMS* picParams) {
 
   // Send frame to be decoded by NVDEC - non-blocking call.
   CUresult result = cuvidDecodePicture(*decoder_.get(), picParams);
-  if (result != CUDA_SUCCESS) {
-    return 0; // Yes, you're reading that right, 0 means error.
-  }
 
-  frameBuffer_.markAsBeingDecoded(/*slotId=*/picParams->CurrPicIdx);
-  return 1;
+  // Yes, you're reading that right, 0 means error, 1 means success
+  return (result == CUDA_SUCCESS);
 }
 
 int BetaCudaDeviceInterface::frameReadyInDisplayOrder(
     CUVIDPARSERDISPINFO* dispInfo) {
-  frameBuffer_.markSlotReadyAndSetInfo(
-      /*slotId=*/dispInfo->picture_index, dispInfo);
-  return 1;
+  readyFrames_.push(*dispInfo);
+  return 1; // success
 }
 
 // Moral equivalent of avcodec_receive_frame().
 int BetaCudaDeviceInterface::receiveFrame(UniqueAVFrame& avFrame) {
-  FrameBuffer::Slot* slot = frameBuffer_.findReadySlotWithLowestPts();
-  if (slot == nullptr) {
+  if (readyFrames_.empty()) {
     // No frame found, instruct caller to try again later after sending more
     // packets.
     return AVERROR(EAGAIN);
   }
+  CUVIDPARSERDISPINFO dispInfo = readyFrames_.front();
+  readyFrames_.pop();
 
   CUVIDPROCPARAMS procParams = {};
-  CUVIDPARSERDISPINFO dispInfo = slot->dispInfo;
   procParams.progressive_frame = dispInfo.progressive_frame;
   procParams.top_field_first = dispInfo.top_field_first;
   procParams.unpaired_field = dispInfo.repeat_first_field < 0;
   CUdeviceptr framePtr = 0;
   unsigned int pitch = 0;
 
-  frameBuffer_.free(slot->slotId);
-
   // We know the frame we want was sent to the hardware decoder, but now we need
   // to "map" it to an "output surface" before we can use its data. This is a
   // blocking calls that waits until the frame is fully decoded and ready to be
@@ -478,7 +472,8 @@ void BetaCudaDeviceInterface::flush() {
 
   isFlushing_ = false;
 
-  frameBuffer_.clear();
+  std::queue<CUVIDPARSERDISPINFO> emptyQueue;
+  std::swap(readyFrames_, emptyQueue);
 
   eofSent_ = false;
 }
@@ -512,57 +507,4 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
       preAllocatedOutputTensor);
 }
 
-void BetaCudaDeviceInterface::FrameBuffer::markAsBeingDecoded(int slotId) {
-  auto it = map_.find(slotId);
-  TORCH_CHECK(
-      it == map_.end(),
-      "Slot ",
-      slotId,
-      " is already occupied. This should never happen.");
-
-  map_.emplace(slotId, Slot(slotId, SlotState::BEING_DECODED));
-}
-
-void BetaCudaDeviceInterface::FrameBuffer::markSlotReadyAndSetInfo(
-    int slotId,
-    CUVIDPARSERDISPINFO* dispInfo) {
-  auto it = map_.find(slotId);
-  TORCH_CHECK(
-      it != map_.end(),
-      "Could not find matching slot with slotId ",
-      slotId,
-      ". This should never happen.");
-
-  TORCH_CHECK(
-      it->second.state == SlotState::BEING_DECODED,
-      "Slot ",
-      slotId,
-      " is not in BEING_DECODED state. This should never happen.");
-  it->second.state = SlotState::READY_FOR_OUTPUT;
-  it->second.dispInfo = *dispInfo;
-}
-
-void BetaCudaDeviceInterface::FrameBuffer::free(int slotId) {
-  auto it = map_.find(slotId);
-  TORCH_CHECK(
-      it != map_.end(),
-      "Tried to free non-existing slot with slotId",
-      slotId,
-      ". This should never happen.");
-  map_.erase(it);
-}
-
-BetaCudaDeviceInterface::FrameBuffer::Slot*
-BetaCudaDeviceInterface::FrameBuffer::findReadySlotWithLowestPts() {
-  Slot* outputSlot = nullptr;
-  for (auto& [_, slot] : map_) {
-    if (slot.state == SlotState::READY_FOR_OUTPUT &&
-        (outputSlot == nullptr ||
-         slot.dispInfo.timestamp < outputSlot->dispInfo.timestamp)) {
-      outputSlot = &slot;
-    }
-  }
-  return outputSlot;
-}
-
 } // namespace facebook::torchcodec

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -64,40 +64,6 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   // Apply bitstream filter, modifies packet in-place
   void applyBSF(ReferenceAVPacket& packet);
 
-  class FrameBuffer {
-   public:
-    enum class SlotState { BEING_DECODED, READY_FOR_OUTPUT };
-
-    struct Slot {
-      CUVIDPARSERDISPINFO dispInfo;
-      SlotState state;
-      int slotId;
-
-      explicit Slot(int id, SlotState s) : state(s), slotId(id) {
-        std::memset(&dispInfo, 0, sizeof(dispInfo));
-        TORCH_CHECK(
-            state == SlotState::BEING_DECODED,
-            "Programmer: are you sure you want to create a slot that is not BEING_DECODED?");
-      }
-    };
-
-    FrameBuffer() = default;
-    ~FrameBuffer() = default;
-
-    void markAsBeingDecoded(int slotId);
-    void markSlotReadyAndSetInfo(int slotId, CUVIDPARSERDISPINFO* dispInfo);
-    void free(int slotId);
-    Slot* findReadySlotWithLowestPts();
-
-    void clear() {
-      map_.clear();
-    }
-
-   private:
-    // Map of slotId to Slot
-    std::unordered_map<int, Slot> map_;
-  };
-
   UniqueAVFrame convertCudaFrameToAVFrame(
       CUdeviceptr framePtr,
       unsigned int pitch,
@@ -107,7 +73,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   UniqueCUvideodecoder decoder_;
   CUVIDEOFORMAT videoFormat_ = {};
 
-  FrameBuffer frameBuffer_;
+  std::queue<CUVIDPARSERDISPINFO> readyFrames_;
 
   bool eofSent_ = false;
 
@@ -168,25 +134,25 @@ class BetaCudaDeviceInterface : public DeviceInterface {
 //   parser has accumulated enough data to decode a frame. We send that frame to
 //   the NVDEC hardware for **async** decoding. While that frame is being
 //   decoded, we store a light reference (a Slot) to that frame in the
-//   frameBuffer_, and mark that slot as BEING_DECODED. The value that uniquely
-//   identifies that frame in the frameBuffer_ is its "slotId", which is given
+//   readyFrames_, and mark that slot as BEING_DECODED. The value that uniquely
+//   identifies that frame in the readyFrames_ is its "slotId", which is given
 //   to us by NVCUVID in the callback parameter: picParams->CurrPicIdx.
 // - frameReadyInDisplayOrder(dispInfo)): triggered **in display order** when a
 //   frame is ready to be "displayed" (returned). When it is triggered, we look
-//   up the corresponding frame/slot in the frameBuffer_, using
+//   up the corresponding frame/slot in the readyFrames_, using
 //   dispInfo->picture_index to match it against a given BEING_DECODED slotId.
 //   We mark that frame/slot as READY_FOR_OUTPUT.
 //   Crucially, this callback also tells us the pts of that frame. We store
 //   the pts and other relevant info the slot.
 //
-// Said differently, from the perspective of the frameBuffer_, at any point in
-// time a slot/frame in the frameBuffer_ can be in 3 states:
-// - empty: no slot for that slotId exists in the frameBuffer_
+// Said differently, from the perspective of the readyFrames_, at any point in
+// time a slot/frame in the readyFrames_ can be in 3 states:
+// - empty: no slot for that slotId exists in the readyFrames_
 // - BEING_DECODED: frameReadyForDecoding was triggered for that frame, and the
 //   frame was sent to NVDEC for async decoding. We don't know its pts because
 //   the parser didn't trigger frameReadyInDisplayOrder() for that frame yet.
 // - READY_FOR_OUTPUT: frameReadyInDisplayOrder was triggered for that frame, it
-//   is decoded and ready to be mapped and returned. We know its pts. 
+//   is decoded and ready to be mapped and returned. We know its pts.
 //
 // Because frameReadyInDisplayOrder is triggered in display order, we know that
 // if a slot is READY_FOR_OUTPUT, then all frames with a lower pts are also
@@ -196,4 +162,4 @@ class BetaCudaDeviceInterface : public DeviceInterface {
 // display order. If no slot is READY_FOR_OUTPUT, then we return EAGAIN to
 // indicate the caller should send more packets.
 //
-// Simple, innit?
+// Simple, innit?