Drop releaseContext

Signed-off-by: Dmitry Rogozhkin <[email protected]>

Author: dvrogozh

src/torchcodec/_core/CudaDevice.cpp

@@ -53,7 +53,7 @@ torch::DeviceIndex getFFMPEGCompatibleDeviceIndex(const torch::Device& device) {
 
 void addToCacheIfCacheHasCapacity(
     const torch::Device& device,
-    AVCodecContext* codecContext) {
+    AVBufferRef* hwContext) {
   torch::DeviceIndex deviceIndex = getFFMPEGCompatibleDeviceIndex(device);
   if (static_cast<int>(deviceIndex) >= MAX_CUDA_GPUS) {
     return;
@@ -64,8 +64,7 @@ void addToCacheIfCacheHasCapacity(
           MAX_CONTEXTS_PER_GPU_IN_CACHE) {
     return;
   }
-  g_cached_hw_device_ctxs[deviceIndex].push_back(codecContext->hw_device_ctx);
-  codecContext->hw_device_ctx = nullptr;
+  g_cached_hw_device_ctxs[deviceIndex].push_back(av_buffer_ref(hwContext));
 }
 
 AVBufferRef* getFromCache(const torch::Device& device) {
@@ -170,17 +169,23 @@ CudaDevice::CudaDevice(const torch::Device& device) : DeviceInterface(device) {
   }
 }
 
-void CudaDevice::releaseContext(AVCodecContext* codecContext) {
-  addToCacheIfCacheHasCapacity(device_, codecContext);
+CudaDevice::~CudaDevice() {
+  if (ctx_) {
+    addToCacheIfCacheHasCapacity(device_, ctx_);
+    av_buffer_unref(&ctx_);
+  }
 }
 
 void CudaDevice::initializeContext(AVCodecContext* codecContext) {
+  TORCH_CHECK(!ctx_, "FFmpeg HW device context already initialized");
+
   // It is important for pytorch itself to create the cuda context. If ffmpeg
   // creates the context it may not be compatible with pytorch.
   // This is a dummy tensor to initialize the cuda context.
   torch::Tensor dummyTensorForCudaInitialization = torch::empty(
       {1}, torch::TensorOptions().dtype(torch::kUInt8).device(device_));
-  codecContext->hw_device_ctx = getCudaContext(device_);
+  ctx_ = getCudaContext(device_);
+  codecContext->hw_device_ctx = av_buffer_ref(ctx_);
   return;
 }
 

src/torchcodec/_core/CudaDevice.h

@@ -14,7 +14,7 @@ class CudaDevice : public DeviceInterface {
  public:
   CudaDevice(const torch::Device& device);
 
-  virtual ~CudaDevice(){};
+  virtual ~CudaDevice();
 
   std::optional<const AVCodec*> findCodec(const AVCodecID& codecId) override;
 
@@ -27,7 +27,8 @@ class CudaDevice : public DeviceInterface {
       std::optional<torch::Tensor> preAllocatedOutputTensor =
           std::nullopt) override;
 
-  void releaseContext(AVCodecContext* codecContext) override;
+ private:
+  AVBufferRef* ctx_ = nullptr;
 };
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/DeviceInterface.h

@@ -46,8 +46,6 @@ class DeviceInterface {
       SingleStreamDecoder::FrameOutput& frameOutput,
       std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt) = 0;
 
-  virtual void releaseContext(AVCodecContext* codecContext) = 0;
-
  protected:
   torch::Device device_;
 };

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -92,15 +92,6 @@ SingleStreamDecoder::SingleStreamDecoder(
   initializeDecoder();
 }
 
-SingleStreamDecoder::~SingleStreamDecoder() {
-  for (auto& [streamIndex, streamInfo] : streamInfos_) {
-    auto& deviceInterface = streamInfo.deviceInterface;
-    if (deviceInterface) {
-      deviceInterface->releaseContext(streamInfo.codecContext.get());
-    }
-  }
-}
-
 void SingleStreamDecoder::initializeDecoder() {
   TORCH_CHECK(!initialized_, "Attempted double initialization.");
 

src/torchcodec/_core/SingleStreamDecoder.h

@@ -23,8 +23,6 @@ class DeviceInterface;
 // Do not call non-const APIs concurrently on the same object.
 class SingleStreamDecoder {
  public:
-  ~SingleStreamDecoder();
-
   // --------------------------------------------------------------------------
   // CONSTRUCTION API
   // --------------------------------------------------------------------------