From c9e019df8196b5c59bfec618657f2420c6b612d5 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Mon, 3 Aug 2026 22:19:43 +0100 Subject: [PATCH 1/2] 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. --- src/torchcodec/_core/BetaCudaDeviceInterface.cpp | 13 +++++++------ src/torchcodec/_core/BetaCudaDeviceInterface.h | 4 ++-- src/torchcodec/_core/ColorConverter.cpp | 2 +- src/torchcodec/_core/ColorConverter.h | 2 +- src/torchcodec/_core/CpuDeviceInterface.cpp | 6 +++--- src/torchcodec/_core/CpuDeviceInterface.h | 6 +++--- src/torchcodec/_core/CudaDeviceInterface.cpp | 15 +++++++++------ src/torchcodec/_core/CudaDeviceInterface.h | 7 ++++--- src/torchcodec/_core/DeviceInterface.h | 7 ++++++- src/torchcodec/_core/SingleStreamDecoder.cpp | 2 +- src/torchcodec/_core/SingleStreamDecoder.h | 2 +- src/torchcodec/_core/color_conversion.cpp | 2 +- src/torchcodec/_core/color_conversion.h | 2 +- .../ThirdPartyInterfaceTest.cpp | 2 +- 14 files changed, 41 insertions(+), 31 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 7f47d327e..211e7c633 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/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 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 || diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index 1812a341c..970932475 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/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 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( diff --git a/src/torchcodec/_core/ColorConverter.cpp b/src/torchcodec/_core/ColorConverter.cpp index 6263af1cd..d1f75a2b6 100644 --- a/src/torchcodec/_core/ColorConverter.cpp +++ b/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); diff --git a/src/torchcodec/_core/ColorConverter.h b/src/torchcodec/_core/ColorConverter.h index cf6368735..c2d8c00ca 100644 --- a/src/torchcodec/_core/ColorConverter.h +++ b/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 device_interface_; diff --git a/src/torchcodec/_core/CpuDeviceInterface.cpp b/src/torchcodec/_core/CpuDeviceInterface.cpp index 46731ea0a..8a9acebd7 100644 --- a/src/torchcodec/_core/CpuDeviceInterface.cpp +++ b/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 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 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(src_av_frame->format); diff --git a/src/torchcodec/_core/CpuDeviceInterface.h b/src/torchcodec/_core/CpuDeviceInterface.h index 1a3166e66..85edde9dc 100644 --- a/src/torchcodec/_core/CpuDeviceInterface.h +++ b/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 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 pre_allocated_output_tensor); diff --git a/src/torchcodec/_core/CudaDeviceInterface.cpp b/src/torchcodec/_core/CudaDeviceInterface.cpp index d1962ac5a..bb7ae02a9 100644 --- a/src/torchcodec/_core/CudaDeviceInterface.cpp +++ b/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 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 diff --git a/src/torchcodec/_core/CudaDeviceInterface.h b/src/torchcodec/_core/CudaDeviceInterface.h index c77870230..18ebce234 100644 --- a/src/torchcodec/_core/CudaDeviceInterface.h +++ b/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 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. diff --git a/src/torchcodec/_core/DeviceInterface.h b/src/torchcodec/_core/DeviceInterface.h index e90915e20..c5de6d3d2 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/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 pre_allocated_output_tensor = std::nullopt) = 0; diff --git a/src/torchcodec/_core/SingleStreamDecoder.cpp b/src/torchcodec/_core/SingleStreamDecoder.cpp index 28eac5e75..2dc72613e 100644 --- a/src/torchcodec/_core/SingleStreamDecoder.cpp +++ b/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 pre_allocated_output_tensor) { // Convert the frame to tensor. FrameOutput frame_output; diff --git a/src/torchcodec/_core/SingleStreamDecoder.h b/src/torchcodec/_core/SingleStreamDecoder.h index e14a31661..b2ce092a7 100644 --- a/src/torchcodec/_core/SingleStreamDecoder.h +++ b/src/torchcodec/_core/SingleStreamDecoder.h @@ -282,7 +282,7 @@ class FORCE_PUBLIC_VISIBILITY SingleStreamDecoder { torch::stable::Tensor& tensor); FrameOutput convert_av_frame_to_frame_output( - UniqueAVFrame& av_frame, + const UniqueAVFrame& av_frame, std::optional pre_allocated_output_tensor = std::nullopt); diff --git a/src/torchcodec/_core/color_conversion.cpp b/src/torchcodec/_core/color_conversion.cpp index 22b9729af..b87b79a10 100644 --- a/src/torchcodec/_core/color_conversion.cpp +++ b/src/torchcodec/_core/color_conversion.cpp @@ -189,7 +189,7 @@ void compute_rgb_to_yuv_matrix( } torch::stable::Tensor convert_yuv_frame_to_rgb( - UniqueAVFrame& av_frame, + const UniqueAVFrame& av_frame, const StableDevice& device, cudaStream_t nvdec_stream, std::optional pre_allocated_output_tensor, diff --git a/src/torchcodec/_core/color_conversion.h b/src/torchcodec/_core/color_conversion.h index 74c940b1e..9614c9034 100644 --- a/src/torchcodec/_core/color_conversion.h +++ b/src/torchcodec/_core/color_conversion.h @@ -84,7 +84,7 @@ void launch_p016_to_rgb16_kernel( // outputDims: desired output size; if the frame was rounded up to even // dimensions, the result is cropped back to outputDims. torch::stable::Tensor convert_yuv_frame_to_rgb( - UniqueAVFrame& av_frame, + const UniqueAVFrame& av_frame, const StableDevice& device, cudaStream_t nvdec_stream, std::optional pre_allocated_output_tensor, diff --git a/test/third-party-interface/ThirdPartyInterfaceTest.cpp b/test/third-party-interface/ThirdPartyInterfaceTest.cpp index 8960ac7b8..bc4e6a615 100644 --- a/test/third-party-interface/ThirdPartyInterfaceTest.cpp +++ b/test/third-party-interface/ThirdPartyInterfaceTest.cpp @@ -24,7 +24,7 @@ class DummyDeviceInterface : public DeviceInterface { } void convert_av_frame_to_frame_output( - UniqueAVFrame& av_frame, + const UniqueAVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor = std::nullopt) override {} From df0e0d21034d44413536bce27e924bc7de55a29d Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Mon, 3 Aug 2026 22:31:59 +0100 Subject: [PATCH 2/2] comments --- src/torchcodec/_core/BetaCudaDeviceInterface.cpp | 2 -- src/torchcodec/_core/DeviceInterface.h | 5 ----- 2 files changed, 7 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 211e7c633..e9138a7d5 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -1003,8 +1003,6 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( // may round them up to even. FrameDims original_dims(av_frame->height, av_frame->width); - // 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) diff --git a/src/torchcodec/_core/DeviceInterface.h b/src/torchcodec/_core/DeviceInterface.h index c5de6d3d2..c5a571cb5 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/src/torchcodec/_core/DeviceInterface.h @@ -98,11 +98,6 @@ 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( const UniqueAVFrame& av_frame, FrameOutput& frame_output,