Don't let device interfaces take ownership of the frame they convert

convert_av_frame_to_frame_output() took a non-const UniqueAVFrame&, which
let implementations steal the caller's frame. Both CUDA interfaces did:
BetaCudaDeviceInterface moved out of it, and CudaDeviceInterface reassigned
it with the result of maybe_convert_av_frame_to_nv12_or_rgb24().

That was invisible under SingleStreamDecoder, whose frame is a loop local
that dies right after conversion. It is not invisible under the
building-block ops, where the frame is owned by a handle that outlives the
call: the frame ends up freed twice, corrupting the heap.

Take the frame by const reference all the way down the conversion path, so
a callee physically cannot move from it or reset it. Where an
implementation needs a different frame, it now owns that one separately and
binds a const reference to whichever of the two applies -- the idiom
convert_audio_av_frame_to_frame_output() already used.

maybe_convert_av_frame_to_nv12_or_rgb24() returns a null frame to mean "no
conversion needed" instead of handing back the input.

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -824,7 +824,7 @@ void BetaCudaDeviceInterface::flush() {
 }
 
 UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
-    UniqueAVFrame& cpu_frame,
+    const UniqueAVFrame& cpu_frame,
     AVPixelFormat target_pix_fmt) {
   // This is called in the context of the CPU fallback: the frame was decoded on
   // the CPU, and in this function we convert that frame into NV12 or P016
@@ -967,7 +967,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
 }
 
 void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
-    UniqueAVFrame& av_frame,
+    const UniqueAVFrame& av_frame,
     FrameOutput& frame_output,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
   if (cpu_fallback_) {
@@ -1003,15 +1003,16 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
   // may round them up to even.
   FrameDims original_dims(av_frame->height, av_frame->width);
 
-  UniqueAVFrame gpu_frame;
+  // On the CPU fallback we own the GPU frame we just created; otherwise the
+  // input frame is already what we need, and it's only borrowed.
+  UniqueAVFrame transferred_frame;
   if (cpu_fallback_) {
     AVPixelFormat target_pix_fmt = (output_dtype_ == OutputDtype::FLOAT32)
         ? AV_PIX_FMT_P016LE
         : AV_PIX_FMT_NV12;
-    gpu_frame = transfer_cpu_frame_to_gpu(av_frame, target_pix_fmt);
-  } else {
-    gpu_frame = std::move(av_frame);
+    transferred_frame = transfer_cpu_frame_to_gpu(av_frame, target_pix_fmt);
   }
+  const UniqueAVFrame& gpu_frame = cpu_fallback_ ? transferred_frame : av_frame;
 
   STD_TORCH_CHECK(
       gpu_frame->format == AV_PIX_FMT_NV12 ||

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -52,7 +52,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
       OutputDtype requested_dtype) const override;
 
   void convert_av_frame_to_frame_output(
-      UniqueAVFrame& av_frame,
+      const UniqueAVFrame& av_frame,
       FrameOutput& frame_output,
       std::optional<torch::stable::Tensor> pre_allocated_output_tensor)
       override;
@@ -92,7 +92,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
       const CUVIDPARSERDISPINFO& disp_info);
 
   UniqueAVFrame transfer_cpu_frame_to_gpu(
-      UniqueAVFrame& cpu_frame,
+      const UniqueAVFrame& cpu_frame,
       AVPixelFormat target_pix_fmt);
 
   void apply_rotation(

src/torchcodec/_core/ColorConverter.cpp

@@ -45,7 +45,7 @@ ColorConverter::ColorConverter(
       /*resized_output_dims=*/std::nullopt);
 }
 
-torch::stable::Tensor ColorConverter::convert(UniqueAVFrame& av_frame) {
+torch::stable::Tensor ColorConverter::convert(const UniqueAVFrame& av_frame) {
   FrameOutput frame_output;
   device_interface_->convert_av_frame_to_frame_output(
       av_frame, frame_output, std::nullopt);

src/torchcodec/_core/ColorConverter.h

@@ -21,7 +21,7 @@ class FORCE_PUBLIC_VISIBILITY ColorConverter {
       const StableDevice& device = StableDevice(kStableCPU),
       std::string_view device_variant = "default");
 
-  torch::stable::Tensor convert(UniqueAVFrame& av_frame);
+  torch::stable::Tensor convert(const UniqueAVFrame& av_frame);
 
  private:
   std::unique_ptr<DeviceInterface> device_interface_;

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -200,7 +200,7 @@ ColorConversionLibrary CpuDeviceInterface::get_color_conversion_library(
 }
 
 void CpuDeviceInterface::convert_av_frame_to_frame_output(
-    UniqueAVFrame& av_frame,
+    const UniqueAVFrame& av_frame,
     FrameOutput& frame_output,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
   STD_TORCH_CHECK(initialized_, "CpuDeviceInterface was not initialized.");
@@ -223,7 +223,7 @@ void CpuDeviceInterface::convert_av_frame_to_frame_output(
 // Dimension order of the preAllocatedOutputTensor must be HWC, regardless of
 // `dimension_order` parameter. It's up to callers to re-shape it if needed.
 void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
-    UniqueAVFrame& av_frame,
+    const UniqueAVFrame& av_frame,
     FrameOutput& frame_output,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
   // Note that we ignore the dimensions from the metadata; we don't even bother
@@ -350,7 +350,7 @@ CpuDeviceInterface::convert_av_frame_to_tensor_using_filter_graph(
 }
 
 void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
-    UniqueAVFrame& src_av_frame,
+    const UniqueAVFrame& src_av_frame,
     FrameOutput& frame_output) {
   AVSampleFormat src_sample_format =
       static_cast<AVSampleFormat>(src_av_frame->format);

src/torchcodec/_core/CpuDeviceInterface.h

@@ -41,7 +41,7 @@ class CpuDeviceInterface : public DeviceInterface {
       override;
 
   void convert_av_frame_to_frame_output(
-      UniqueAVFrame& av_frame,
+      const UniqueAVFrame& av_frame,
       FrameOutput& frame_output,
       std::optional<torch::stable::Tensor> pre_allocated_output_tensor)
       override;
@@ -59,11 +59,11 @@ class CpuDeviceInterface : public DeviceInterface {
 
  private:
   void convert_audio_av_frame_to_frame_output(
-      UniqueAVFrame& src_av_frame,
+      const UniqueAVFrame& src_av_frame,
       FrameOutput& frame_output);
 
   void convert_video_av_frame_to_frame_output(
-      UniqueAVFrame& av_frame,
+      const UniqueAVFrame& av_frame,
       FrameOutput& frame_output,
       std::optional<torch::stable::Tensor> pre_allocated_output_tensor);
 

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -146,7 +146,7 @@ void CudaDeviceInterface::register_hardware_device_with_codec(
 }
 
 UniqueAVFrame CudaDeviceInterface::maybe_convert_av_frame_to_nv12_or_rgb24(
-    UniqueAVFrame& av_frame) {
+    const UniqueAVFrame& av_frame) {
   // We need FFmpeg filters to handle those conversion cases which are not
   // directly implemented in CUDA or CPU device interface (in case of a
   // fallback).
@@ -155,7 +155,7 @@ UniqueAVFrame CudaDeviceInterface::maybe_convert_av_frame_to_nv12_or_rgb24(
   // skipping filters context as CPU device interface will handle everything for
   // us.
   if (av_frame->format != AV_PIX_FMT_CUDA) {
-    return std::move(av_frame);
+    return UniqueAVFrame{};
   }
 
   auto hw_frames_ctx =
@@ -169,7 +169,7 @@ UniqueAVFrame CudaDeviceInterface::maybe_convert_av_frame_to_nv12_or_rgb24(
 
   // If the frame is already in NV12 format, we don't need to do anything.
   if (actual_format == AV_PIX_FMT_NV12) {
-    return std::move(av_frame);
+    return UniqueAVFrame{};
   }
 
   AVPixelFormat output_format;
@@ -237,18 +237,21 @@ UniqueAVFrame CudaDeviceInterface::maybe_convert_av_frame_to_nv12_or_rgb24(
 }
 
 void CudaDeviceInterface::convert_av_frame_to_frame_output(
-    UniqueAVFrame& av_frame,
+    const UniqueAVFrame& input_av_frame,
     FrameOutput& frame_output,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
   validate_pre_allocated_tensor_shape(
       pre_allocated_output_tensor,
-      FrameDims(av_frame->height, av_frame->width));
+      FrameDims(input_av_frame->height, input_av_frame->width));
 
   has_decoded_frame_ = true;
 
   // All of our CUDA decoding assumes NV12 format. We handle non-NV12 formats by
   // converting them to NV12.
-  av_frame = maybe_convert_av_frame_to_nv12_or_rgb24(av_frame);
+  UniqueAVFrame converted_av_frame =
+      maybe_convert_av_frame_to_nv12_or_rgb24(input_av_frame);
+  const UniqueAVFrame& av_frame =
+      converted_av_frame ? converted_av_frame : input_av_frame;
 
   if (av_frame->format != AV_PIX_FMT_CUDA) {
     // The frame's format is AV_PIX_FMT_CUDA if and only if its content is on

src/torchcodec/_core/CudaDeviceInterface.h

@@ -41,7 +41,7 @@ class CudaDeviceInterface : public DeviceInterface {
       AVCodecContext* codec_context) override;
 
   void convert_av_frame_to_frame_output(
-      UniqueAVFrame& av_frame,
+      const UniqueAVFrame& av_frame,
       FrameOutput& frame_output,
       std::optional<torch::stable::Tensor> pre_allocated_output_tensor)
       override;
@@ -63,9 +63,10 @@ class CudaDeviceInterface : public DeviceInterface {
  private:
   // Our CUDA decoding code assumes NV12 format. In order to handle other
   // kinds of input, we need to convert them to NV12. Our current implementation
-  // does this using filtergraph.
+  // does this using filtergraph. Returns a null frame when no conversion is
+  // needed, i.e. when the input frame can be used as-is.
   UniqueAVFrame maybe_convert_av_frame_to_nv12_or_rgb24(
-      UniqueAVFrame& av_frame);
+      const UniqueAVFrame& av_frame);
 
   // We sometimes encounter frames that cannot be decoded on the CUDA device.
   // Rather than erroring out, we decode them on the CPU.

src/torchcodec/_core/DeviceInterface.h

@@ -98,8 +98,13 @@ class DeviceInterface {
     return requested_dtype;
   }
 
+  // The caller retains ownership of the frame, and keeps using it after this
+  // call returns (the building-block ops hand out frames that live in Python).
+  // Implementations must not take ownership of it, and must not free it.
+  // Implementations needing a different frame, e.g. after a format conversion,
+  // must allocate their own and leave this one alone.
   virtual void convert_av_frame_to_frame_output(
-      UniqueAVFrame& av_frame,
+      const UniqueAVFrame& av_frame,
       FrameOutput& frame_output,
       std::optional<torch::stable::Tensor> pre_allocated_output_tensor =
           std::nullopt) = 0;

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -1573,7 +1573,7 @@ UniqueAVFrame SingleStreamDecoder::decode_av_frame(
 // --------------------------------------------------------------------------
 
 FrameOutput SingleStreamDecoder::convert_av_frame_to_frame_output(
-    UniqueAVFrame& av_frame,
+    const UniqueAVFrame& av_frame,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
   // Convert the frame to tensor.
   FrameOutput frame_output;