From 8916180d2b9143e785945fc1662a21f17f76b669 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Mon, 3 Aug 2026 21:35:11 +0100 Subject: [PATCH 01/14] WIP --- .../_core/BetaCudaDeviceInterface.cpp | 97 ++++++++++++++++--- .../_core/BetaCudaDeviceInterface.h | 6 ++ src/torchcodec/_core/ColorConverter.cpp | 19 ++-- src/torchcodec/_core/CpuDeviceInterface.cpp | 3 +- src/torchcodec/_core/DeviceInterface.h | 3 + src/torchcodec/_core/PacketDecoder.cpp | 21 +++- .../decoders/_blocks/_color_converter.py | 8 +- .../decoders/_blocks/_packet_decoder.py | 7 +- 8 files changed, 135 insertions(+), 29 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 7f47d327e..e577b7586 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -289,6 +289,10 @@ void BetaCudaDeviceInterface::initialize_video( const VideoStreamOptions& video_stream_options, const std::vector>& transforms, const std::optional& resized_output_dims) { + // TODO_API_BREAKDOWN ewwwww + if (!av_stream){ + return; + } STD_TORCH_CHECK(av_stream != nullptr, "AVStream cannot be null"); rotation_ = rotation_from_degrees(get_rotation_from_stream(av_stream)); output_dtype_ = video_stream_options.output_dtype; @@ -585,6 +589,7 @@ int BetaCudaDeviceInterface::send_eof_packet() { int BetaCudaDeviceInterface::send_cuvid_packet( CUVIDSOURCEDATAPACKET& cuvid_packet) { CUresult result = cuvidParseVideoData(video_parser_, &cuvid_packet); + printf("cuvidParseVideoData returned %d\n", result); return result == CUDA_SUCCESS ? AVSUCCESS : AVERROR_EXTERNAL; } @@ -794,14 +799,77 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( reinterpret_cast(frame_ptr + (pitch * even_height)); av_frame->data[2] = nullptr; av_frame->data[3] = nullptr; - av_frame->linesize[0] = pitch; - av_frame->linesize[1] = pitch; + // TODO_API_BREAKDOWN_CUDA: Check range before cast? + av_frame->linesize[0] = static_cast(pitch); + av_frame->linesize[1] = static_cast(pitch); av_frame->linesize[2] = 0; av_frame->linesize[3] = 0; return av_frame; } +void nvdec_info_free_callback( + [[maybe_unused]] void* opaque, + [[maybe_unused]] uint8_t* data) { + printf("Freeing standalone frame attached data\n"); + fflush(stdout); + // delete reinterpret_cast(data); +} + +void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { + // TODO_API_BREAKDOWN_CUDA: stongly assumes NVDEC frame, we might want to + // account for CPU fallback frames as well + + // Roughly, the number of bytes an NV12 image takes is: + // num_bytes = len(Y) + len(UV) + // = num_pixels + num_pixels / 2 + // = num_pixels * 3 / 2 + // + // To make it correct, we should use num_pixels = pitch * height, not + // num_pixels = pitch * height. The pitch value also accounts for the data + // size (uint8 vs uint16) so this is also correct for P016. + int64_t even_height = + static_cast(round_up_to_even(av_frame->height)); + int64_t pitch = static_cast(av_frame->linesize[0]); + int64_t num_bytes = pitch * even_height * 3 / 2; + + // TODO_API_BREAKDOWN_CUDA: How the hell do we know that the underlying + // storage isn't freed at the end of this scope?? + auto storage = + torch::stable::empty({num_bytes}, kStableUInt8, std::nullopt, device_); + + // TODO_API_BREAKDOWN_CUDA: Sync with the nvdec stream before copying? + cudaStream_t stream = get_current_cuda_stream(device_.index()); + cudaError_t err = cudaMemcpyAsync( + storage.mutable_data_ptr(), + av_frame->data[0], + static_cast(num_bytes), + cudaMemcpyDeviceToDevice, + stream); + STD_TORCH_CHECK( + err == cudaSuccess, + "Failed to copy NVDEC surface: ", + cudaGetErrorString(err)); + + // TODO_API_BREAKDOWN_CUDA: Should we unmap here? Or let the next + // receive_frame() call do it? + // unmap_previous_frame(); + + auto y_plane = static_cast(storage.mutable_data_ptr()); + av_frame->data[0] = y_plane; + av_frame->data[1] = y_plane + (pitch * even_height); + printf("Returning copied frame\n"); + fflush(stdout); + + auto attached_data = new StandAloneFrameAttachedData(); + // TODO_API_BREAKDOWN_CUDA: We don't *really* need to std::move it I guess? + attached_data->storage = std::move(storage); + + // av_frame->opaque_ref = av_buffer_create(reinterpret_cast(attached_data), sizeof(StandAloneFrameAttachedData), nvdec_info_free_callback, nullptr, 0); + // Intentionally leak storage, just todebug. + av_frame->opaque_ref = av_buffer_create(reinterpret_cast(attached_data), sizeof(StandAloneFrameAttachedData), nullptr, nullptr, 0); +} + void BetaCudaDeviceInterface::flush() { if (cpu_fallback_) { cpu_fallback_->flush(); @@ -810,10 +878,10 @@ void BetaCudaDeviceInterface::flush() { // The NVCUVID docs mention that after seeking, i.e. when flush() is called, // we should send a packet with the CUVID_PKT_DISCONTINUITY flag. The docs - // don't say whether this should be an empty packet, or whether it should be a - // flag on the next non-empty packet. It doesn't matter: neither work :) - // Sending an EOF packet, however, does work. So we do that. And we re-set the - // eofSent_ flag to false because that's not a true EOF notification. + // don't say whether this should be an empty packet, or whether it should be + // a flag on the next non-empty packet. It doesn't matter: neither work :) + // Sending an EOF packet, however, does work. So we do that. And we re-set + // the eofSent_ flag to false because that's not a true EOF notification. send_eof_packet(); eof_sent_ = false; @@ -826,17 +894,17 @@ void BetaCudaDeviceInterface::flush() { UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu( 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 + // 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 // format and send it to the GPU. // We do that in 2 steps: // - First we convert the input CPU frame into an intermediate NV12/P016 CPU // frame using sws_scale. // - Then we allocate GPU memory and copy the CPU frame to the GPU. This // is what we return. - // Since NV12/P016 require even dimensions, the returned frame will have even - // (rounded up) width and height, even if the original CPU frame had odd - // dimensions. + // Since NV12/P016 require even dimensions, the returned frame will have + // even (rounded up) width and height, even if the original CPU frame had + // odd dimensions. STD_TORCH_CHECK(cpu_frame != nullptr, "CPU frame cannot be null"); // NV12 = 1 byte per sample, P016 = 2 bytes per sample @@ -950,9 +1018,9 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu( "Failed to copy frame properties: ", get_ffmpeg_error_string_from_error_code(ret)); - // We need to make sure the CUDA memory is freed properly. Since we allocated - // it ourselves, FFmpeg doesn't know how to free it. We associate a `free` - // callback via opaque_ref that will be called by av_frame_free(). + // We need to make sure the CUDA memory is freed properly. Since we + // allocated it ourselves, FFmpeg doesn't know how to free it. We associate + // a `free` callback via opaque_ref that will be called by av_frame_free(). gpu_frame->opaque_ref = av_buffer_create( nullptr, // data - we don't need any 0, // data size @@ -971,6 +1039,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( FrameOutput& frame_output, std::optional pre_allocated_output_tensor) { if (cpu_fallback_) { + // When the CPU fallback happens, we'll try to run the color-conversion on // GPU by sending those CPU frames to the GPU as NV12 or P016 (See // transferCpuFrameToGpu() below). However, it's not always diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index 1812a341c..e65b3fd5d 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.h @@ -33,6 +33,9 @@ #include "nvcuvid_include/nvcuvid.h" namespace facebook::torchcodec { + struct StandAloneFrameAttachedData{ + torch::stable::Tensor storage; + }; class BetaCudaDeviceInterface : public DeviceInterface { public: @@ -91,6 +94,9 @@ class BetaCudaDeviceInterface : public DeviceInterface { unsigned int pitch, const CUVIDPARSERDISPINFO& disp_info); + + void make_frame_standalone(UniqueAVFrame& av_frame) override; + UniqueAVFrame transfer_cpu_frame_to_gpu( UniqueAVFrame& cpu_frame, AVPixelFormat target_pix_fmt); diff --git a/src/torchcodec/_core/ColorConverter.cpp b/src/torchcodec/_core/ColorConverter.cpp index 6263af1cd..3bd3f8ea0 100644 --- a/src/torchcodec/_core/ColorConverter.cpp +++ b/src/torchcodec/_core/ColorConverter.cpp @@ -27,15 +27,16 @@ ColorConverter::ColorConverter( options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet options.device = device; - // No user transforms and no stream: the converter is stream-agnostic and - // derives everything it needs from each frame. - // - // TODO_API_BREAKDOWN Need to refac/rethink all this. It seems unnatural that - // the color-converter needs its own device_interface_, but at the same time - // the color-conversion *must* be third-party aware, and the only way to - // achieve that for now is via the interface. - // This will become very relevant when we tackle CUDA, so we can defer until - // then. For now this is an OK hack. + // TODO_API_BREAKDOWN It seems unnatural that the color-converter needs its + // own device_interface_, but at the same time the color-conversion *must* be + // third-party aware, and the only way to achieve that for now is via the + // interface. Should at the very least write a note about this design that now + // the DeviceInterface has different modes: decode only, color-convert only, + // and decode+color-convert (which used to be the only mode). + + // TODO_API_BREAKDOWN: we shouldn't call initialize_video here, this is for + // the decoding+color-convert mode. We should do something cleaner e.g. + // initialize_color_convertion_only() std::vector> no_transforms; device_interface_->initialize_video( /*av_stream=*/nullptr, diff --git a/src/torchcodec/_core/CpuDeviceInterface.cpp b/src/torchcodec/_core/CpuDeviceInterface.cpp index 46731ea0a..9f64b2fdf 100644 --- a/src/torchcodec/_core/CpuDeviceInterface.cpp +++ b/src/torchcodec/_core/CpuDeviceInterface.cpp @@ -91,7 +91,8 @@ void CpuDeviceInterface::initialize_video( // time_base_ = av_stream->time_base; // but now that avStrean can be null (to create a standalone color converter) // we need this workaround. This is bad, we need to preserve the previous - // check somehow. + // check somehow. See corresponding TODO in color-converter and packet decoder + // code. time_base_ = (av_stream != nullptr) ? av_stream->time_base : AVRational{1, AV_TIME_BASE}; av_media_type_ = AVMEDIA_TYPE_VIDEO; diff --git a/src/torchcodec/_core/DeviceInterface.h b/src/torchcodec/_core/DeviceInterface.h index e90915e20..99bc9b148 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/src/torchcodec/_core/DeviceInterface.h @@ -138,6 +138,9 @@ class DeviceInterface { return avcodec_receive_frame(codec_context_.get(), av_frame.get()); } + virtual void make_frame_standalone([[maybe_unused]] UniqueAVFrame& av_frame) { + }; + // Flush remaining frames from decoder virtual void flush() { STD_TORCH_CHECK( diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index 23e726f8e..122322871 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -71,6 +71,21 @@ PacketDecoder::PacketDecoder( codec_context_ = create_and_open_codec_context( stream, av_codec, device_interface_.get(), ffmpeg_thread_count); device_interface_->initialize(codec_context_); + + VideoStreamOptions options; + options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet + options.device = device; + + // TODO_API_BREAKDOWN: This isn't right, it's needed only for the NVDEC + // interface. This should probably be initialize_video_only - there's a + // sibling TODO in the ColorConverter code (about color-conversion only.) + std::vector> no_transforms; + device_interface_->initialize_video( + stream, + demuxer.format_context(), + options, + no_transforms, + /*resized_output_dims=*/std::nullopt); } int PacketDecoder::send_packet(AVPacket* packet) { @@ -90,7 +105,11 @@ int PacketDecoder::send_eof() { } int PacketDecoder::receive_frame(UniqueAVFrame& av_frame) { - return device_interface_->receive_frame(av_frame); + int status = device_interface_->receive_frame(av_frame); + if (status == AVSUCCESS) { + device_interface_->make_frame_standalone(av_frame); + } + return status; } } // namespace facebook::torchcodec diff --git a/src/torchcodec/decoders/_blocks/_color_converter.py b/src/torchcodec/decoders/_blocks/_color_converter.py index 69e3bf3c7..c6d6156bf 100644 --- a/src/torchcodec/decoders/_blocks/_color_converter.py +++ b/src/torchcodec/decoders/_blocks/_color_converter.py @@ -30,8 +30,12 @@ class ColorConverter: block is intentionally stream-agnostic. """ - def __init__(self): - self._handle = _blocks_create_color_converter() + # TODO_API_BREAKDOWN: device default should be None + # TODO_API_BREAKDOWN: add checks for coupling between device param of + # PacketDecoder and ColorConverter. What if one is CPU and the other is + # CUDA? What if they're different CUDA devices? Maybe we should just error. + def __init__(self, device="cpu"): + self._handle = _blocks_create_color_converter(device=device) def convert(self, decoded_frame: DecodedFrame) -> Frame: data = _blocks_convert_frame(self._handle, decoded_frame._handle) diff --git a/src/torchcodec/decoders/_blocks/_packet_decoder.py b/src/torchcodec/decoders/_blocks/_packet_decoder.py index 506096062..53bc795ba 100644 --- a/src/torchcodec/decoders/_blocks/_packet_decoder.py +++ b/src/torchcodec/decoders/_blocks/_packet_decoder.py @@ -31,8 +31,11 @@ class PacketDecoder: on your own threads. """ - def __init__(self, demuxer: Demuxer): - self._handle = _blocks_create_packet_decoder(demuxer._handle, num_threads=1) + # TODO_API_BREAKDOWN: device default should be None + def __init__(self, demuxer: Demuxer, device="cpu"): + self._handle = _blocks_create_packet_decoder( + demuxer._handle, num_threads=1, device=device + ) def _drain(self) -> list[DecodedFrame]: frames = [] From 6899cb5966b7ec1fb254111d7bc390e319e2d69a Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Mon, 3 Aug 2026 23:16:35 +0100 Subject: [PATCH 02/14] Pass AVFrames by reference, not by smart pointer Everything downstream of an AVFrame's owner took `UniqueAVFrame&` or `const UniqueAVFrame&`. That's a constraint on the caller's storage rather than a statement about what the function does, and it has two costs. A non-const `UniqueAVFrame&` lets a callee take ownership of the caller's frame. Both CUDA interfaces did: BetaCudaDeviceInterface moved out of it, CudaDeviceInterface reassigned it. That is invisible under SingleStreamDecoder, whose frame is a loop local that dies right after conversion, but the building-block ops own their frame in a handle that outlives the call, so the frame gets freed twice. `const UniqueAVFrame&` doesn't allow the steal, but it still forces anyone holding a plain AVFrame* -- which is what the ops' tensor handle really is -- to manufacture a unique_ptr just to make the call, and manufacturing a second owner for an already-owned object is its own bug factory. So: functions that only look at a frame now take `const AVFrame&`, and the one that writes to it takes `AVFrame&`. Ownership stays with whoever actually owns the frame. Producers (receive_frame) keep `UniqueAVFrame&` because they really do hand back ownership. With that, the ops layer needs no ownership sleight-of-hand: wrap_pointer_to_tensor() gains a deleter parameter, so the one generic handle covers Demuxer/PacketDecoder/ColorConverter and the FFmpeg types too, and both bespoke wrap_*_pointer_to_tensor() functions go away. Demuxer::next_packet() returns UniqueAVPacket rather than a raw pointer plus a comment telling the caller to free it. The encoder is left alone: it owns and mutates its frames, and it uses a null frame as the flush signal, so a reference is the wrong shape there. --- .../_core/BetaCudaDeviceInterface.cpp | 41 ++++----- .../_core/BetaCudaDeviceInterface.h | 4 +- src/torchcodec/_core/ColorConverter.cpp | 2 +- src/torchcodec/_core/ColorConverter.h | 2 +- src/torchcodec/_core/CpuDeviceInterface.cpp | 41 +++++---- src/torchcodec/_core/CpuDeviceInterface.h | 8 +- src/torchcodec/_core/CudaDeviceInterface.cpp | 43 ++++----- src/torchcodec/_core/CudaDeviceInterface.h | 7 +- src/torchcodec/_core/Demuxer.cpp | 8 +- src/torchcodec/_core/Demuxer.h | 7 +- src/torchcodec/_core/DeviceInterface.cpp | 10 +-- src/torchcodec/_core/DeviceInterface.h | 4 +- src/torchcodec/_core/Encoder.cpp | 6 +- src/torchcodec/_core/FFMPEGCommon.cpp | 82 +++++++++-------- src/torchcodec/_core/FFMPEGCommon.h | 20 +++-- src/torchcodec/_core/FilterGraph.cpp | 4 +- src/torchcodec/_core/FilterGraph.h | 2 +- src/torchcodec/_core/SingleStreamDecoder.cpp | 27 +++--- src/torchcodec/_core/SingleStreamDecoder.h | 4 +- src/torchcodec/_core/SwScale.cpp | 12 +-- src/torchcodec/_core/SwScale.h | 4 +- src/torchcodec/_core/color_conversion.cpp | 26 +++--- src/torchcodec/_core/color_conversion.h | 2 +- src/torchcodec/_core/custom_ops.cpp | 87 ++++--------------- .../ThirdPartyInterfaceTest.cpp | 2 +- 25 files changed, 201 insertions(+), 254 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 7f47d327e..82ece2b7b 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -752,7 +752,8 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( // Note that we used to rely on videoFormat_.frame_rate for this, but that // proved less accurate than FFmpeg. set_duration( - av_frame, compute_safe_duration(frame_rate_avg_from_ffmpeg_, time_base_)); + *av_frame, + compute_safe_duration(frame_rate_avg_from_ffmpeg_, time_base_)); // We need to assign the frame colorspace. This is crucial for proper color // conversion. NVCUVID stores that in the matrix_coefficients field, but @@ -824,7 +825,7 @@ void BetaCudaDeviceInterface::flush() { } UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu( - UniqueAVFrame& cpu_frame, + const AVFrame& 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 @@ -838,15 +839,14 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu( // (rounded up) width and height, even if the original CPU frame had odd // dimensions. - STD_TORCH_CHECK(cpu_frame != nullptr, "CPU frame cannot be null"); // NV12 = 1 byte per sample, P016 = 2 bytes per sample STD_TORCH_CHECK( target_pix_fmt == AV_PIX_FMT_NV12 || target_pix_fmt == AV_PIX_FMT_P016LE, "targetPixFmt must be NV12 or P016LE"); int bytes_per_sample = (target_pix_fmt == AV_PIX_FMT_P016LE) ? 2 : 1; - int width = cpu_frame->width; - int height = cpu_frame->height; + int width = cpu_frame.width; + int height = cpu_frame.height; int even_width = round_up_to_even(width); int even_height = round_up_to_even(height); @@ -868,8 +868,8 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu( SwsConfig sws_config( width, height, - static_cast(cpu_frame->format), - cpu_frame->colorspace, + static_cast(cpu_frame.format), + cpu_frame.colorspace, even_width, even_height, target_pix_fmt); @@ -881,8 +881,8 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu( int converted_height = sws_scale( sws_context_.get(), - cpu_frame->data, - cpu_frame->linesize, + cpu_frame.data, + cpu_frame.linesize, 0, height, intermediate_cpu_frame->data, @@ -944,7 +944,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu( "Failed to copy UV plane to GPU: ", cudaGetErrorString(err)); - ret = av_frame_copy_props(gpu_frame.get(), cpu_frame.get()); + ret = av_frame_copy_props(gpu_frame.get(), &cpu_frame); STD_TORCH_CHECK( ret >= 0, "Failed to copy frame properties: ", @@ -967,7 +967,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu( } void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( - UniqueAVFrame& av_frame, + const AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor) { if (cpu_fallback_) { @@ -979,7 +979,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( // do the color conversion on the CPU and then send the full RGB frame to // the GPU. const AVPixFmtDescriptor* desc = - av_pix_fmt_desc_get(static_cast(av_frame->format)); + av_pix_fmt_desc_get(static_cast(av_frame.format)); bool is444 = desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0; if (is444) { FrameOutput cpu_frame_output; @@ -1001,28 +1001,29 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( // Capture original dimensions before transferCpuFrameToGpu() // may round them up to even. - FrameDims original_dims(av_frame->height, av_frame->width); + 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 we only observe it. + 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 AVFrame& gpu_frame = cpu_fallback_ ? *transferred_frame : av_frame; STD_TORCH_CHECK( - gpu_frame->format == AV_PIX_FMT_NV12 || - gpu_frame->format == AV_PIX_FMT_P016LE, + gpu_frame.format == AV_PIX_FMT_NV12 || + gpu_frame.format == AV_PIX_FMT_P016LE, "Expected NV12 or P016LE format frame"); cudaStream_t nvdec_stream = get_current_cuda_stream(device_.index()); auto convert_frame = [&](std::optional pre_alloc) -> torch::stable::Tensor { - bool is_p016 = (gpu_frame->format == AV_PIX_FMT_P016LE); + bool is_p016 = (gpu_frame.format == AV_PIX_FMT_P016LE); int bit_depth = 8; if (is_p016) { bit_depth = cpu_fallback_ diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index 1812a341c..e2d0e3518 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 AVFrame& 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 AVFrame& 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..2e06dd617 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 AVFrame& 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..d06f2d4a2 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 AVFrame& av_frame); private: std::unique_ptr device_interface_; diff --git a/src/torchcodec/_core/CpuDeviceInterface.cpp b/src/torchcodec/_core/CpuDeviceInterface.cpp index 46731ea0a..b6268adbd 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 AVFrame& 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 AVFrame& 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 @@ -239,7 +239,7 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output( // Both cases cause problems for our batch APIs, as we allocate // FrameBatchOutputs based on the the stream metadata. But single-frame APIs // can still work in such situations, so they should. - auto input_dims = FrameDims(av_frame->height, av_frame->width); + auto input_dims = FrameDims(av_frame.height, av_frame.width); auto output_dims = resized_output_dims_.value_or(input_dims); if (pre_allocated_output_tensor.has_value()) { @@ -264,12 +264,12 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output( pre_allocated_output_tensor.value_or(allocate_empty_hwc_tensor( output_dims, kStableCPU, video_stream_options_.output_dtype)); - auto av_frame_format = static_cast(av_frame->format); + auto av_frame_format = static_cast(av_frame.format); SwsConfig sws_config( - av_frame->width, - av_frame->height, + av_frame.width, + av_frame.height, av_frame_format, - av_frame->colorspace, + av_frame.colorspace, output_dims.width, output_dims.height, output_pixel_format_); @@ -326,15 +326,15 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output( torch::stable::Tensor CpuDeviceInterface::convert_av_frame_to_tensor_using_filter_graph( - const UniqueAVFrame& av_frame, + const AVFrame& av_frame, const FrameDims& output_dims) { - auto av_frame_format = static_cast(av_frame->format); + auto av_frame_format = static_cast(av_frame.format); FiltersConfig filters_config( - av_frame->width, - av_frame->height, + av_frame.width, + av_frame.height, av_frame_format, - av_frame->sample_aspect_ratio, + av_frame.sample_aspect_ratio, output_dims.width, output_dims.height, output_pixel_format_, @@ -346,17 +346,17 @@ CpuDeviceInterface::convert_av_frame_to_tensor_using_filter_graph( std::make_unique(filters_config, video_stream_options_); prev_filters_config_ = std::move(filters_config); } - return rgb_av_frame_to_tensor(filter_graph_->convert(av_frame)); + return rgb_av_frame_to_tensor(*filter_graph_->convert(av_frame)); } void CpuDeviceInterface::convert_audio_av_frame_to_frame_output( - UniqueAVFrame& src_av_frame, + const AVFrame& src_av_frame, FrameOutput& frame_output) { AVSampleFormat src_sample_format = - static_cast(src_av_frame->format); + static_cast(src_av_frame.format); AVSampleFormat out_sample_format = AV_SAMPLE_FMT_FLTP; - int src_sample_rate = src_av_frame->sample_rate; + int src_sample_rate = src_av_frame.sample_rate; int out_sample_rate = audio_stream_options_.sample_rate.value_or(src_sample_rate); @@ -397,10 +397,9 @@ void CpuDeviceInterface::convert_audio_av_frame_to_frame_output( out_sample_rate, out_num_channels); } - const UniqueAVFrame& av_frame = - must_convert ? converted_av_frame : src_av_frame; + const AVFrame& av_frame = must_convert ? *converted_av_frame : src_av_frame; - AVSampleFormat format = static_cast(av_frame->format); + AVSampleFormat format = static_cast(av_frame.format); STD_TORCH_CHECK( format == out_sample_format, "Something went wrong, the frame didn't get converted to the desired format. ", @@ -419,7 +418,7 @@ void CpuDeviceInterface::convert_audio_av_frame_to_frame_output( num_channels, " instead."); - auto num_samples = av_frame->nb_samples; + auto num_samples = av_frame.nb_samples; frame_output.data = torch::stable::empty({num_channels, num_samples}); @@ -431,7 +430,7 @@ void CpuDeviceInterface::convert_audio_av_frame_to_frame_output( ++channel, output_channel_data += num_bytes_per_channel) { std::memcpy( output_channel_data, - av_frame->extended_data[channel], + av_frame.extended_data[channel], num_bytes_per_channel); } } diff --git a/src/torchcodec/_core/CpuDeviceInterface.h b/src/torchcodec/_core/CpuDeviceInterface.h index 1a3166e66..2d57eaf7b 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 AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor) override; @@ -59,16 +59,16 @@ class CpuDeviceInterface : public DeviceInterface { private: void convert_audio_av_frame_to_frame_output( - UniqueAVFrame& src_av_frame, + const AVFrame& src_av_frame, FrameOutput& frame_output); void convert_video_av_frame_to_frame_output( - UniqueAVFrame& av_frame, + const AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor); torch::stable::Tensor convert_av_frame_to_tensor_using_filter_graph( - const UniqueAVFrame& av_frame, + const AVFrame& av_frame, const FrameDims& output_dims); ColorConversionLibrary get_color_conversion_library( diff --git a/src/torchcodec/_core/CudaDeviceInterface.cpp b/src/torchcodec/_core/CudaDeviceInterface.cpp index d1962ac5a..1e8c2c34f 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 AVFrame& 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). @@ -154,12 +154,12 @@ UniqueAVFrame CudaDeviceInterface::maybe_convert_av_frame_to_nv12_or_rgb24( // Input frame is on CPU, we will just pass it to CPU device interface, so // 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); + if (av_frame.format != AV_PIX_FMT_CUDA) { + return UniqueAVFrame{}; } auto hw_frames_ctx = - reinterpret_cast(av_frame->hw_frames_ctx->data); + reinterpret_cast(av_frame.hw_frames_ctx->data); STD_TORCH_CHECK( hw_frames_ctx != nullptr, "The AVFrame does not have a 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; @@ -198,19 +198,19 @@ UniqueAVFrame CudaDeviceInterface::maybe_convert_av_frame_to_nv12_or_rgb24( } enum AVPixelFormat frame_format = - static_cast(av_frame->format); + static_cast(av_frame.format); auto new_config = std::make_unique( - av_frame->width, - av_frame->height, + av_frame.width, + av_frame.height, frame_format, - av_frame->sample_aspect_ratio, - av_frame->width, - av_frame->height, + av_frame.sample_aspect_ratio, + av_frame.width, + av_frame.height, output_format, filters.str(), time_base_, - av_buffer_ref(av_frame->hw_frames_ctx)); + av_buffer_ref(av_frame.hw_frames_ctx)); if (!nv12_conversion_ || *nv12_conversion_config_ != *new_config) { nv12_conversion_ = @@ -237,20 +237,23 @@ UniqueAVFrame CudaDeviceInterface::maybe_convert_av_frame_to_nv12_or_rgb24( } void CudaDeviceInterface::convert_av_frame_to_frame_output( - UniqueAVFrame& av_frame, + const AVFrame& 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 AVFrame& av_frame = + converted_av_frame ? *converted_av_frame : input_av_frame; - if (av_frame->format != AV_PIX_FMT_CUDA) { + 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 // the GPU. In this branch, the frame is on the CPU. There are two possible // reasons: @@ -266,7 +269,7 @@ void CudaDeviceInterface::convert_av_frame_to_frame_output( // CUDA device when we're done. enum AVPixelFormat frame_format = - static_cast(av_frame->format); + static_cast(av_frame.format); FrameOutput cpu_frame_output; if (frame_format == AV_PIX_FMT_RGB24) { @@ -302,10 +305,10 @@ void CudaDeviceInterface::convert_av_frame_to_frame_output( // because this is what our color conversion kernel expects. This SHOULD // be enforced by our call to maybeConvertAVFrameToNV12OrRGB24() above. STD_TORCH_CHECK( - av_frame->hw_frames_ctx != nullptr, + av_frame.hw_frames_ctx != nullptr, "The AVFrame does not have a hw_frames_ctx. This should never happen"); AVHWFramesContext* hw_frames_ctx = - reinterpret_cast(av_frame->hw_frames_ctx->data); + reinterpret_cast(av_frame.hw_frames_ctx->data); STD_TORCH_CHECK( hw_frames_ctx != nullptr, "The AVFrame does not have a valid hw_frames_ctx. This should never happen"); @@ -338,7 +341,7 @@ void CudaDeviceInterface::convert_av_frame_to_frame_output( device_, nvdec_stream, pre_allocated_output_tensor, - FrameDims(av_frame->height, av_frame->width), + FrameDims(av_frame.height, av_frame.width), /*isP016=*/false, /*bitDepth=*/8, cached_color_matrix_); diff --git a/src/torchcodec/_core/CudaDeviceInterface.h b/src/torchcodec/_core/CudaDeviceInterface.h index c77870230..d449ccb97 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 AVFrame& 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 AVFrame& 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/Demuxer.cpp b/src/torchcodec/_core/Demuxer.cpp index 4f4ff4c6e..ba46d6450 100644 --- a/src/torchcodec/_core/Demuxer.cpp +++ b/src/torchcodec/_core/Demuxer.cpp @@ -68,12 +68,12 @@ Demuxer::Demuxer( } } -AVPacket* Demuxer::next_packet() { +UniqueAVPacket Demuxer::next_packet() { ReferenceAVPacket packet(auto_packet_); int status = read_next_packet(format_context_.get(), active_stream_index_, packet); if (status == AVERROR_EOF) { - return nullptr; + return UniqueAVPacket{}; } STD_TORCH_CHECK( status >= AVSUCCESS, @@ -82,9 +82,9 @@ AVPacket* Demuxer::next_packet() { // Move the reference out into a fresh, independent packet the caller owns. // This is what makes the packet safe to hand to another thread. - AVPacket* owned = av_packet_alloc(); + UniqueAVPacket owned(av_packet_alloc()); STD_TORCH_CHECK(owned != nullptr, "Failed to allocate AVPacket"); - av_packet_move_ref(owned, packet.get()); + av_packet_move_ref(owned.get(), packet.get()); return owned; } diff --git a/src/torchcodec/_core/Demuxer.h b/src/torchcodec/_core/Demuxer.h index e0cfbed3e..204d08673 100644 --- a/src/torchcodec/_core/Demuxer.h +++ b/src/torchcodec/_core/Demuxer.h @@ -31,10 +31,9 @@ class FORCE_PUBLIC_VISIBILITY Demuxer { const std::string& file_path, std::optional stream_index = std::nullopt); - // Returns the next packet for the active stream as a freshly-allocated, - // owning AVPacket (the caller takes ownership and must av_packet_free it), or - // nullptr at end of stream. - AVPacket* next_packet(); + // Returns the next packet for the active stream as a freshly-allocated + // packet, or a null packet at end of stream. + UniqueAVPacket next_packet(); AVStream* active_stream() const { return stream_; diff --git a/src/torchcodec/_core/DeviceInterface.cpp b/src/torchcodec/_core/DeviceInterface.cpp index 4ad7e8fe1..484b86728 100644 --- a/src/torchcodec/_core/DeviceInterface.cpp +++ b/src/torchcodec/_core/DeviceInterface.cpp @@ -116,16 +116,16 @@ std::unique_ptr create_device_interface( "'"); } -torch::stable::Tensor rgb_av_frame_to_tensor(const UniqueAVFrame& av_frame) { - auto format = static_cast(av_frame->format); +torch::stable::Tensor rgb_av_frame_to_tensor(const AVFrame& av_frame) { + auto format = static_cast(av_frame.format); STD_TORCH_CHECK( format == AV_PIX_FMT_RGB24 || format == AV_PIX_FMT_RGB48, "Expected RGB24 or RGB48 format, got ", (av_get_pix_fmt_name(format) ? av_get_pix_fmt_name(format) : "unknown")); - int height = av_frame->height; - int width = av_frame->width; - AVFrame* cloned_av_frame = av_frame_clone(av_frame.get()); + int height = av_frame.height; + int width = av_frame.width; + AVFrame* cloned_av_frame = av_frame_clone(&av_frame); auto deleter = [cloned_av_frame](void*) { UniqueAVFrame av_frame_to_delete(cloned_av_frame); diff --git a/src/torchcodec/_core/DeviceInterface.h b/src/torchcodec/_core/DeviceInterface.h index e90915e20..ca976fd36 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/src/torchcodec/_core/DeviceInterface.h @@ -99,7 +99,7 @@ class DeviceInterface { } virtual void convert_av_frame_to_frame_output( - UniqueAVFrame& av_frame, + const AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor = std::nullopt) = 0; @@ -207,6 +207,6 @@ create_device_interface( const StableDevice& device, const std::string_view variant = "default"); -torch::stable::Tensor rgb_av_frame_to_tensor(const UniqueAVFrame& av_frame); +torch::stable::Tensor rgb_av_frame_to_tensor(const AVFrame& av_frame); } // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/Encoder.cpp b/src/torchcodec/_core/Encoder.cpp index e3665574c..11808e091 100644 --- a/src/torchcodec/_core/Encoder.cpp +++ b/src/torchcodec/_core/Encoder.cpp @@ -794,7 +794,7 @@ UniqueAVFrame MultiStreamEncoder::maybe_convert_audio_av_frame( AudioStream& audio_stream) { if (static_cast(av_frame->format) == audio_stream.av_codec_context->sample_fmt && - get_num_channels(av_frame) == audio_stream.out_num_channels && + get_num_channels(*av_frame) == audio_stream.out_num_channels && av_frame->sample_rate == audio_stream.out_sample_rate) { // Note: the clone references the same underlying data, it's a cheap copy. return UniqueAVFrame(av_frame_clone(av_frame.get())); @@ -806,7 +806,7 @@ UniqueAVFrame MultiStreamEncoder::maybe_convert_audio_av_frame( audio_stream.av_codec_context->sample_fmt, av_frame->sample_rate, audio_stream.out_sample_rate, - av_frame, + *av_frame, audio_stream.out_num_channels)); } // convertAudioAVFrameSamples uses avFrame's extended_data field, so we ensure @@ -817,7 +817,7 @@ UniqueAVFrame MultiStreamEncoder::maybe_convert_audio_av_frame( "Codec context data and extended_data pointers differ, this is unexpected."); UniqueAVFrame converted_av_frame = convert_audio_av_frame_samples( audio_stream.swr_context, - av_frame, + *av_frame, audio_stream.av_codec_context->sample_fmt, audio_stream.out_sample_rate, audio_stream.out_num_channels); diff --git a/src/torchcodec/_core/FFMPEGCommon.cpp b/src/torchcodec/_core/FFMPEGCommon.cpp index b7b6c1550..dc5b6e8d0 100644 --- a/src/torchcodec/_core/FFMPEGCommon.cpp +++ b/src/torchcodec/_core/FFMPEGCommon.cpp @@ -54,11 +54,11 @@ std::string get_ffmpeg_error_string_from_error_code(int error_code) { return std::string(error_buffer); } -int64_t get_duration(const UniqueAVFrame& av_frame) { +int64_t get_duration(const AVFrame& av_frame) { #if LIBAVUTIL_VERSION_MAJOR < 58 - return av_frame->pkt_duration; + return av_frame.pkt_duration; #else - return av_frame->duration; + return av_frame.duration; #endif } @@ -71,15 +71,15 @@ int64_t get_pts_or_dts(ReferenceAVPacket& packet) { return packet->pts == INT64_MIN ? packet->dts : packet->pts; } -int64_t get_pts_or_dts(const UniqueAVFrame& av_frame) { - return av_frame->pts == INT64_MIN ? av_frame->pkt_dts : av_frame->pts; +int64_t get_pts_or_dts(const AVFrame& av_frame) { + return av_frame.pts == INT64_MIN ? av_frame.pkt_dts : av_frame.pts; } -void set_duration(const UniqueAVFrame& av_frame, int64_t duration) { +void set_duration(AVFrame& av_frame, int64_t duration) { #if LIBAVUTIL_VERSION_MAJOR < 58 - av_frame->pkt_duration = duration; + av_frame.pkt_duration = duration; #else - av_frame->duration = duration; + av_frame.duration = duration; #endif } @@ -148,20 +148,18 @@ const AVSampleFormat* get_supported_output_sample_formats( return supported_sample_formats; } -int get_num_channels(const UniqueAVFrame& av_frame) { +int get_num_channels(const AVFrame& av_frame) { #if LIBAVFILTER_VERSION_MAJOR > 8 || \ (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44) - return av_frame->ch_layout.nb_channels; + return av_frame.ch_layout.nb_channels; #else - int num_channels = - av_get_channel_layout_nb_channels(av_frame->channel_layout); + int num_channels = av_get_channel_layout_nb_channels(av_frame.channel_layout); // Handle FFmpeg 4 bug where channel_layout and num_channels are 0 or unset - // Set values based on av_frame->channels which appears to be correct + // Set values based on av_frame.channels which appears to be correct // to allow successful initialization of SwrContext - if (num_channels == 0 && av_frame->channels > 0) { - av_frame->channel_layout = - av_get_default_channel_layout(av_frame->channels); - num_channels = av_frame->channels; + if (num_channels == 0 && av_frame.channels > 0) { + av_frame.channel_layout = av_get_default_channel_layout(av_frame.channels); + num_channels = av_frame.channels; } return num_channels; #endif @@ -200,15 +198,15 @@ void set_default_channel_layout( #endif } -void set_default_channel_layout(UniqueAVFrame& av_frame, int num_channels) { +void set_default_channel_layout(AVFrame& av_frame, int num_channels) { #if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4 AVChannelLayout channel_layout; av_channel_layout_default(&channel_layout, num_channels); - av_frame->ch_layout = channel_layout; + av_frame.ch_layout = channel_layout; #else uint64_t channel_layout = av_get_default_channel_layout(num_channels); - av_frame->channel_layout = channel_layout; - av_frame->channels = num_channels; + av_frame.channel_layout = channel_layout; + av_frame.channels = num_channels; #endif } @@ -300,10 +298,10 @@ namespace { // - the default channel layout with out_num_channels otherwise. AVChannelLayout get_output_channel_layout( int out_num_channels, - const UniqueAVFrame& src_av_frame) { + const AVFrame& src_av_frame) { AVChannelLayout out_layout; if (out_num_channels == get_num_channels(src_av_frame)) { - out_layout = src_av_frame->ch_layout; + out_layout = src_av_frame.ch_layout; } else { av_channel_layout_default(&out_layout, out_num_channels); } @@ -315,10 +313,10 @@ AVChannelLayout get_output_channel_layout( // Same as above int64_t get_output_channel_layout( int out_num_channels, - const UniqueAVFrame& src_av_frame) { + const AVFrame& src_av_frame) { int64_t out_layout; if (out_num_channels == get_num_channels(src_av_frame)) { - out_layout = src_av_frame->channel_layout; + out_layout = src_av_frame.channel_layout; } else { out_layout = av_get_default_channel_layout(out_num_channels); } @@ -330,21 +328,21 @@ int64_t get_output_channel_layout( // Sets dst_av_frame' channel layout to get_output_channel_layout(): see doc // above void set_channel_layout( - UniqueAVFrame& dst_av_frame, - const UniqueAVFrame& src_av_frame, + AVFrame& dst_av_frame, + const AVFrame& src_av_frame, int out_num_channels) { #if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4 AVChannelLayout out_layout = get_output_channel_layout(out_num_channels, src_av_frame); - auto status = av_channel_layout_copy(&dst_av_frame->ch_layout, &out_layout); + auto status = av_channel_layout_copy(&dst_av_frame.ch_layout, &out_layout); STD_TORCH_CHECK( status == AVSUCCESS, "Couldn't copy channel layout to av_frame: ", get_ffmpeg_error_string_from_error_code(status)); #else - dst_av_frame->channel_layout = + dst_av_frame.channel_layout = get_output_channel_layout(out_num_channels, src_av_frame); - dst_av_frame->channels = out_num_channels; + dst_av_frame.channels = out_num_channels; #endif } @@ -358,7 +356,7 @@ UniqueAVFrame allocate_av_frame( av_frame->nb_samples = num_samples; av_frame->sample_rate = sample_rate; - set_default_channel_layout(av_frame, num_channels); + set_default_channel_layout(*av_frame, num_channels); av_frame->format = sample_format; auto status = av_frame_get_buffer(av_frame.get(), 0); @@ -380,7 +378,7 @@ SwrContext* create_swr_context( AVSampleFormat out_sample_format, int src_sample_rate, int out_sample_rate, - const UniqueAVFrame& src_av_frame, + const AVFrame& src_av_frame, int out_num_channels) { SwrContext* swr_context = nullptr; int status = AVSUCCESS; @@ -392,7 +390,7 @@ SwrContext* create_swr_context( &out_layout, out_sample_format, out_sample_rate, - &src_av_frame->ch_layout, + &src_av_frame.ch_layout, src_sample_format, src_sample_rate, 0, @@ -410,7 +408,7 @@ SwrContext* create_swr_context( out_layout, out_sample_format, out_sample_rate, - src_av_frame->channel_layout, + src_av_frame.channel_layout, src_sample_format, src_sample_rate, 0, @@ -493,7 +491,7 @@ AVFilterContext* create_av_filter_context_with_options( UniqueAVFrame convert_audio_av_frame_samples( const UniqueSwrContext& swr_context, - const UniqueAVFrame& src_av_frame, + const AVFrame& src_av_frame, AVSampleFormat out_sample_format, int out_sample_rate, int out_num_channels) { @@ -502,11 +500,11 @@ UniqueAVFrame convert_audio_av_frame_samples( converted_av_frame, "Could not allocate frame for sample format conversion."); - converted_av_frame->pts = src_av_frame->pts; + converted_av_frame->pts = src_av_frame.pts; converted_av_frame->format = static_cast(out_sample_format); converted_av_frame->sample_rate = out_sample_rate; - int src_sample_rate = src_av_frame->sample_rate; + int src_sample_rate = src_av_frame.sample_rate; if (src_sample_rate != out_sample_rate) { // Note that this is an upper bound on the number of output samples. // `swr_convert()` will likely not fill convertedAVFrame with that many @@ -518,15 +516,15 @@ UniqueAVFrame convert_audio_av_frame_samples( // tighter bound. converted_av_frame->nb_samples = av_rescale_rnd( swr_get_delay(swr_context.get(), src_sample_rate) + - src_av_frame->nb_samples, + src_av_frame.nb_samples, out_sample_rate, src_sample_rate, AV_ROUND_UP); } else { - converted_av_frame->nb_samples = src_av_frame->nb_samples; + converted_av_frame->nb_samples = src_av_frame.nb_samples; } - set_channel_layout(converted_av_frame, src_av_frame, out_num_channels); + set_channel_layout(*converted_av_frame, src_av_frame, out_num_channels); auto status = av_frame_get_buffer(converted_av_frame.get(), 0); STD_TORCH_CHECK( @@ -543,8 +541,8 @@ UniqueAVFrame convert_audio_av_frame_samples( converted_av_frame->extended_data, converted_av_frame->nb_samples, static_cast( - const_cast(src_av_frame->extended_data)), - src_av_frame->nb_samples); + const_cast(src_av_frame.extended_data)), + src_av_frame.nb_samples); // numConvertedSamples can be 0 if we're downsampling by a great factor and // the first frame doesn't contain a lot of samples. It should be handled // properly by the caller. diff --git a/src/torchcodec/_core/FFMPEGCommon.h b/src/torchcodec/_core/FFMPEGCommon.h index 88bb04ceb..e6d57eadf 100644 --- a/src/torchcodec/_core/FFMPEGCommon.h +++ b/src/torchcodec/_core/FFMPEGCommon.h @@ -83,6 +83,8 @@ inline SharedAVCodecContext make_shared_av_codec_context(AVCodecContext* ctx) { using UniqueAVFrame = std::unique_ptr>; +using UniqueAVPacket = + std::unique_ptr>; using UniqueAVFilterGraph = std::unique_ptr< AVFilterGraph, Deleterp>; @@ -206,20 +208,20 @@ std::string get_ffmpeg_error_string_from_error_code(int error_code); // Returns duration from the frame. Abstracted into a function because the // struct member representing duration has changed across the versions we // support. -int64_t get_duration(const UniqueAVFrame& frame); -void set_duration(const UniqueAVFrame& frame, int64_t duration); +int64_t get_duration(const AVFrame& frame); +void set_duration(AVFrame& frame, int64_t duration); // pts accessors that fall back to dts when pts is unset (INT64_MIN). See the // definitions for details. int64_t get_pts_or_dts(ReferenceAVPacket& packet); -int64_t get_pts_or_dts(const UniqueAVFrame& av_frame); +int64_t get_pts_or_dts(const AVFrame& av_frame); const int* get_supported_sample_rates(const AVCodec& av_codec); const AVSampleFormat* get_supported_output_sample_formats( const AVCodec& av_codec); const AVPixelFormat* get_supported_pixel_formats(const AVCodec& av_codec); -int get_num_channels(const UniqueAVFrame& av_frame); +int get_num_channels(const AVFrame& av_frame); int get_num_channels(const SharedAVCodecContext& av_codec_context); int get_num_channels(const AVCodecParameters* codecpar); @@ -227,13 +229,13 @@ void set_default_channel_layout( UniqueAVCodecContext& av_codec_context, int num_channels); -void set_default_channel_layout(UniqueAVFrame& av_frame, int num_channels); +void set_default_channel_layout(AVFrame& av_frame, int num_channels); void validate_num_channels(const AVCodec& av_codec, int num_channels); void set_channel_layout( - UniqueAVFrame& dst_av_frame, - const UniqueAVFrame& src_av_frame, + AVFrame& dst_av_frame, + const AVFrame& src_av_frame, int desired_num_channels); UniqueAVFrame allocate_av_frame( @@ -247,7 +249,7 @@ SwrContext* create_swr_context( AVSampleFormat desired_sample_format, int src_sample_rate, int desired_sample_rate, - const UniqueAVFrame& src_av_frame, + const AVFrame& src_av_frame, int desired_num_channels); // Converts, if needed: @@ -257,7 +259,7 @@ SwrContext* create_swr_context( // createSwrContext must have been previously called with matching parameters. UniqueAVFrame convert_audio_av_frame_samples( const UniqueSwrContext& swr_context, - const UniqueAVFrame& src_av_frame, + const AVFrame& src_av_frame, AVSampleFormat desired_sample_format, int desired_sample_rate, int desired_num_channels); diff --git a/src/torchcodec/_core/FilterGraph.cpp b/src/torchcodec/_core/FilterGraph.cpp index 36c3fa4ae..f3b8d05be 100644 --- a/src/torchcodec/_core/FilterGraph.cpp +++ b/src/torchcodec/_core/FilterGraph.cpp @@ -150,8 +150,8 @@ FilterGraph::FilterGraph( ", provided filters: " + filters_config.filtergraph_str); } -UniqueAVFrame FilterGraph::convert(const UniqueAVFrame& av_frame) { - int status = av_buffersrc_write_frame(source_context_, av_frame.get()); +UniqueAVFrame FilterGraph::convert(const AVFrame& av_frame) { + int status = av_buffersrc_write_frame(source_context_, &av_frame); STD_TORCH_CHECK( status >= AVSUCCESS, "Failed to add frame to buffer source context"); diff --git a/src/torchcodec/_core/FilterGraph.h b/src/torchcodec/_core/FilterGraph.h index 650ea584d..356122e27 100644 --- a/src/torchcodec/_core/FilterGraph.h +++ b/src/torchcodec/_core/FilterGraph.h @@ -48,7 +48,7 @@ class FilterGraph { const FiltersConfig& filters_config, const VideoStreamOptions& video_stream_options); - UniqueAVFrame convert(const UniqueAVFrame& av_frame); + UniqueAVFrame convert(const AVFrame& av_frame); private: UniqueAVFilterGraph filter_graph_; diff --git a/src/torchcodec/_core/SingleStreamDecoder.cpp b/src/torchcodec/_core/SingleStreamDecoder.cpp index 28eac5e75..f6a770dbd 100644 --- a/src/torchcodec/_core/SingleStreamDecoder.cpp +++ b/src/torchcodec/_core/SingleStreamDecoder.cpp @@ -677,12 +677,11 @@ FrameOutput SingleStreamDecoder::get_next_frame() { FrameOutput SingleStreamDecoder::get_next_frame_internal( std::optional pre_allocated_output_tensor) { validate_active_stream(); - UniqueAVFrame av_frame = - decode_av_frame([this](const UniqueAVFrame& av_frame) { - return get_pts_or_dts(av_frame) >= cursor_; - }); + UniqueAVFrame av_frame = decode_av_frame([this](const AVFrame& av_frame) { + return get_pts_or_dts(av_frame) >= cursor_; + }); return convert_av_frame_to_frame_output( - av_frame, pre_allocated_output_tensor); + *av_frame, pre_allocated_output_tensor); } FrameOutput SingleStreamDecoder::get_frame_at_index(int64_t frame_index) { @@ -867,7 +866,7 @@ FrameOutput SingleStreamDecoder::get_frame_played_at(double seconds) { set_cursor_pts_in_seconds(seconds); UniqueAVFrame av_frame = - decode_av_frame([seconds, this](const UniqueAVFrame& av_frame) { + decode_av_frame([seconds, this](const AVFrame& av_frame) { StreamInfo& stream_info = stream_infos_[active_stream_index_]; double frame_start_time = pts_to_seconds(get_pts_or_dts(av_frame), stream_info.time_base); @@ -888,7 +887,7 @@ FrameOutput SingleStreamDecoder::get_frame_played_at(double seconds) { }); // Convert the frame to tensor. - FrameOutput frame_output = convert_av_frame_to_frame_output(av_frame); + FrameOutput frame_output = convert_av_frame_to_frame_output(*av_frame); frame_output.data = maybe_permute_and_convert_to_float32(frame_output.data); return frame_output; } @@ -1241,12 +1240,12 @@ AudioFramesOutput SingleStreamDecoder::get_frames_played_in_range_audio( while (!finished) { try { UniqueAVFrame av_frame = - decode_av_frame([start_pts, stop_pts](const UniqueAVFrame& av_frame) { + decode_av_frame([start_pts, stop_pts](const AVFrame& av_frame) { return start_pts < get_pts_or_dts(av_frame) + get_duration(av_frame) && stop_pts > get_pts_or_dts(av_frame); }); - auto frame_output = convert_av_frame_to_frame_output(av_frame); + auto frame_output = convert_av_frame_to_frame_output(*av_frame); if (!first_frame_pts_seconds.has_value()) { first_frame_pts_seconds = frame_output.pts_seconds; } @@ -1458,7 +1457,7 @@ void SingleStreamDecoder::maybe_seek_to_before_desired_pts() { // -------------------------------------------------------------------------- UniqueAVFrame SingleStreamDecoder::decode_av_frame( - std::function filter_function) { + std::function filter_function) { validate_active_stream(); reset_decode_stats(); @@ -1484,7 +1483,7 @@ UniqueAVFrame SingleStreamDecoder::decode_av_frame( decode_stats_.num_frames_received_by_decoder++; // Is this the kind of frame we're looking for? - if (status == AVSUCCESS && filter_function(av_frame)) { + if (status == AVSUCCESS && filter_function(*av_frame)) { // Yes, this is the frame we'll return; break out of the decoding loop. break; } else if (status == AVSUCCESS) { @@ -1562,8 +1561,8 @@ UniqueAVFrame SingleStreamDecoder::decode_av_frame( // received as frames. Eventually we will either hit AVERROR_EOF from // av_receive_frame() or the user will have seeked to a different location // in the file and that will flush the decoder. - last_decoded_av_frame_pts_ = get_pts_or_dts(av_frame); - last_decoded_av_frame_duration_ = get_duration(av_frame); + last_decoded_av_frame_pts_ = get_pts_or_dts(*av_frame); + last_decoded_av_frame_duration_ = get_duration(*av_frame); return av_frame; } @@ -1573,7 +1572,7 @@ UniqueAVFrame SingleStreamDecoder::decode_av_frame( // -------------------------------------------------------------------------- FrameOutput SingleStreamDecoder::convert_av_frame_to_frame_output( - UniqueAVFrame& av_frame, + const AVFrame& 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..c1ef2a1bd 100644 --- a/src/torchcodec/_core/SingleStreamDecoder.h +++ b/src/torchcodec/_core/SingleStreamDecoder.h @@ -270,7 +270,7 @@ class FORCE_PUBLIC_VISIBILITY SingleStreamDecoder { void maybe_seek_to_before_desired_pts(); UniqueAVFrame decode_av_frame( - std::function filter_function); + std::function filter_function); FrameOutput get_next_frame_internal( std::optional pre_allocated_output_tensor = @@ -282,7 +282,7 @@ class FORCE_PUBLIC_VISIBILITY SingleStreamDecoder { torch::stable::Tensor& tensor); FrameOutput convert_av_frame_to_frame_output( - UniqueAVFrame& av_frame, + const AVFrame& av_frame, std::optional pre_allocated_output_tensor = std::nullopt); diff --git a/src/torchcodec/_core/SwScale.cpp b/src/torchcodec/_core/SwScale.cpp index 48be27043..f0e585eba 100644 --- a/src/torchcodec/_core/SwScale.cpp +++ b/src/torchcodec/_core/SwScale.cpp @@ -52,7 +52,7 @@ SwScale::SwScale(const SwsConfig& config, int sws_flags) } int SwScale::convert( - const UniqueAVFrame& av_frame, + const AVFrame& av_frame, torch::stable::Tensor& output_tensor) { // When resizing is needed, we do sws_scale twice: first convert to output // RGB at original resolution, then resize in output RGB space. This ensures @@ -85,19 +85,19 @@ int SwScale::convert( int color_converted_height = sws_scale( color_conversion_sws_context_.get(), - av_frame->data, - av_frame->linesize, + av_frame.data, + av_frame.linesize, 0, - av_frame->height, + av_frame.height, color_converted_pointers, color_converted_linesizes); STD_TORCH_CHECK( - color_converted_height == av_frame->height, + color_converted_height == av_frame.height, "Color conversion swscale pass failed: colorConvertedHeight != avFrame->height: ", color_converted_height, " != ", - av_frame->height); + av_frame.height); if (needs_resize_) { uint8_t* src_pointers[4] = { diff --git a/src/torchcodec/_core/SwScale.h b/src/torchcodec/_core/SwScale.h index b76d743f4..3b307d579 100644 --- a/src/torchcodec/_core/SwScale.h +++ b/src/torchcodec/_core/SwScale.h @@ -29,9 +29,7 @@ class SwScale { // >8-bit. SwScale(const SwsConfig& config, int sws_flags = SWS_BILINEAR); - int convert( - const UniqueAVFrame& av_frame, - torch::stable::Tensor& output_tensor); + int convert(const AVFrame& av_frame, torch::stable::Tensor& output_tensor); const SwsConfig& get_config() const { return config_; diff --git a/src/torchcodec/_core/color_conversion.cpp b/src/torchcodec/_core/color_conversion.cpp index 22b9729af..0efa6d699 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 AVFrame& av_frame, const StableDevice& device, cudaStream_t nvdec_stream, std::optional pre_allocated_output_tensor, @@ -203,8 +203,8 @@ torch::stable::Tensor convert_yuv_frame_to_rgb( // Dimensions may be odd (NVDEC display area for VP9 etc.). NV12/P016 // color conversion requires even dimensions, so we round up to even // for the kernel, then crop to outputDims. - int even_height = round_up_to_even(av_frame->height); - int even_width = round_up_to_even(av_frame->width); + int even_height = round_up_to_even(av_frame.height); + int even_width = round_up_to_even(av_frame.width); int out_height = output_dims.height; int out_width = output_dims.width; @@ -227,33 +227,33 @@ torch::stable::Tensor convert_yuv_frame_to_rgb( maybe_update_color_matrix( cached_color_matrix, - av_frame->colorspace, - av_frame->color_range, + av_frame.colorspace, + av_frame.color_range, bit_depth, out_scale); if (is_p016) { launch_p016_to_rgb16_kernel( - reinterpret_cast(av_frame->data[0]), - reinterpret_cast(av_frame->data[1]), + reinterpret_cast(av_frame.data[0]), + reinterpret_cast(av_frame.data[1]), dst.mutable_data_ptr(), even_width, even_height, - av_frame->linesize[0], - av_frame->linesize[1], + av_frame.linesize[0], + av_frame.linesize[1], validate_int64_to_int(dst.stride(0) * 2, "dst.stride(0)*2"), bit_depth, cached_color_matrix.matrix, stream); } else { launch_nv12_to_rgb_kernel( - av_frame->data[0], - av_frame->data[1], + av_frame.data[0], + av_frame.data[1], dst.mutable_data_ptr(), even_width, even_height, - av_frame->linesize[0], - av_frame->linesize[1], + av_frame.linesize[0], + av_frame.linesize[1], validate_int64_to_int(dst.stride(0), "dst.stride(0)"), cached_color_matrix.matrix, stream); diff --git a/src/torchcodec/_core/color_conversion.h b/src/torchcodec/_core/color_conversion.h index 74c940b1e..ad3ae29ef 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 AVFrame& av_frame, const StableDevice& device, cudaStream_t nvdec_stream, std::optional pre_allocated_output_tensor, diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp index 76634bc21..73351c99c 100644 --- a/src/torchcodec/_core/custom_ops.cpp +++ b/src/torchcodec/_core/custom_ops.cpp @@ -182,13 +182,17 @@ SingleStreamDecoder* unwrap_tensor_to_get_decoder( } // Generic pointer<->tensor laundering for the building-block handle types -// (Demuxer / PacketDecoder / ColorConverter). Same trick as -// wrap_decoder_pointer_to_tensor: the tensor's data pointer IS the raw pointer, -// with a deleter that deletes the owned object when the handle is dropped. -template -torch::stable::Tensor wrap_pointer_to_tensor(std::unique_ptr ptr) { +// (Demuxer / PacketDecoder / ColorConverter / AVPacket / AVFrame). Same trick +// as wrap_decoder_pointer_to_tensor: the tensor's data pointer IS the raw +// pointer, and the tensor is the owner. Taking the unique_ptr by value is what +// makes the ownership transfer explicit at the call site; the unique_ptr's own +// deleter is what frees the object, so FFmpeg types work here as long as they +// arrive in their UniqueAVXxx alias. +template +torch::stable::Tensor wrap_pointer_to_tensor(std::unique_ptr ptr) { + D object_deleter = ptr.get_deleter(); T* raw = ptr.release(); - auto deleter = [raw](void*) { delete raw; }; + auto deleter = [raw, object_deleter](void*) { object_deleter(raw); }; int64_t sizes[] = {static_cast(sizeof(T*))}; int64_t strides[] = {1}; return torch::stable::from_blob( @@ -206,56 +210,6 @@ T* unwrap_tensor_to_pointer(torch::stable::Tensor& tensor) { return static_cast(tensor.mutable_data_ptr()); } -// Opaque packet/frame handles: launder a raw AVPacket*/AVFrame* through a [1] -// int64 CPU tensor whose data pointer IS the raw pointer, with a deleter that -// frees it. Thread-movable, process-local. -torch::stable::Tensor wrap_packet_pointer_to_tensor(AVPacket* packet) { - auto deleter = [packet](void*) { - AVPacket* p = packet; - av_packet_free(&p); - }; - int64_t sizes[] = {1}; - int64_t strides[] = {1}; - return torch::stable::from_blob( - packet, - {sizes, 1}, - {strides, 1}, - StableDevice(kStableCPU), - kStableInt64, - deleter); -} - -AVPacket* unwrap_tensor_to_packet(torch::stable::Tensor& tensor) { - STD_TORCH_CHECK(tensor.is_contiguous(), "packet handle must be contiguous"); - return static_cast(tensor.mutable_data_ptr()); -} - -torch::stable::Tensor wrap_frame_pointer_to_tensor(AVFrame* frame) { - // Owning handle: the frame is freed when the handle tensor's refcount drops. - // ColorConverter borrows the frame during conversion (on CPU it does not free - // it), so the handle stays the sole owner and there is no leak even if a - // frame is never converted. (GPU conversion would consume the frame; GPU is - // not exposed through these ops yet.) - auto deleter = [frame](void*) { - AVFrame* f = frame; - av_frame_free(&f); - }; - int64_t sizes[] = {1}; - int64_t strides[] = {1}; - return torch::stable::from_blob( - frame, - {sizes, 1}, - {strides, 1}, - StableDevice(kStableCPU), - kStableInt64, - deleter); -} - -AVFrame* unwrap_tensor_to_frame(torch::stable::Tensor& tensor) { - STD_TORCH_CHECK(tensor.is_contiguous(), "frame handle must be contiguous"); - return static_cast(tensor.mutable_data_ptr()); -} - torch::stable::Tensor wrap_multi_stream_encoder_pointer_to_tensor( std::unique_ptr unique_encoder) { MultiStreamEncoder* encoder = unique_encoder.release(); @@ -850,11 +804,11 @@ using OpsPacketOutput = std::tuple; OpsPacketOutput _blocks_demuxer_next_packet(torch::stable::Tensor& demuxer) { Demuxer* demuxer_ptr = unwrap_tensor_to_pointer(demuxer); - AVPacket* packet = demuxer_ptr->next_packet(); + UniqueAVPacket packet = demuxer_ptr->next_packet(); if (packet == nullptr) { return std::make_tuple(torch::stable::full({1}, 0, kStableInt64), true); } - return std::make_tuple(wrap_packet_pointer_to_tensor(packet), false); + return std::make_tuple(wrap_pointer_to_tensor(std::move(packet)), false); } torch::stable::Tensor _blocks_create_packet_decoder( @@ -877,7 +831,7 @@ int64_t _blocks_packet_decoder_send_packet( torch::stable::Tensor& decoder, torch::stable::Tensor& packet) { PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer(decoder); - AVPacket* raw_packet = unwrap_tensor_to_packet(packet); + AVPacket* raw_packet = unwrap_tensor_to_pointer(packet); return static_cast(decoder_ptr->send_packet(raw_packet)); } @@ -908,11 +862,10 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame( 0.0); } AVRational time_base = decoder_ptr->time_base(); - double pts_seconds = pts_to_seconds(get_pts_or_dts(av_frame), time_base); - double duration_seconds = pts_to_seconds(get_duration(av_frame), time_base); - AVFrame* raw_frame = av_frame.release(); + double pts_seconds = pts_to_seconds(get_pts_or_dts(*av_frame), time_base); + double duration_seconds = pts_to_seconds(get_duration(*av_frame), time_base); return std::make_tuple( - wrap_frame_pointer_to_tensor(raw_frame), + wrap_pointer_to_tensor(std::move(av_frame)), static_cast(0), pts_seconds, duration_seconds); @@ -932,13 +885,7 @@ torch::stable::Tensor _blocks_convert_frame( torch::stable::Tensor& frame) { ColorConverter* converter_ptr = unwrap_tensor_to_pointer(converter); - AVFrame* raw_frame = unwrap_tensor_to_frame(frame); - // Borrow the frame for conversion, then release() so the handle keeps - // ownership and frees it when its tensor is dropped (CPU path). - UniqueAVFrame borrowed(raw_frame); - torch::stable::Tensor data = converter_ptr->convert(borrowed); - borrowed.release(); - return data; + return converter_ptr->convert(*unwrap_tensor_to_pointer(frame)); } // For testing only. We need to implement this operation as a core library diff --git a/test/third-party-interface/ThirdPartyInterfaceTest.cpp b/test/third-party-interface/ThirdPartyInterfaceTest.cpp index 8960ac7b8..a8c341f46 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 AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor = std::nullopt) override {} From b10d36c875d059f2bf417869a0b039485f48d750 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 15:07:14 +0100 Subject: [PATCH 03/14] Keep the AVFrame observers const on FFmpeg 4 and 5 Two things only show up when building against the older FFmpeg headers. get_num_channels() was patching av_frame.channel_layout when FFmpeg 4 left it unset, so it was a mutator wearing an observer's name, and the layout fix-up was a side effect that swresample setup silently relied on. Pull the fix-up into get_channel_layout() and call that from the two places that actually need a layout; get_num_channels() just counts. swr_alloc_set_opts2() only became const-correct in FFmpeg 6 (libswresample 4.12). Cast for the older headers, which don't modify the layout either. --- src/torchcodec/_core/FFMPEGCommon.cpp | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/torchcodec/_core/FFMPEGCommon.cpp b/src/torchcodec/_core/FFMPEGCommon.cpp index dc5b6e8d0..52ae31be0 100644 --- a/src/torchcodec/_core/FFMPEGCommon.cpp +++ b/src/torchcodec/_core/FFMPEGCommon.cpp @@ -148,17 +148,27 @@ const AVSampleFormat* get_supported_output_sample_formats( return supported_sample_formats; } +#if !( \ + LIBAVFILTER_VERSION_MAJOR > 8 || \ + (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44)) +// FFmpeg 4 leaves channel_layout unset (0) on some decoded frames even though +// .channels is correct. Everything we feed to swresample needs a real layout, +// so fall back to the default layout for that channel count. +int64_t get_channel_layout(const AVFrame& av_frame) { + if (av_frame.channel_layout != 0 || av_frame.channels <= 0) { + return static_cast(av_frame.channel_layout); + } + return av_get_default_channel_layout(av_frame.channels); +} +#endif + int get_num_channels(const AVFrame& av_frame) { #if LIBAVFILTER_VERSION_MAJOR > 8 || \ (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44) return av_frame.ch_layout.nb_channels; #else int num_channels = av_get_channel_layout_nb_channels(av_frame.channel_layout); - // Handle FFmpeg 4 bug where channel_layout and num_channels are 0 or unset - // Set values based on av_frame.channels which appears to be correct - // to allow successful initialization of SwrContext if (num_channels == 0 && av_frame.channels > 0) { - av_frame.channel_layout = av_get_default_channel_layout(av_frame.channels); num_channels = av_frame.channels; } return num_channels; @@ -316,7 +326,7 @@ int64_t get_output_channel_layout( const AVFrame& src_av_frame) { int64_t out_layout; if (out_num_channels == get_num_channels(src_av_frame)) { - out_layout = src_av_frame.channel_layout; + out_layout = get_channel_layout(src_av_frame); } else { out_layout = av_get_default_channel_layout(out_num_channels); } @@ -390,7 +400,10 @@ SwrContext* create_swr_context( &out_layout, out_sample_format, out_sample_rate, - &src_av_frame.ch_layout, + // swr_alloc_set_opts2() only became const-correct in FFmpeg 6 + // (libswresample 4.12): before that it asks for a non-const layout that + // it doesn't modify. + const_cast(&src_av_frame.ch_layout), src_sample_format, src_sample_rate, 0, @@ -408,7 +421,7 @@ SwrContext* create_swr_context( out_layout, out_sample_format, out_sample_rate, - src_av_frame.channel_layout, + get_channel_layout(src_av_frame), src_sample_format, src_sample_rate, 0, From 5ccac346199f3b27b149f1e77d1eb9a34cb079cc Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 15:22:55 +0100 Subject: [PATCH 04/14] rework FFmpeg macros --- src/torchcodec/_core/FFMPEGCommon.cpp | 45 ++++++++++++--------------- src/torchcodec/_core/FFMPEGCommon.h | 27 ++++++++++++++++ 2 files changed, 47 insertions(+), 25 deletions(-) diff --git a/src/torchcodec/_core/FFMPEGCommon.cpp b/src/torchcodec/_core/FFMPEGCommon.cpp index 52ae31be0..3726131eb 100644 --- a/src/torchcodec/_core/FFMPEGCommon.cpp +++ b/src/torchcodec/_core/FFMPEGCommon.cpp @@ -55,10 +55,10 @@ std::string get_ffmpeg_error_string_from_error_code(int error_code) { } int64_t get_duration(const AVFrame& av_frame) { -#if LIBAVUTIL_VERSION_MAJOR < 58 - return av_frame.pkt_duration; -#else +#if FFMPEG_HAS_FRAME_DURATION return av_frame.duration; +#else + return av_frame.pkt_duration; #endif } @@ -76,16 +76,16 @@ int64_t get_pts_or_dts(const AVFrame& av_frame) { } void set_duration(AVFrame& av_frame, int64_t duration) { -#if LIBAVUTIL_VERSION_MAJOR < 58 - av_frame.pkt_duration = duration; -#else +#if FFMPEG_HAS_FRAME_DURATION av_frame.duration = duration; +#else + av_frame.pkt_duration = duration; #endif } const int* get_supported_sample_rates(const AVCodec& av_codec) { const int* supported_sample_rates = nullptr; -#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1 +#if FFMPEG_HAS_SUPPORTED_CONFIG int num_sample_rates = 0; int ret = avcodec_get_supported_config( nullptr, @@ -106,7 +106,7 @@ const int* get_supported_sample_rates(const AVCodec& av_codec) { const AVPixelFormat* get_supported_pixel_formats(const AVCodec& av_codec) { const AVPixelFormat* supported_pixel_formats = nullptr; -#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1 +#if FFMPEG_HAS_SUPPORTED_CONFIG int num_pixel_formats = 0; int ret = avcodec_get_supported_config( nullptr, @@ -128,7 +128,7 @@ const AVPixelFormat* get_supported_pixel_formats(const AVCodec& av_codec) { const AVSampleFormat* get_supported_output_sample_formats( const AVCodec& av_codec) { const AVSampleFormat* supported_sample_formats = nullptr; -#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1 +#if FFMPEG_HAS_SUPPORTED_CONFIG int num_sample_formats = 0; int ret = avcodec_get_supported_config( nullptr, @@ -148,9 +148,7 @@ const AVSampleFormat* get_supported_output_sample_formats( return supported_sample_formats; } -#if !( \ - LIBAVFILTER_VERSION_MAJOR > 8 || \ - (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44)) +#if !FFMPEG_HAS_CH_LAYOUT // FFmpeg 4 leaves channel_layout unset (0) on some decoded frames even though // .channels is correct. Everything we feed to swresample needs a real layout, // so fall back to the default layout for that channel count. @@ -163,8 +161,7 @@ int64_t get_channel_layout(const AVFrame& av_frame) { #endif int get_num_channels(const AVFrame& av_frame) { -#if LIBAVFILTER_VERSION_MAJOR > 8 || \ - (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44) +#if FFMPEG_HAS_CH_LAYOUT return av_frame.ch_layout.nb_channels; #else int num_channels = av_get_channel_layout_nb_channels(av_frame.channel_layout); @@ -176,8 +173,7 @@ int get_num_channels(const AVFrame& av_frame) { } int get_num_channels(const SharedAVCodecContext& av_codec_context) { -#if LIBAVFILTER_VERSION_MAJOR > 8 || \ - (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44) +#if FFMPEG_HAS_CH_LAYOUT return av_codec_context->ch_layout.nb_channels; #else return av_codec_context->channels; @@ -186,8 +182,7 @@ int get_num_channels(const SharedAVCodecContext& av_codec_context) { int get_num_channels(const AVCodecParameters* codecpar) { STD_TORCH_CHECK(codecpar != nullptr, "codecpar is null"); -#if LIBAVFILTER_VERSION_MAJOR > 8 || \ - (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44) +#if FFMPEG_HAS_CH_LAYOUT return codecpar->ch_layout.nb_channels; #else return codecpar->channels; @@ -197,7 +192,7 @@ int get_num_channels(const AVCodecParameters* codecpar) { void set_default_channel_layout( UniqueAVCodecContext& av_codec_context, int num_channels) { -#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4 +#if FFMPEG_HAS_CH_LAYOUT AVChannelLayout channel_layout; av_channel_layout_default(&channel_layout, num_channels); av_codec_context->ch_layout = channel_layout; @@ -209,7 +204,7 @@ void set_default_channel_layout( } void set_default_channel_layout(AVFrame& av_frame, int num_channels) { -#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4 +#if FFMPEG_HAS_CH_LAYOUT AVChannelLayout channel_layout; av_channel_layout_default(&channel_layout, num_channels); av_frame.ch_layout = channel_layout; @@ -221,7 +216,7 @@ void set_default_channel_layout(AVFrame& av_frame, int num_channels) { } void validate_num_channels(const AVCodec& av_codec, int num_channels) { -#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1 +#if FFMPEG_HAS_SUPPORTED_CONFIG std::stringstream supported_num_channels; const AVChannelLayout* supported_layouts = nullptr; int num_layouts = 0; @@ -246,7 +241,7 @@ void validate_num_channels(const AVCodec& av_codec, int num_channels) { return; } } -#elif LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4 +#elif FFMPEG_HAS_CH_LAYOUT if (av_codec.ch_layouts == nullptr) { // If we can't validate, we must assume it'll be fine. If not, FFmpeg will // eventually raise. @@ -301,7 +296,7 @@ void validate_num_channels(const AVCodec& av_codec, int num_channels) { } namespace { -#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4 +#if FFMPEG_HAS_CH_LAYOUT // Returns: // - the src_av_frame's channel layout if src_av_frame has out_num_channels @@ -341,7 +336,7 @@ void set_channel_layout( AVFrame& dst_av_frame, const AVFrame& src_av_frame, int out_num_channels) { -#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4 +#if FFMPEG_HAS_CH_LAYOUT AVChannelLayout out_layout = get_output_channel_layout(out_num_channels, src_av_frame); auto status = av_channel_layout_copy(&dst_av_frame.ch_layout, &out_layout); @@ -392,7 +387,7 @@ SwrContext* create_swr_context( int out_num_channels) { SwrContext* swr_context = nullptr; int status = AVSUCCESS; -#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4 +#if FFMPEG_HAS_CH_LAYOUT AVChannelLayout out_layout = get_output_channel_layout(out_num_channels, src_av_frame); status = swr_alloc_set_opts2( diff --git a/src/torchcodec/_core/FFMPEGCommon.h b/src/torchcodec/_core/FFMPEGCommon.h index e6d57eadf..86381bb9f 100644 --- a/src/torchcodec/_core/FFMPEGCommon.h +++ b/src/torchcodec/_core/FFMPEGCommon.h @@ -31,6 +31,33 @@ extern "C" { #include } +// FFmpeg 5.1 replaced the .channels + .channel_layout pair on AVFrame and +// AVCodecContext with a single AVChannelLayout .ch_layout, and added the +// av_channel_layout_* / swr_alloc_set_opts2() APIs that go with it. +// libavutil 57.24 is the real marker, but libavfilter 8.44 is the equivalent +// and is what this codebase has always tested against. +#if LIBAVFILTER_VERSION_MAJOR > 8 || \ + (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44) +#define FFMPEG_HAS_CH_LAYOUT 1 +#else +#define FFMPEG_HAS_CH_LAYOUT 0 +#endif + +// FFmpeg 7.1 added avcodec_get_supported_config(), replacing the codec's +// pix_fmts / sample_fmts / supported_samplerates / ch_layouts arrays. +#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) +#define FFMPEG_HAS_SUPPORTED_CONFIG 1 +#else +#define FFMPEG_HAS_SUPPORTED_CONFIG 0 +#endif + +// FFmpeg 6 renamed AVFrame.pkt_duration to AVFrame.duration. +#if LIBAVUTIL_VERSION_MAJOR < 58 +#define FFMPEG_HAS_FRAME_DURATION 0 +#else +#define FFMPEG_HAS_FRAME_DURATION 1 +#endif + namespace facebook::torchcodec { // FFMPEG uses special delete functions for some structures. These template From 63f2aeda9e3adc4fa9f7e56559791687fd75ff66 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 15:45:19 +0100 Subject: [PATCH 05/14] TODOs --- src/torchcodec/_core/BetaCudaDeviceInterface.cpp | 2 -- src/torchcodec/_core/Demuxer.cpp | 2 ++ src/torchcodec/_core/custom_ops.cpp | 5 ++--- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 82ece2b7b..f6a8f2432 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 we only observe it. UniqueAVFrame transferred_frame; if (cpu_fallback_) { AVPixelFormat target_pix_fmt = (output_dtype_ == OutputDtype::FLOAT32) diff --git a/src/torchcodec/_core/Demuxer.cpp b/src/torchcodec/_core/Demuxer.cpp index ba46d6450..333a2306b 100644 --- a/src/torchcodec/_core/Demuxer.cpp +++ b/src/torchcodec/_core/Demuxer.cpp @@ -69,6 +69,8 @@ Demuxer::Demuxer( } UniqueAVPacket Demuxer::next_packet() { + // TODO_API_BREAKDOWN: Not a fan of the ReferenceAVPacket / AutoAVPacket / + // UniqueAVPacket dance here. Can we simplify? ReferenceAVPacket packet(auto_packet_); int status = read_next_packet(format_context_.get(), active_stream_index_, packet); diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp index 73351c99c..78037e4c8 100644 --- a/src/torchcodec/_core/custom_ops.cpp +++ b/src/torchcodec/_core/custom_ops.cpp @@ -185,9 +185,7 @@ SingleStreamDecoder* unwrap_tensor_to_get_decoder( // (Demuxer / PacketDecoder / ColorConverter / AVPacket / AVFrame). Same trick // as wrap_decoder_pointer_to_tensor: the tensor's data pointer IS the raw // pointer, and the tensor is the owner. Taking the unique_ptr by value is what -// makes the ownership transfer explicit at the call site; the unique_ptr's own -// deleter is what frees the object, so FFmpeg types work here as long as they -// arrive in their UniqueAVXxx alias. +// makes the ownership transfer explicit at the call site. template torch::stable::Tensor wrap_pointer_to_tensor(std::unique_ptr ptr) { D object_deleter = ptr.get_deleter(); @@ -832,6 +830,7 @@ int64_t _blocks_packet_decoder_send_packet( torch::stable::Tensor& packet) { PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer(decoder); AVPacket* raw_packet = unwrap_tensor_to_pointer(packet); + // TODO_API_BREAKDOWN: Do we really need this to be a raw AVPacket*? return static_cast(decoder_ptr->send_packet(raw_packet)); } From 5152a550670bbfc2cfd8bb45795e3388f5b0587d Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 16:40:41 +0100 Subject: [PATCH 06/14] Proper destructor --- .../_core/BetaCudaDeviceInterface.cpp | 29 +++++++++---------- .../_core/BetaCudaDeviceInterface.h | 7 ++--- 2 files changed, 16 insertions(+), 20 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 47c8e1aad..b58f2e2d0 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -270,6 +270,12 @@ void cuda_buffer_free_callback(void* opaque, [[maybe_unused]] uint8_t* data) { cudaFree(opaque); } +void standalone_frame_free_callback( + [[maybe_unused]] void* opaque, + uint8_t* data) { + delete reinterpret_cast(data); +} + } // namespace BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device) @@ -290,7 +296,7 @@ void BetaCudaDeviceInterface::initialize_video( const std::vector>& transforms, const std::optional& resized_output_dims) { // TODO_API_BREAKDOWN ewwwww - if (!av_stream){ + if (!av_stream) { return; } STD_TORCH_CHECK(av_stream != nullptr, "AVStream cannot be null"); @@ -809,14 +815,6 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( return av_frame; } -void nvdec_info_free_callback( - [[maybe_unused]] void* opaque, - [[maybe_unused]] uint8_t* data) { - printf("Freeing standalone frame attached data\n"); - fflush(stdout); - // delete reinterpret_cast(data); -} - void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { // TODO_API_BREAKDOWN_CUDA: stongly assumes NVDEC frame, we might want to // account for CPU fallback frames as well @@ -834,8 +832,6 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { int64_t pitch = static_cast(av_frame->linesize[0]); int64_t num_bytes = pitch * even_height * 3 / 2; - // TODO_API_BREAKDOWN_CUDA: How the hell do we know that the underlying - // storage isn't freed at the end of this scope?? auto storage = torch::stable::empty({num_bytes}, kStableUInt8, std::nullopt, device_); @@ -865,10 +861,12 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { auto attached_data = new StandAloneFrameAttachedData(); // TODO_API_BREAKDOWN_CUDA: We don't *really* need to std::move it I guess? attached_data->storage = std::move(storage); - - // av_frame->opaque_ref = av_buffer_create(reinterpret_cast(attached_data), sizeof(StandAloneFrameAttachedData), nvdec_info_free_callback, nullptr, 0); - // Intentionally leak storage, just todebug. - av_frame->opaque_ref = av_buffer_create(reinterpret_cast(attached_data), sizeof(StandAloneFrameAttachedData), nullptr, nullptr, 0); + av_frame->opaque_ref = av_buffer_create( + reinterpret_cast(attached_data), + sizeof(StandAloneFrameAttachedData), + standalone_frame_free_callback, + nullptr, + 0); } void BetaCudaDeviceInterface::flush() { @@ -1039,7 +1037,6 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( FrameOutput& frame_output, std::optional pre_allocated_output_tensor) { if (cpu_fallback_) { - // When the CPU fallback happens, we'll try to run the color-conversion on // GPU by sending those CPU frames to the GPU as NV12 or P016 (See // transferCpuFrameToGpu() below). However, it's not always diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index 55d50c2c8..9d70394ca 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.h @@ -33,9 +33,9 @@ #include "nvcuvid_include/nvcuvid.h" namespace facebook::torchcodec { - struct StandAloneFrameAttachedData{ - torch::stable::Tensor storage; - }; +struct StandAloneFrameAttachedData { + torch::stable::Tensor storage; +}; class BetaCudaDeviceInterface : public DeviceInterface { public: @@ -94,7 +94,6 @@ class BetaCudaDeviceInterface : public DeviceInterface { unsigned int pitch, const CUVIDPARSERDISPINFO& disp_info); - void make_frame_standalone(UniqueAVFrame& av_frame) override; UniqueAVFrame transfer_cpu_frame_to_gpu( From 4f4bd45caeb1c0c9511d91d706190e5ef5f9cafc Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 17:14:11 +0100 Subject: [PATCH 07/14] Add tests, support CPU fallback - I think --- .../_core/BetaCudaDeviceInterface.cpp | 23 ++++- test/test_decoders.py | 90 ++++++++++++------- 2 files changed, 78 insertions(+), 35 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 6f038016d..1a0ef0aa0 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -815,7 +815,16 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( return av_frame; } +// TODO_API_BREAKDOWN_CUDA: Does this even nede to be a method? Maybe it can be +// a function that just lives in the PacketDecoder so we don't need to expose +// another API to the DeviceInterface? void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { + + if (!(av_frame->format == AV_PIX_FMT_P016LE || av_frame->format == AV_PIX_FMT_NV12)) { + // The CPU frames are already standalone, so we don't need to do anything. + return; + } + // TODO_API_BREAKDOWN_CUDA: stongly assumes NVDEC frame, we might want to // account for CPU fallback frames as well @@ -834,6 +843,8 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { auto storage = torch::stable::empty({num_bytes}, kStableUInt8, std::nullopt, device_); + printf("MAKING FRAME STANDLONE\n"); + fflush(stdout); // TODO_API_BREAKDOWN_CUDA: Sync with the nvdec stream before copying? cudaStream_t stream = get_current_cuda_stream(device_.index()); @@ -1036,7 +1047,10 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( const AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor) { - if (cpu_fallback_) { + + bool cpu_fallback = av_frame.format != AV_PIX_FMT_NV12 && av_frame.format != AV_PIX_FMT_P016LE; + + if (cpu_fallback) { // When the CPU fallback happens, we'll try to run the color-conversion on // GPU by sending those CPU frames to the GPU as NV12 or P016 (See // transferCpuFrameToGpu() below). However, it's not always @@ -1048,6 +1062,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( av_pix_fmt_desc_get(static_cast(av_frame.format)); bool is444 = desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0; if (is444) { + // TODO_API_BREAKDOWN we need to handle this FrameOutput cpu_frame_output; cpu_fallback_->convert_av_frame_to_frame_output( av_frame, cpu_frame_output); @@ -1070,13 +1085,13 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( FrameDims original_dims(av_frame.height, av_frame.width); UniqueAVFrame transferred_frame; - if (cpu_fallback_) { + if (cpu_fallback) { AVPixelFormat target_pix_fmt = (output_dtype_ == OutputDtype::FLOAT32) ? AV_PIX_FMT_P016LE : AV_PIX_FMT_NV12; transferred_frame = transfer_cpu_frame_to_gpu(av_frame, target_pix_fmt); } - const AVFrame& gpu_frame = cpu_fallback_ ? *transferred_frame : av_frame; + const AVFrame& gpu_frame = cpu_fallback ? *transferred_frame : av_frame; STD_TORCH_CHECK( gpu_frame.format == AV_PIX_FMT_NV12 || @@ -1090,7 +1105,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( bool is_p016 = (gpu_frame.format == AV_PIX_FMT_P016LE); int bit_depth = 8; if (is_p016) { - bit_depth = cpu_fallback_ + bit_depth = cpu_fallback ? codec_context_->bits_per_raw_sample : static_cast(video_format_.bit_depth_luma_minus8) + 8; } diff --git a/test/test_decoders.py b/test/test_decoders.py index 70a2bbf95..9336bc615 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -3275,14 +3275,20 @@ def test_multiple_calls_with_backward_seeks(self): assert wav_samples.pts_seconds == audio_samples.pts_seconds +# The blocks support CPU and the NVDEC CUDA backend, and nothing else. In +# particular they never use the "ffmpeg" CUDA backend, so this is deliberately +# not all_supported_devices(). +def _block_devices(): + return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda)) + + +@pytest.mark.parametrize("device", _block_devices()) class TestBlocks: - def test_block_output_types(self): + def test_block_output_types(self, device): # Demuxer yields Packets, PacketDecoder yields DecodedFrames, and # ColorConverter yields Frames with the expected shape/dtype. - demuxer = Demuxer(NASA_VIDEO.path) - decoder = PacketDecoder(demuxer) - converter = ColorConverter() + demuxer, decoder, converter = self._make_blocks(NASA_VIDEO.path, device) num_packets = 0 for packet in demuxer: @@ -3295,6 +3301,7 @@ def test_block_output_types(self): assert frame.data.ndim == 3 # CHW assert frame.data.shape[0] == 3 # channels first assert frame.data.dtype == torch.uint8 + assert frame.data.device.type == device assert frame.duration_seconds >= 0 assert num_packets > 0 @@ -3351,42 +3358,42 @@ def drain(): return drain() - def _decoded_frames(self, path): - # demux + decode, as a single generator of DecodedFrames (pts order). + @staticmethod + def _make_blocks(path, device): demuxer = Demuxer(path) - decoder = PacketDecoder(demuxer) + decoder = PacketDecoder(demuxer, device=device) + # The converter is standalone on every device: it is not bound to the + # decoder that produced the frames. + converter = ColorConverter(device=device) + return demuxer, decoder, converter + + def _decoded_frames(self, path, device): + # demux + decode, as a single generator of DecodedFrames (pts order). + demuxer, decoder, _ = self._make_blocks(path, device) return self._decode(decoder, self._demux(demuxer)) - def _decode_sequential(self, path): + def _decode_sequential(self, path, device): # demux -> decode -> color-convert, all on the calling thread. - demuxer = Demuxer(path) - decoder = PacketDecoder(demuxer) - converter = ColorConverter() + demuxer, decoder, converter = self._make_blocks(path, device) return list( self._convert(converter, self._decode(decoder, self._demux(demuxer))) ) - def _decode_prefetch_frames(self, path): + def _decode_prefetch_frames(self, path, device): # [demux + decode] on one thread || [color-convert] on another. - demuxer = Demuxer(path) - decoder = PacketDecoder(demuxer) - converter = ColorConverter() + demuxer, decoder, converter = self._make_blocks(path, device) frames = self.prefetch(self._decode(decoder, self._demux(demuxer))) return list(self._convert(converter, frames)) - def _decode_prefetch_packets(self, path): + def _decode_prefetch_packets(self, path, device): # [demux] on one thread || [decode + color-convert] on another. - demuxer = Demuxer(path) - decoder = PacketDecoder(demuxer) - converter = ColorConverter() + demuxer, decoder, converter = self._make_blocks(path, device) packets = self.prefetch(self._demux(demuxer)) return list(self._convert(converter, self._decode(decoder, packets))) - def _decode_prefetch_packets_and_frames(self, path): + def _decode_prefetch_packets_and_frames(self, path, device): # [demux] || [decode] || [color-convert], each on its own thread. - demuxer = Demuxer(path) - decoder = PacketDecoder(demuxer) - converter = ColorConverter() + demuxer, decoder, converter = self._make_blocks(path, device) packets = self.prefetch(self._demux(demuxer)) frames = self.prefetch(self._decode(decoder, packets)) return list(self._convert(converter, frames)) @@ -3402,7 +3409,18 @@ def _to_frame_batch(self, frames): ), ) - @pytest.mark.parametrize("video", (NASA_VIDEO, BT709_FULL_RANGE)) + @pytest.mark.parametrize( + "video", + ( + NASA_VIDEO, + BT709_FULL_RANGE, + NASA_VIDEO_HDR, TEST_SRC_2_720P_HDR, TEST_SRC_2_12BIT_HDR, + # NVDEC can't decode this one (too small), so on CUDA this covers + # the CPU-fallback path: the decoder hands out CPU frames and the + # converter has to notice and upload them itself. + H265_VIDEO, + ), + ) @pytest.mark.parametrize( "decode_method", ( @@ -3413,10 +3431,11 @@ def _to_frame_batch(self, frames): ), ids=lambda f: f.__name__.removeprefix("_decode_"), ) - def test_matches_video_decoder(self, video, decode_method): - got = self._to_frame_batch(decode_method(self, video.path)) - ref = VideoDecoder(video.path).get_all_frames() + def test_matches_video_decoder(self, video, decode_method, device): + got = self._to_frame_batch(decode_method(self, video.path, device)) + ref = VideoDecoder(video.path, device=device).get_all_frames() + assert got.data.device.type == device assert got.data.shape == ref.data.shape torch.testing.assert_close(got.data, ref.data, atol=0, rtol=0) torch.testing.assert_close(got.pts_seconds, ref.pts_seconds, atol=0, rtol=0) @@ -3424,13 +3443,13 @@ def test_matches_video_decoder(self, video, decode_method): got.duration_seconds, ref.duration_seconds, atol=0, rtol=0 ) - def test_color_converter_reused_across_videos(self): + def test_color_converter_reused_across_videos(self, device): # A single unbound ColorConverter must correctly convert frames from two # different videos - here interleaved frame-by-frame, so the converter # switches input resolution/format on every call. - converter = ColorConverter() + converter = ColorConverter(device=device) videos = [NASA_VIDEO, BT709_FULL_RANGE] - generators = [self._decoded_frames(v.path) for v in videos] + generators = [self._decoded_frames(v.path, device) for v in videos] outputs = [[] for _ in videos] done = [False] * len(videos) @@ -3446,10 +3465,19 @@ def test_color_converter_reused_across_videos(self): for video, frames in zip(videos, outputs): got = self._to_frame_batch(frames) - ref = VideoDecoder(video.path).get_all_frames() + ref = VideoDecoder(video.path, device=device).get_all_frames() assert got.data.shape == ref.data.shape torch.testing.assert_close(got.data, ref.data, atol=0, rtol=0) + def test_set_cuda_backend_is_a_noop(self, device): + # The blocks always use the NVDEC CUDA backend. Asking for the "ffmpeg" + # one changes nothing, rather than silently producing something else. + # TODO_API_BREAKDOWN: let's just error? + with set_cuda_backend("ffmpeg"): + got = self._to_frame_batch(self._decode_sequential(NASA_VIDEO.path, device)) + ref = self._to_frame_batch(self._decode_sequential(NASA_VIDEO.path, device)) + torch.testing.assert_close(got.data, ref.data, atol=0, rtol=0) + # Small helpers to avoid having to always specify the same skip marks and decode_fn def _jpeg_param(*values): From 13dc67bebf65de3febefb6e569bd3e2e54854de3 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 17:24:32 +0100 Subject: [PATCH 08/14] Ignore set_cuda_backend --- .../_core/BetaCudaDeviceInterface.cpp | 21 +++++++------------ src/torchcodec/_core/ColorConverter.cpp | 6 ++---- src/torchcodec/_core/ColorConverter.h | 3 +-- src/torchcodec/_core/DeviceInterface.h | 2 +- src/torchcodec/_core/PacketDecoder.cpp | 3 +-- src/torchcodec/_core/PacketDecoder.h | 1 - src/torchcodec/_core/custom_ops.cpp | 21 +++++++------------ test/test_decoders.py | 4 +++- 8 files changed, 24 insertions(+), 37 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 1a0ef0aa0..e1e17a236 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -595,7 +595,6 @@ int BetaCudaDeviceInterface::send_eof_packet() { int BetaCudaDeviceInterface::send_cuvid_packet( CUVIDSOURCEDATAPACKET& cuvid_packet) { CUresult result = cuvidParseVideoData(video_parser_, &cuvid_packet); - printf("cuvidParseVideoData returned %d\n", result); return result == CUDA_SUCCESS ? AVSUCCESS : AVERROR_EXTERNAL; } @@ -819,11 +818,11 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( // a function that just lives in the PacketDecoder so we don't need to expose // another API to the DeviceInterface? void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { - - if (!(av_frame->format == AV_PIX_FMT_P016LE || av_frame->format == AV_PIX_FMT_NV12)) { - // The CPU frames are already standalone, so we don't need to do anything. - return; - } + if (!(av_frame->format == AV_PIX_FMT_P016LE || + av_frame->format == AV_PIX_FMT_NV12)) { + // The CPU frames are already standalone, so we don't need to do anything. + return; + } // TODO_API_BREAKDOWN_CUDA: stongly assumes NVDEC frame, we might want to // account for CPU fallback frames as well @@ -843,8 +842,6 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { auto storage = torch::stable::empty({num_bytes}, kStableUInt8, std::nullopt, device_); - printf("MAKING FRAME STANDLONE\n"); - fflush(stdout); // TODO_API_BREAKDOWN_CUDA: Sync with the nvdec stream before copying? cudaStream_t stream = get_current_cuda_stream(device_.index()); @@ -866,8 +863,6 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { auto y_plane = static_cast(storage.mutable_data_ptr()); av_frame->data[0] = y_plane; av_frame->data[1] = y_plane + (pitch * even_height); - printf("Returning copied frame\n"); - fflush(stdout); auto attached_data = new StandAloneFrameAttachedData(); // TODO_API_BREAKDOWN_CUDA: We don't *really* need to std::move it I guess? @@ -1047,9 +1042,9 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( const AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor) { - - bool cpu_fallback = av_frame.format != AV_PIX_FMT_NV12 && av_frame.format != AV_PIX_FMT_P016LE; - + bool cpu_fallback = av_frame.format != AV_PIX_FMT_NV12 && + av_frame.format != AV_PIX_FMT_P016LE; + if (cpu_fallback) { // When the CPU fallback happens, we'll try to run the color-conversion on // GPU by sending those CPU frames to the GPU as NV12 or P016 (See diff --git a/src/torchcodec/_core/ColorConverter.cpp b/src/torchcodec/_core/ColorConverter.cpp index eb2218179..262d1946c 100644 --- a/src/torchcodec/_core/ColorConverter.cpp +++ b/src/torchcodec/_core/ColorConverter.cpp @@ -15,10 +15,8 @@ namespace facebook::torchcodec { -ColorConverter::ColorConverter( - const StableDevice& device, - std::string_view device_variant) { - device_interface_ = create_device_interface(device, device_variant); +ColorConverter::ColorConverter(const StableDevice& device) { + device_interface_ = create_device_interface(device); STD_TORCH_CHECK( device_interface_ != nullptr, "Failed to create device interface. This should never happen, please report."); diff --git a/src/torchcodec/_core/ColorConverter.h b/src/torchcodec/_core/ColorConverter.h index d06f2d4a2..5c66d1f1d 100644 --- a/src/torchcodec/_core/ColorConverter.h +++ b/src/torchcodec/_core/ColorConverter.h @@ -18,8 +18,7 @@ namespace facebook::torchcodec { class FORCE_PUBLIC_VISIBILITY ColorConverter { public: explicit ColorConverter( - const StableDevice& device = StableDevice(kStableCPU), - std::string_view device_variant = "default"); + const StableDevice& device = StableDevice(kStableCPU)); torch::stable::Tensor convert(const AVFrame& av_frame); diff --git a/src/torchcodec/_core/DeviceInterface.h b/src/torchcodec/_core/DeviceInterface.h index 61d24759e..a37b292d7 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/src/torchcodec/_core/DeviceInterface.h @@ -203,7 +203,7 @@ TORCHCODEC_THIRD_PARTY_API bool register_device_interface( FORCE_PUBLIC_VISIBILITY void validate_device_interface( const std::string& device, - const std::string& variant); + const std::string& variant = "default"); TORCHCODEC_THIRD_PARTY_API std::unique_ptr create_device_interface( diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index 122322871..f05cee09e 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -58,9 +58,8 @@ const AVCodec* find_decoder( PacketDecoder::PacketDecoder( const Demuxer& demuxer, const StableDevice& device, - std::string_view device_variant, std::optional ffmpeg_thread_count) { - device_interface_ = create_device_interface(device, device_variant); + device_interface_ = create_device_interface(device); STD_TORCH_CHECK( device_interface_ != nullptr, "Failed to create device interface. This should never happen, please report."); diff --git a/src/torchcodec/_core/PacketDecoder.h b/src/torchcodec/_core/PacketDecoder.h index 929698ac6..a22307874 100644 --- a/src/torchcodec/_core/PacketDecoder.h +++ b/src/torchcodec/_core/PacketDecoder.h @@ -33,7 +33,6 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder { explicit PacketDecoder( const Demuxer& demuxer, const StableDevice& device = StableDevice(kStableCPU), - std::string_view device_variant = "default", std::optional ffmpeg_thread_count = std::nullopt); // Feed one packet to the decoder. Borrows `packet` (does not take ownership). diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp index f5c789b63..5b6d18f1a 100644 --- a/src/torchcodec/_core/custom_ops.cpp +++ b/src/torchcodec/_core/custom_ops.cpp @@ -77,14 +77,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) { "_blocks_create_demuxer(str filename, int? stream_index=None) -> Tensor"); m.def("_blocks_demuxer_next_packet(Tensor(a!) demuxer) -> (Tensor, bool)"); m.def( - "_blocks_create_packet_decoder(Tensor demuxer, *, int? num_threads=None, str device=\"cpu\", str device_variant=\"default\") -> Tensor"); + "_blocks_create_packet_decoder(Tensor demuxer, *, int? num_threads=None, str device=\"cpu\") -> Tensor"); m.def( "_blocks_packet_decoder_send_packet(Tensor(a!) decoder, Tensor packet) -> int"); m.def("_blocks_packet_decoder_send_eof(Tensor(a!) decoder) -> int"); m.def( "_blocks_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float)"); - m.def( - "_blocks_create_color_converter(str device=\"cpu\", str device_variant=\"default\") -> Tensor"); + m.def("_blocks_create_color_converter(str device=\"cpu\") -> Tensor"); m.def("_blocks_convert_frame(Tensor(a!) converter, Tensor frame) -> Tensor"); m.def("_get_key_frame_indices(Tensor(a!) decoder) -> Tensor"); m.def("get_json_metadata(Tensor(a!) decoder) -> str"); @@ -812,16 +811,15 @@ OpsPacketOutput _blocks_demuxer_next_packet(torch::stable::Tensor& demuxer) { torch::stable::Tensor _blocks_create_packet_decoder( torch::stable::Tensor& demuxer, std::optional num_threads, - std::string device, - std::string device_variant) { + std::string device) { Demuxer* demuxer_ptr = unwrap_tensor_to_pointer(demuxer); - validate_device_interface(device, device_variant); + validate_device_interface(device); std::optional thread_count; if (num_threads.has_value()) { thread_count = static_cast(num_threads.value()); } auto decoder = std::make_unique( - *demuxer_ptr, StableDevice(device), device_variant, thread_count); + *demuxer_ptr, StableDevice(device), thread_count); return wrap_pointer_to_tensor(std::move(decoder)); } @@ -870,12 +868,9 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame( duration_seconds); } -torch::stable::Tensor _blocks_create_color_converter( - std::string device, - std::string device_variant) { - validate_device_interface(device, device_variant); - auto converter = - std::make_unique(StableDevice(device), device_variant); +torch::stable::Tensor _blocks_create_color_converter(std::string device) { + validate_device_interface(device); + auto converter = std::make_unique(StableDevice(device)); return wrap_pointer_to_tensor(std::move(converter)); } diff --git a/test/test_decoders.py b/test/test_decoders.py index 9336bc615..a1c615866 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -3414,7 +3414,9 @@ def _to_frame_batch(self, frames): ( NASA_VIDEO, BT709_FULL_RANGE, - NASA_VIDEO_HDR, TEST_SRC_2_720P_HDR, TEST_SRC_2_12BIT_HDR, + NASA_VIDEO_HDR, + TEST_SRC_2_720P_HDR, + TEST_SRC_2_12BIT_HDR, # NVDEC can't decode this one (too small), so on CUDA this covers # the CPU-fallback path: the decoder hands out CPU frames and the # converter has to notice and upload them itself. From 086ce09fd78ed0c7d8761ab850edc181c9fb41c7 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 18:42:01 +0100 Subject: [PATCH 09/14] Hacky workaround --- test/test_decoders.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/test/test_decoders.py b/test/test_decoders.py index a1c615866..a66c92d2f 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -3381,8 +3381,20 @@ def _decode_sequential(self, path, device): def _decode_prefetch_frames(self, path, device): # [demux + decode] on one thread || [color-convert] on another. - demuxer, decoder, converter = self._make_blocks(path, device) - frames = self.prefetch(self._decode(decoder, self._demux(demuxer))) + demuxer = Demuxer(path) + converter = ColorConverter(device=device) + + def demux_and_decode(): + # Constructed here so the PacketDecoder (and thus the CUDA context it + # binds via its device interface) lives on the prefetch worker + # thread that actually runs the cuvid* calls. + # TODO_API_BREAKDOWN_CUDA: this is a temporary workaroun for a real + # issue - needs fixing. Things should work when the objects are + # constructed in a different thread than where they're consumed. + decoder = PacketDecoder(demuxer, device=device) + yield from self._decode(decoder, self._demux(demuxer)) + + frames = self.prefetch(demux_and_decode()) return list(self._convert(converter, frames)) def _decode_prefetch_packets(self, path, device): @@ -3393,9 +3405,19 @@ def _decode_prefetch_packets(self, path, device): def _decode_prefetch_packets_and_frames(self, path, device): # [demux] || [decode] || [color-convert], each on its own thread. - demuxer, decoder, converter = self._make_blocks(path, device) + demuxer = Demuxer(path) + converter = ColorConverter(device=device) packets = self.prefetch(self._demux(demuxer)) - frames = self.prefetch(self._decode(decoder, packets)) + + def decode(packets): + # Constructed here so the PacketDecoder (and thus the CUDA context it + # binds via its device interface) lives on the prefetch worker + # thread that actually runs the cuvid* calls. + # TODO_API_BREAKDOWN_CUDA: same TODO as ab + decoder = PacketDecoder(demuxer, device=device) + yield from self._decode(decoder, packets) + + frames = self.prefetch(decode(packets)) return list(self._convert(converter, frames)) def _to_frame_batch(self, frames): From 560bf6d0e47efbf0adcdcef61226d3a065de9d85 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 19:10:28 +0100 Subject: [PATCH 10/14] Add TODOs --- .../_core/BetaCudaDeviceInterface.cpp | 20 +++++++++++++++---- test/test_decoders.py | 2 +- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index e1e17a236..8b07ae3e3 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -824,9 +824,6 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { return; } - // TODO_API_BREAKDOWN_CUDA: stongly assumes NVDEC frame, we might want to - // account for CPU fallback frames as well - // Roughly, the number of bytes an NV12 image takes is: // num_bytes = len(Y) + len(UV) // = num_pixels + num_pixels / 2 @@ -843,7 +840,12 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { auto storage = torch::stable::empty({num_bytes}, kStableUInt8, std::nullopt, device_); - // TODO_API_BREAKDOWN_CUDA: Sync with the nvdec stream before copying? + // TODO_API_BREAKDOWN_CUDA: This copies on the current stream, but the frame + // was produced by NVDEC on a potentially different stream. If they differ we + // must wait on the producing stream before copying, otherwise we read the + // surface before NVDEC finished writing it. We also need to record the + // producing stream on the attached data below, so the ColorConverter - which + // may run on yet another thread/stream - can wait on it before converting. cudaStream_t stream = get_current_cuda_stream(device_.index()); cudaError_t err = cudaMemcpyAsync( storage.mutable_data_ptr(), @@ -1042,6 +1044,9 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( const AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor) { + + // TODO_API_BREAKDOWN_CUDA is that accurate and safe? Can there be a CPU NV12 + // frame in our code? Should we create a helper used in the make_standalone function too? bool cpu_fallback = av_frame.format != AV_PIX_FMT_NV12 && av_frame.format != AV_PIX_FMT_P016LE; @@ -1093,8 +1098,15 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( gpu_frame.format == AV_PIX_FMT_P016LE, "Expected NV12 or P016LE format frame"); + // TODO_API_BREAKDOWN_CUDA: In convert-only mode the frame was produced on + // another interface, possibly on another thread and hence another stream. + // Using the current stream is only correct when this same interface decoded + // the frame. We should read the producing stream off the frame (via the + // attached data set in make_frame_standalone) and wait on it here. cudaStream_t nvdec_stream = get_current_cuda_stream(device_.index()); + // TODO_API_BREAKDOWN: we don't suppor output_dtype so some of that is not + // execrcized. auto convert_frame = [&](std::optional pre_alloc) -> torch::stable::Tensor { bool is_p016 = (gpu_frame.format == AV_PIX_FMT_P016LE); diff --git a/test/test_decoders.py b/test/test_decoders.py index a66c92d2f..8e39c56f4 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -3413,7 +3413,7 @@ def decode(packets): # Constructed here so the PacketDecoder (and thus the CUDA context it # binds via its device interface) lives on the prefetch worker # thread that actually runs the cuvid* calls. - # TODO_API_BREAKDOWN_CUDA: same TODO as ab + # TODO_API_BREAKDOWN_CUDA: same TODO as above. decoder = PacketDecoder(demuxer, device=device) yield from self._decode(decoder, packets) From e04e070d06c13daf1972316f6f7d429d96b66a26 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 4 Aug 2026 23:34:42 +0100 Subject: [PATCH 11/14] Fix stream sync issue --- .../_core/BetaCudaDeviceInterface.cpp | 40 +++++++++++-------- .../_core/BetaCudaDeviceInterface.h | 1 + src/torchcodec/_core/color_conversion.cpp | 9 ++++- src/torchcodec/_core/color_conversion.h | 2 +- test/test_decoders.py | 5 ++- 5 files changed, 36 insertions(+), 21 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 8b07ae3e3..b51a922bc 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -698,6 +698,8 @@ int BetaCudaDeviceInterface::receive_frame(UniqueAVFrame& av_frame) { // color-converted (with a copy), or that's a frame that was discarded in // SingleStreamDecoder. Either way, the underlying output surface can be // safely re-used. + // TODO_API_BREAKDOWN: We should update this comment slightly to now account + // for the frame copy we do in make_frame_standalone() unmap_previous_frame(); CUresult result = cuvidMapVideoFrame( *decoder_.get(), @@ -840,19 +842,16 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { auto storage = torch::stable::empty({num_bytes}, kStableUInt8, std::nullopt, device_); - // TODO_API_BREAKDOWN_CUDA: This copies on the current stream, but the frame - // was produced by NVDEC on a potentially different stream. If they differ we - // must wait on the producing stream before copying, otherwise we read the - // surface before NVDEC finished writing it. We also need to record the - // producing stream on the attached data below, so the ColorConverter - which - // may run on yet another thread/stream - can wait on it before converting. - cudaStream_t stream = get_current_cuda_stream(device_.index()); + // TODO_API_BREAKDOWN_CUDA: I suspect we don't need to wait on the nvdec + // stream here, because we can only arrive here from a path where the frame + // has already been mapped so its data is available - worth double checking. + cudaStream_t current_stream = get_current_cuda_stream(device_.index()); cudaError_t err = cudaMemcpyAsync( storage.mutable_data_ptr(), av_frame->data[0], static_cast(num_bytes), cudaMemcpyDeviceToDevice, - stream); + current_stream); STD_TORCH_CHECK( err == cudaSuccess, "Failed to copy NVDEC surface: ", @@ -867,6 +866,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { av_frame->data[1] = y_plane + (pitch * even_height); auto attached_data = new StandAloneFrameAttachedData(); + attached_data->producer_stream = current_stream; // TODO_API_BREAKDOWN_CUDA: We don't *really* need to std::move it I guess? attached_data->storage = std::move(storage); av_frame->opaque_ref = av_buffer_create( @@ -1044,9 +1044,9 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( const AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor) { - // TODO_API_BREAKDOWN_CUDA is that accurate and safe? Can there be a CPU NV12 - // frame in our code? Should we create a helper used in the make_standalone function too? + // frame in our code? Should we create a helper used in the make_standalone + // function too? bool cpu_fallback = av_frame.format != AV_PIX_FMT_NV12 && av_frame.format != AV_PIX_FMT_P016LE; @@ -1098,12 +1098,18 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( gpu_frame.format == AV_PIX_FMT_P016LE, "Expected NV12 or P016LE format frame"); - // TODO_API_BREAKDOWN_CUDA: In convert-only mode the frame was produced on - // another interface, possibly on another thread and hence another stream. - // Using the current stream is only correct when this same interface decoded - // the frame. We should read the producing stream off the frame (via the - // attached data set in make_frame_standalone) and wait on it here. - cudaStream_t nvdec_stream = get_current_cuda_stream(device_.index()); + // TODO_API_BREAKDOWN: Cleanup how we get the attached data? Make it more + // robust? Should we couple it to a flag on the interface saying "I'm + // color-conversion only, I absolutely expect frames to be standalone"? + cudaStream_t producer_stream; + if (av_frame.opaque_ref != nullptr && + av_frame.opaque_ref->size == sizeof(StandAloneFrameAttachedData)) { + auto attached_data = reinterpret_cast( + av_frame.opaque_ref->data); + producer_stream = attached_data->producer_stream; + } else { + producer_stream = get_current_cuda_stream(device_.index()); + } // TODO_API_BREAKDOWN: we don't suppor output_dtype so some of that is not // execrcized. @@ -1119,7 +1125,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( return convert_yuv_frame_to_rgb( gpu_frame, device_, - nvdec_stream, + producer_stream, pre_alloc, original_dims, is_p016, diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index 9d70394ca..3ae019ec1 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.h @@ -34,6 +34,7 @@ namespace facebook::torchcodec { struct StandAloneFrameAttachedData { + cudaStream_t producer_stream = nullptr; torch::stable::Tensor storage; }; diff --git a/src/torchcodec/_core/color_conversion.cpp b/src/torchcodec/_core/color_conversion.cpp index 0efa6d699..98ffcc35f 100644 --- a/src/torchcodec/_core/color_conversion.cpp +++ b/src/torchcodec/_core/color_conversion.cpp @@ -191,7 +191,7 @@ void compute_rgb_to_yuv_matrix( torch::stable::Tensor convert_yuv_frame_to_rgb( const AVFrame& av_frame, const StableDevice& device, - cudaStream_t nvdec_stream, + cudaStream_t producer_stream, std::optional pre_allocated_output_tensor, const FrameDims& output_dims, bool is_p016, @@ -221,9 +221,14 @@ torch::stable::Tensor convert_yuv_frame_to_rgb( FrameDims(out_height, out_width), device, out_dtype); } + // TODO_API_BREAKDOWN: This may not be the semantic that we want: this will + // wait for all ongoin work on the producer stream to finish. But maybe the + // producer stream produced the frame data a long time ago, and lots of + // kernels have been launched on it already. We'd be waiting on those to + // finish even though the data we need is already available. cudaStream_t stream = get_current_cuda_stream(device.index()); sync_streams( - /*runningStream=*/nvdec_stream, /*waitingStream=*/stream); + /*runningStream=*/producer_stream, /*waitingStream=*/stream); maybe_update_color_matrix( cached_color_matrix, diff --git a/src/torchcodec/_core/color_conversion.h b/src/torchcodec/_core/color_conversion.h index ad3ae29ef..f2965acf8 100644 --- a/src/torchcodec/_core/color_conversion.h +++ b/src/torchcodec/_core/color_conversion.h @@ -86,7 +86,7 @@ void launch_p016_to_rgb16_kernel( torch::stable::Tensor convert_yuv_frame_to_rgb( const AVFrame& av_frame, const StableDevice& device, - cudaStream_t nvdec_stream, + cudaStream_t producer_stream, std::optional pre_allocated_output_tensor, const FrameDims& output_dims, bool is_p016, diff --git a/test/test_decoders.py b/test/test_decoders.py index 8e39c56f4..3cc7d0f1f 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -3282,9 +3282,9 @@ def _block_devices(): return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda)) -@pytest.mark.parametrize("device", _block_devices()) class TestBlocks: + @pytest.mark.parametrize("device", _block_devices()) def test_block_output_types(self, device): # Demuxer yields Packets, PacketDecoder yields DecodedFrames, and # ColorConverter yields Frames with the expected shape/dtype. @@ -3455,6 +3455,7 @@ def _to_frame_batch(self, frames): ), ids=lambda f: f.__name__.removeprefix("_decode_"), ) + @pytest.mark.parametrize("device", _block_devices()) def test_matches_video_decoder(self, video, decode_method, device): got = self._to_frame_batch(decode_method(self, video.path, device)) ref = VideoDecoder(video.path, device=device).get_all_frames() @@ -3467,6 +3468,7 @@ def test_matches_video_decoder(self, video, decode_method, device): got.duration_seconds, ref.duration_seconds, atol=0, rtol=0 ) + @pytest.mark.parametrize("device", _block_devices()) def test_color_converter_reused_across_videos(self, device): # A single unbound ColorConverter must correctly convert frames from two # different videos - here interleaved frame-by-frame, so the converter @@ -3493,6 +3495,7 @@ def test_color_converter_reused_across_videos(self, device): assert got.data.shape == ref.data.shape torch.testing.assert_close(got.data, ref.data, atol=0, rtol=0) + @pytest.mark.parametrize("device", _block_devices()) def test_set_cuda_backend_is_a_noop(self, device): # The blocks always use the NVDEC CUDA backend. Asking for the "ffmpeg" # one changes nothing, rather than silently producing something else. From d115fcc0892bfd2b748a65c2a194db1d6ab9d386 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 5 Aug 2026 14:33:09 +0100 Subject: [PATCH 12/14] Add Cuda Context guard --- .../_core/BetaCudaDeviceInterface.cpp | 74 ++++++++++++++----- src/torchcodec/_core/ColorConverter.cpp | 4 +- src/torchcodec/_core/CpuDeviceInterface.cpp | 2 +- src/torchcodec/_core/Demuxer.cpp | 2 +- src/torchcodec/_core/PacketDecoder.cpp | 4 +- src/torchcodec/_core/custom_ops.cpp | 2 +- .../decoders/_blocks/_color_converter.py | 12 +-- .../decoders/_blocks/_packet_decoder.py | 6 +- test/test_decoders.py | 33 ++------- 9 files changed, 80 insertions(+), 59 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index b51a922bc..c149acc66 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -266,6 +266,8 @@ std::optional get_nvdec_surface_format( // Callback for freeing CUDA memory associated with AVFrame see where it's used // for more details. +// TODO_API_BREAKDOWN P2: Should we align this with the other free callback +// below? Why did we use cudaMalloc? Can we just allocate with torch?? void cuda_buffer_free_callback(void* opaque, [[maybe_unused]] uint8_t* data) { cudaFree(opaque); } @@ -276,6 +278,36 @@ void standalone_frame_free_callback( delete reinterpret_cast(data); } +class CudaContextGuard { + // There's one CUDA context per process per device. But new threads aren't + // bound to a context. The binding often happens automatically when calling + // CUDA APIs (like cudaFree), but some APIs like the NVCUVID ones that we use + // here aren't automatically binding. + // So for a thread to be able to use NVCUVID APIs, it must have a context + // bound to it, and we have to enforce that binding manually. + // That's what this guard does: it calls cudaFree(nullptr), which is a common + // near-free way to force the CUDA runtime to bind the context for the current + // thread. And this call must happen within a device guard to make sure we're + // binding the context of the device this interface is using. + // We must call this guard in every public method of the interface that uses + // NVCUVID APIs, because these methods can, in theory, be called from any + // thread. + // Note that none of this was an issue before when our only entry-point was + // the SingleStreamDecoder: all the entry-points were called from the same + // thread. Now that we have split the APIs in different blocks (PacketDecoder, + // ColorConverter), each of these blocks can be on different threads - and + // importantly, they can be created in the main thread (where the context is + // bound by our call to initialize_cuda_context_with_pytorch()), but then used + // in a different thread that doesn't have the context. + public: + explicit CudaContextGuard(int device_index) : device_guard_(device_index) { + cudaFree(nullptr); + } + + private: + StableDeviceGuard device_guard_; +}; + } // namespace BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device) @@ -284,6 +316,9 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device) STD_TORCH_CHECK( device_.type() == kStableCUDA, "Unsupported device: must be CUDA"); + // Note: now that we have the CudaContextGuard, we might not need to do that + // anymore. The comment says we need pytorch to create the context - maybe + // that's true, but that's a very old comment now. initialize_cuda_context_with_pytorch(device_); nvcuvid_available_ = load_nvcuvid_library(); @@ -295,11 +330,12 @@ void BetaCudaDeviceInterface::initialize_video( const VideoStreamOptions& video_stream_options, const std::vector>& transforms, const std::optional& resized_output_dims) { - // TODO_API_BREAKDOWN ewwwww + // TODO_API_BREAKDOWN P0 if (!av_stream) { return; } STD_TORCH_CHECK(av_stream != nullptr, "AVStream cannot be null"); + CudaContextGuard context_guard(device_.index()); rotation_ = rotation_from_degrees(get_rotation_from_stream(av_stream)); output_dtype_ = video_stream_options.output_dtype; @@ -403,6 +439,7 @@ void BetaCudaDeviceInterface::send_seqhdr_packet() { } BetaCudaDeviceInterface::~BetaCudaDeviceInterface() { + CudaContextGuard context_guard(device_.index()); if (decoder_) { // DALI doesn't seem to do any particular cleanup of the decoder before // sending it to the cache, so we probably don't need to do anything either. @@ -554,6 +591,7 @@ int BetaCudaDeviceInterface::stream_property_change( // Moral equivalent of avcodec_send_packet(). Here, we pass the AVPacket down to // the NVCUVID parser. int BetaCudaDeviceInterface::send_packet(ReferenceAVPacket& packet) { + CudaContextGuard context_guard(device_.index()); if (cpu_fallback_) { return cpu_fallback_->send_packet(packet); } @@ -581,6 +619,7 @@ int BetaCudaDeviceInterface::send_packet(ReferenceAVPacket& packet) { } int BetaCudaDeviceInterface::send_eof_packet() { + CudaContextGuard context_guard(device_.index()); if (cpu_fallback_) { return cpu_fallback_->send_eof_packet(); } @@ -656,6 +695,7 @@ int BetaCudaDeviceInterface::frame_ready_in_display_order( // Moral equivalent of avcodec_receive_frame(). int BetaCudaDeviceInterface::receive_frame(UniqueAVFrame& av_frame) { + CudaContextGuard context_guard(device_.index()); if (cpu_fallback_) { return cpu_fallback_->receive_frame(av_frame); } @@ -698,8 +738,8 @@ int BetaCudaDeviceInterface::receive_frame(UniqueAVFrame& av_frame) { // color-converted (with a copy), or that's a frame that was discarded in // SingleStreamDecoder. Either way, the underlying output surface can be // safely re-used. - // TODO_API_BREAKDOWN: We should update this comment slightly to now account - // for the frame copy we do in make_frame_standalone() + // TODO_API_BREAKDOWN P1: We should update this comment slightly to now + // account for the frame copy we do in make_frame_standalone() unmap_previous_frame(); CUresult result = cuvidMapVideoFrame( *decoder_.get(), @@ -807,7 +847,7 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( reinterpret_cast(frame_ptr + (pitch * even_height)); av_frame->data[2] = nullptr; av_frame->data[3] = nullptr; - // TODO_API_BREAKDOWN_CUDA: Check range before cast? + // TODO_API_BREAKDOWN_CUDA P2: Check range before cast? av_frame->linesize[0] = static_cast(pitch); av_frame->linesize[1] = static_cast(pitch); av_frame->linesize[2] = 0; @@ -816,10 +856,8 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( return av_frame; } -// TODO_API_BREAKDOWN_CUDA: Does this even nede to be a method? Maybe it can be -// a function that just lives in the PacketDecoder so we don't need to expose -// another API to the DeviceInterface? void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { + CudaContextGuard context_guard(device_.index()); if (!(av_frame->format == AV_PIX_FMT_P016LE || av_frame->format == AV_PIX_FMT_NV12)) { // The CPU frames are already standalone, so we don't need to do anything. @@ -832,7 +870,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { // = num_pixels * 3 / 2 // // To make it correct, we should use num_pixels = pitch * height, not - // num_pixels = pitch * height. The pitch value also accounts for the data + // num_pixels = width * height. The pitch value also accounts for the data // size (uint8 vs uint16) so this is also correct for P016. int64_t even_height = static_cast(round_up_to_even(av_frame->height)); @@ -842,7 +880,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { auto storage = torch::stable::empty({num_bytes}, kStableUInt8, std::nullopt, device_); - // TODO_API_BREAKDOWN_CUDA: I suspect we don't need to wait on the nvdec + // TODO_API_BREAKDOWN_CUDA P1: I suspect we don't need to wait on the nvdec // stream here, because we can only arrive here from a path where the frame // has already been mapped so its data is available - worth double checking. cudaStream_t current_stream = get_current_cuda_stream(device_.index()); @@ -857,7 +895,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { "Failed to copy NVDEC surface: ", cudaGetErrorString(err)); - // TODO_API_BREAKDOWN_CUDA: Should we unmap here? Or let the next + // TODO_API_BREAKDOWN_CUDA P2: Should we unmap here? Or let the next // receive_frame() call do it? // unmap_previous_frame(); @@ -867,7 +905,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { auto attached_data = new StandAloneFrameAttachedData(); attached_data->producer_stream = current_stream; - // TODO_API_BREAKDOWN_CUDA: We don't *really* need to std::move it I guess? + // TODO_API_BREAKDOWN_CUDA P2: We don't *really* need to std::move it I guess? attached_data->storage = std::move(storage); av_frame->opaque_ref = av_buffer_create( reinterpret_cast(attached_data), @@ -878,6 +916,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { } void BetaCudaDeviceInterface::flush() { + CudaContextGuard context_guard(device_.index()); if (cpu_fallback_) { cpu_fallback_->flush(); return; @@ -1044,9 +1083,10 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( const AVFrame& av_frame, FrameOutput& frame_output, std::optional pre_allocated_output_tensor) { - // TODO_API_BREAKDOWN_CUDA is that accurate and safe? Can there be a CPU NV12 - // frame in our code? Should we create a helper used in the make_standalone - // function too? + CudaContextGuard context_guard(device_.index()); + // TODO_API_BREAKDOWN_CUDA P0 is that accurate and safe? Can there be a CPU + // NV12 frame in our code? Should we create a helper used in the + // make_standalone function too? bool cpu_fallback = av_frame.format != AV_PIX_FMT_NV12 && av_frame.format != AV_PIX_FMT_P016LE; @@ -1062,7 +1102,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( av_pix_fmt_desc_get(static_cast(av_frame.format)); bool is444 = desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0; if (is444) { - // TODO_API_BREAKDOWN we need to handle this + // TODO_API_BREAKDOWN P1: we need to handle this FrameOutput cpu_frame_output; cpu_fallback_->convert_av_frame_to_frame_output( av_frame, cpu_frame_output); @@ -1098,7 +1138,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( gpu_frame.format == AV_PIX_FMT_P016LE, "Expected NV12 or P016LE format frame"); - // TODO_API_BREAKDOWN: Cleanup how we get the attached data? Make it more + // TODO_API_BREAKDOWN P1: Cleanup how we get the attached data? Make it more // robust? Should we couple it to a flag on the interface saying "I'm // color-conversion only, I absolutely expect frames to be standalone"? cudaStream_t producer_stream; @@ -1111,7 +1151,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( producer_stream = get_current_cuda_stream(device_.index()); } - // TODO_API_BREAKDOWN: we don't suppor output_dtype so some of that is not + // TODO_API_BREAKDOWN P1: we don't suppor output_dtype so some of that is not // execrcized. auto convert_frame = [&](std::optional pre_alloc) -> torch::stable::Tensor { diff --git a/src/torchcodec/_core/ColorConverter.cpp b/src/torchcodec/_core/ColorConverter.cpp index 262d1946c..1655dab38 100644 --- a/src/torchcodec/_core/ColorConverter.cpp +++ b/src/torchcodec/_core/ColorConverter.cpp @@ -25,14 +25,14 @@ ColorConverter::ColorConverter(const StableDevice& device) { options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet options.device = device; - // TODO_API_BREAKDOWN It seems unnatural that the color-converter needs its + // TODO_API_BREAKDOWN P1 It seems unnatural that the color-converter needs its // own device_interface_, but at the same time the color-conversion *must* be // third-party aware, and the only way to achieve that for now is via the // interface. Should at the very least write a note about this design that now // the DeviceInterface has different modes: decode only, color-convert only, // and decode+color-convert (which used to be the only mode). - // TODO_API_BREAKDOWN: we shouldn't call initialize_video here, this is for + // TODO_API_BREAKDOWN P0: we shouldn't call initialize_video here, this is for // the decoding+color-convert mode. We should do something cleaner e.g. // initialize_color_convertion_only() std::vector> no_transforms; diff --git a/src/torchcodec/_core/CpuDeviceInterface.cpp b/src/torchcodec/_core/CpuDeviceInterface.cpp index 426294b9f..d0d8e2c37 100644 --- a/src/torchcodec/_core/CpuDeviceInterface.cpp +++ b/src/torchcodec/_core/CpuDeviceInterface.cpp @@ -86,7 +86,7 @@ void CpuDeviceInterface::initialize_video( const VideoStreamOptions& video_stream_options, const std::vector>& transforms, const std::optional& resized_output_dims) { - // TODO_API_BREAKDOWN this used to be: + // TODO_API_BREAKDOWN P0 this used to be: // STD_TORCH_CHECK(av_stream != nullptr, "avStream is null"); // time_base_ = av_stream->time_base; // but now that avStrean can be null (to create a standalone color converter) diff --git a/src/torchcodec/_core/Demuxer.cpp b/src/torchcodec/_core/Demuxer.cpp index 333a2306b..9835d0493 100644 --- a/src/torchcodec/_core/Demuxer.cpp +++ b/src/torchcodec/_core/Demuxer.cpp @@ -69,7 +69,7 @@ Demuxer::Demuxer( } UniqueAVPacket Demuxer::next_packet() { - // TODO_API_BREAKDOWN: Not a fan of the ReferenceAVPacket / AutoAVPacket / + // TODO_API_BREAKDOWN P2: Not a fan of the ReferenceAVPacket / AutoAVPacket / // UniqueAVPacket dance here. Can we simplify? ReferenceAVPacket packet(auto_packet_); int status = diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index f05cee09e..6cecf90d7 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -8,7 +8,7 @@ namespace facebook::torchcodec { -// TODO_API_BREAKDOWN: we should make sure the block APIs can dispatch to +// TODO_API_BREAKDOWN P1: we should make sure the block APIs can dispatch to // third-party extensions - all of them. SharedAVCodecContext create_and_open_codec_context( @@ -75,7 +75,7 @@ PacketDecoder::PacketDecoder( options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet options.device = device; - // TODO_API_BREAKDOWN: This isn't right, it's needed only for the NVDEC + // TODO_API_BREAKDOWN P0: This isn't right, it's needed only for the NVDEC // interface. This should probably be initialize_video_only - there's a // sibling TODO in the ColorConverter code (about color-conversion only.) std::vector> no_transforms; diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp index 5b6d18f1a..a2d30a113 100644 --- a/src/torchcodec/_core/custom_ops.cpp +++ b/src/torchcodec/_core/custom_ops.cpp @@ -827,7 +827,7 @@ int64_t _blocks_packet_decoder_send_packet( torch::stable::Tensor& decoder, torch::stable::Tensor& packet) { PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer(decoder); - // TODO_API_BREAKDOWN: Do we really need this to be a raw AVPacket*? + // TODO_API_BREAKDOWN P1: Do we really need this to be a raw AVPacket*? AVPacket* raw_packet = unwrap_tensor_to_pointer(packet); return static_cast(decoder_ptr->send_packet(raw_packet)); } diff --git a/src/torchcodec/decoders/_blocks/_color_converter.py b/src/torchcodec/decoders/_blocks/_color_converter.py index c6d6156bf..786893e95 100644 --- a/src/torchcodec/decoders/_blocks/_color_converter.py +++ b/src/torchcodec/decoders/_blocks/_color_converter.py @@ -11,10 +11,10 @@ from ._frame import DecodedFrame -# TODO_API_BREAKDOWN support output_dtype? -# TODO_API_BREAKDOWN Expose output_shape? - -# TODO_API_BREAKDOWN We need to support rotation metadata!! +# TODO_API_BREAKDOWN FEAT support output_dtype? +# TODO_API_BREAKDOWN FEAT We need to support rotation metadata!! +# TODO_API_BREAKDOWN FEAT Implement seeking? +# TODO_API_BREAKDOWN FEAT Implement range-getting (start/end time) for decoding? class ColorConverter: @@ -30,8 +30,8 @@ class ColorConverter: block is intentionally stream-agnostic. """ - # TODO_API_BREAKDOWN: device default should be None - # TODO_API_BREAKDOWN: add checks for coupling between device param of + # TODO_API_BREAKDOWN P1: device default should be None + # TODO_API_BREAKDOWN P1: add checks for coupling between device param of # PacketDecoder and ColorConverter. What if one is CPU and the other is # CUDA? What if they're different CUDA devices? Maybe we should just error. def __init__(self, device="cpu"): diff --git a/src/torchcodec/decoders/_blocks/_packet_decoder.py b/src/torchcodec/decoders/_blocks/_packet_decoder.py index 53bc795ba..7191d0cfc 100644 --- a/src/torchcodec/decoders/_blocks/_packet_decoder.py +++ b/src/torchcodec/decoders/_blocks/_packet_decoder.py @@ -17,7 +17,7 @@ from ._frame import DecodedFrame, Packet -# TODO_API_BREAKDOWN revisit every single docstring / comments at some point. +# TODO_API_BREAKDOWN P2 revisit every single docstring / comments at some point. class PacketDecoder: @@ -31,7 +31,7 @@ class PacketDecoder: on your own threads. """ - # TODO_API_BREAKDOWN: device default should be None + # TODO_API_BREAKDOWN P1: device default should be None, here and everywhere else def __init__(self, demuxer: Demuxer, device="cpu"): self._handle = _blocks_create_packet_decoder( demuxer._handle, num_threads=1, device=device @@ -56,7 +56,7 @@ def decode(self, packet: Packet) -> list[DecodedFrame]: raise RuntimeError(f"Failed to send packet to decoder (status {status})") return self._drain() - # TODO_API_BREAKDOWN maybe this shouldn't be called flush, at least not + # TODO_API_BREAKDOWN P2 maybe this shouldn't be called flush, at least not # as-is. It's not the same flush as the FFmpeg decoder buffer flush. def flush(self) -> list[DecodedFrame]: """Signal end-of-stream and return all remaining buffered frames. Call diff --git a/test/test_decoders.py b/test/test_decoders.py index 3cc7d0f1f..5530e0c2e 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -3381,20 +3381,9 @@ def _decode_sequential(self, path, device): def _decode_prefetch_frames(self, path, device): # [demux + decode] on one thread || [color-convert] on another. - demuxer = Demuxer(path) - converter = ColorConverter(device=device) + demuxer, decoder, converter = self._make_blocks(path, device) - def demux_and_decode(): - # Constructed here so the PacketDecoder (and thus the CUDA context it - # binds via its device interface) lives on the prefetch worker - # thread that actually runs the cuvid* calls. - # TODO_API_BREAKDOWN_CUDA: this is a temporary workaroun for a real - # issue - needs fixing. Things should work when the objects are - # constructed in a different thread than where they're consumed. - decoder = PacketDecoder(demuxer, device=device) - yield from self._decode(decoder, self._demux(demuxer)) - - frames = self.prefetch(demux_and_decode()) + frames = self.prefetch(self._decode(decoder, self._demux(demuxer))) return list(self._convert(converter, frames)) def _decode_prefetch_packets(self, path, device): @@ -3405,19 +3394,9 @@ def _decode_prefetch_packets(self, path, device): def _decode_prefetch_packets_and_frames(self, path, device): # [demux] || [decode] || [color-convert], each on its own thread. - demuxer = Demuxer(path) - converter = ColorConverter(device=device) + demuxer, decoder, converter = self._make_blocks(path, device) packets = self.prefetch(self._demux(demuxer)) - - def decode(packets): - # Constructed here so the PacketDecoder (and thus the CUDA context it - # binds via its device interface) lives on the prefetch worker - # thread that actually runs the cuvid* calls. - # TODO_API_BREAKDOWN_CUDA: same TODO as above. - decoder = PacketDecoder(demuxer, device=device) - yield from self._decode(decoder, packets) - - frames = self.prefetch(decode(packets)) + frames = self.prefetch(self._decode(decoder, packets)) return list(self._convert(converter, frames)) def _to_frame_batch(self, frames): @@ -3431,6 +3410,8 @@ def _to_frame_batch(self, frames): ), ) + # TODO_API_BREAKDOWN P0: We need to test all assets. Generally we need more + # tests for all features / edge cases that were eventually fixed. @pytest.mark.parametrize( "video", ( @@ -3499,7 +3480,7 @@ def test_color_converter_reused_across_videos(self, device): def test_set_cuda_backend_is_a_noop(self, device): # The blocks always use the NVDEC CUDA backend. Asking for the "ffmpeg" # one changes nothing, rather than silently producing something else. - # TODO_API_BREAKDOWN: let's just error? + # TODO_API_BREAKDOWN P2: let's just error? with set_cuda_backend("ffmpeg"): got = self._to_frame_batch(self._decode_sequential(NASA_VIDEO.path, device)) ref = self._to_frame_batch(self._decode_sequential(NASA_VIDEO.path, device)) From 0c50c0ccf1b822b368f5fb282f43ed3c17f065c9 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 5 Aug 2026 15:51:58 +0100 Subject: [PATCH 13/14] Split initialize_video into initialize_video_decoding and initialize_color_conversion --- .../_core/BetaCudaDeviceInterface.cpp | 25 +++++++++-------- .../_core/BetaCudaDeviceInterface.h | 5 +++- src/torchcodec/_core/ColorConverter.cpp | 16 ++--------- src/torchcodec/_core/CpuDeviceInterface.cpp | 17 +++++------ src/torchcodec/_core/CpuDeviceInterface.h | 11 ++++++-- src/torchcodec/_core/CudaDeviceInterface.cpp | 15 ++++++---- src/torchcodec/_core/CudaDeviceInterface.h | 11 ++++---- src/torchcodec/_core/DeviceInterface.h | 28 +++++++++++++++---- src/torchcodec/_core/PacketDecoder.cpp | 12 ++------ 9 files changed, 76 insertions(+), 64 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index c149acc66..75b7c0998 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -324,16 +324,21 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device) nvcuvid_available_ = load_nvcuvid_library(); } -void BetaCudaDeviceInterface::initialize_video( - const AVStream* av_stream, - const UniqueDecodingAVFormatContext& av_format_ctx, +void BetaCudaDeviceInterface::initialize_color_conversion( const VideoStreamOptions& video_stream_options, const std::vector>& transforms, const std::optional& resized_output_dims) { - // TODO_API_BREAKDOWN P0 - if (!av_stream) { - return; + output_dtype_ = video_stream_options.output_dtype; + if (cpu_fallback_) { + cpu_fallback_->initialize_color_conversion( + video_stream_options, transforms, resized_output_dims); } +} + +void BetaCudaDeviceInterface::initialize_video_decoding( + const AVStream* av_stream, + const UniqueDecodingAVFormatContext& av_format_ctx, + const VideoStreamOptions& video_stream_options) { STD_TORCH_CHECK(av_stream != nullptr, "AVStream cannot be null"); CudaContextGuard context_guard(device_.index()); rotation_ = rotation_from_degrees(get_rotation_from_stream(av_stream)); @@ -354,12 +359,8 @@ void BetaCudaDeviceInterface::initialize_video( STD_TORCH_CHECK( cpu_fallback_ != nullptr, "Failed to create CPU device interface"); cpu_fallback_->initialize(codec_context_); - cpu_fallback_->initialize_video( - av_stream, - av_format_ctx, - video_stream_options, - transforms, - resized_output_dims); + cpu_fallback_->initialize_video_decoding( + av_stream, av_format_ctx, video_stream_options); return; } diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index 3ae019ec1..e9af08cdf 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.h @@ -45,9 +45,12 @@ class BetaCudaDeviceInterface : public DeviceInterface { void initialize(const SharedAVCodecContext& codec_context) override; - void initialize_video( + void initialize_video_decoding( const AVStream* av_stream, const UniqueDecodingAVFormatContext& av_format_ctx, + const VideoStreamOptions& video_stream_options) override; + + void initialize_color_conversion( const VideoStreamOptions& video_stream_options, const std::vector>& transforms, const std::optional& resized_output_dims) override; diff --git a/src/torchcodec/_core/ColorConverter.cpp b/src/torchcodec/_core/ColorConverter.cpp index 1655dab38..196e32a05 100644 --- a/src/torchcodec/_core/ColorConverter.cpp +++ b/src/torchcodec/_core/ColorConverter.cpp @@ -28,20 +28,10 @@ ColorConverter::ColorConverter(const StableDevice& device) { // TODO_API_BREAKDOWN P1 It seems unnatural that the color-converter needs its // own device_interface_, but at the same time the color-conversion *must* be // third-party aware, and the only way to achieve that for now is via the - // interface. Should at the very least write a note about this design that now - // the DeviceInterface has different modes: decode only, color-convert only, - // and decode+color-convert (which used to be the only mode). - - // TODO_API_BREAKDOWN P0: we shouldn't call initialize_video here, this is for - // the decoding+color-convert mode. We should do something cleaner e.g. - // initialize_color_convertion_only() + // interface. std::vector> no_transforms; - device_interface_->initialize_video( - /*av_stream=*/nullptr, - UniqueDecodingAVFormatContext{}, - options, - no_transforms, - /*resized_output_dims=*/std::nullopt); + device_interface_->initialize_color_conversion( + options, no_transforms, /*resized_output_dims=*/std::nullopt); } torch::stable::Tensor ColorConverter::convert(const AVFrame& av_frame) { diff --git a/src/torchcodec/_core/CpuDeviceInterface.cpp b/src/torchcodec/_core/CpuDeviceInterface.cpp index d0d8e2c37..0adb74d09 100644 --- a/src/torchcodec/_core/CpuDeviceInterface.cpp +++ b/src/torchcodec/_core/CpuDeviceInterface.cpp @@ -80,21 +80,18 @@ void CpuDeviceInterface::initialize(const SharedAVCodecContext& codec_context) { codec_context_ = codec_context; } -void CpuDeviceInterface::initialize_video( +void CpuDeviceInterface::initialize_video_decoding( const AVStream* av_stream, [[maybe_unused]] const UniqueDecodingAVFormatContext& av_format_ctx, + [[maybe_unused]] const VideoStreamOptions& video_stream_options) { + STD_TORCH_CHECK(av_stream != nullptr, "avStream is null"); + time_base_ = av_stream->time_base; +} + +void CpuDeviceInterface::initialize_color_conversion( const VideoStreamOptions& video_stream_options, const std::vector>& transforms, const std::optional& resized_output_dims) { - // TODO_API_BREAKDOWN P0 this used to be: - // STD_TORCH_CHECK(av_stream != nullptr, "avStream is null"); - // time_base_ = av_stream->time_base; - // but now that avStrean can be null (to create a standalone color converter) - // we need this workaround. This is bad, we need to preserve the previous - // check somehow. See corresponding TODO in color-converter and packet decoder - // code. - time_base_ = (av_stream != nullptr) ? av_stream->time_base - : AVRational{1, AV_TIME_BASE}; av_media_type_ = AVMEDIA_TYPE_VIDEO; video_stream_options_ = video_stream_options; resized_output_dims_ = resized_output_dims; diff --git a/src/torchcodec/_core/CpuDeviceInterface.h b/src/torchcodec/_core/CpuDeviceInterface.h index 2d57eaf7b..556698330 100644 --- a/src/torchcodec/_core/CpuDeviceInterface.h +++ b/src/torchcodec/_core/CpuDeviceInterface.h @@ -27,9 +27,12 @@ class CpuDeviceInterface : public DeviceInterface { virtual void initialize(const SharedAVCodecContext& codec_context) override; - virtual void initialize_video( + virtual void initialize_video_decoding( const AVStream* av_stream, const UniqueDecodingAVFormatContext& av_format_ctx, + const VideoStreamOptions& video_stream_options) override; + + virtual void initialize_color_conversion( const VideoStreamOptions& video_stream_options, const std::vector>& transforms, const std::optional& resized_output_dims) override; @@ -76,7 +79,11 @@ class CpuDeviceInterface : public DeviceInterface { const FrameDims& output_dims) const; VideoStreamOptions video_stream_options_; - AVRational time_base_; + // Default used when color conversion runs standalone (no stream to derive it + // from, e.g. the ColorConverter block API). initialize_video_decoding() + // overrides it from the stream when there is one. Its value doesn't matter on + // color-conversion-only mode, but filtergraph still expects it. + AVRational time_base_ = {1, AV_TIME_BASE}; AVPixelFormat output_pixel_format_; // If the resized output dimensions are present, then we always use those as diff --git a/src/torchcodec/_core/CudaDeviceInterface.cpp b/src/torchcodec/_core/CudaDeviceInterface.cpp index 1e8c2c34f..d0e04a545 100644 --- a/src/torchcodec/_core/CudaDeviceInterface.cpp +++ b/src/torchcodec/_core/CudaDeviceInterface.cpp @@ -114,17 +114,13 @@ void CudaDeviceInterface::initialize( codec_context_ = codec_context; } -void CudaDeviceInterface::initialize_video( +void CudaDeviceInterface::initialize_video_decoding( const AVStream* av_stream, const UniqueDecodingAVFormatContext& av_format_ctx, - const VideoStreamOptions& video_stream_options, - [[maybe_unused]] const std::vector>& transforms, - [[maybe_unused]] const std::optional& resized_output_dims) { + [[maybe_unused]] const VideoStreamOptions& video_stream_options) { STD_TORCH_CHECK(av_stream != nullptr, "avStream is null"); time_base_ = av_stream->time_base; - video_stream_options_ = video_stream_options; - // TODO: Ideally, we should keep all interface implementations independent. cpu_interface_ = create_device_interface(kStableCPU); STD_TORCH_CHECK( cpu_interface_ != nullptr, "Failed to create CPU device interface"); @@ -137,6 +133,13 @@ void CudaDeviceInterface::initialize_video( /*resizedOutputDims=*/std::nullopt); } +void CudaDeviceInterface::initialize_color_conversion( + const VideoStreamOptions& video_stream_options, + [[maybe_unused]] const std::vector>& transforms, + [[maybe_unused]] const std::optional& resized_output_dims) { + video_stream_options_ = video_stream_options; +} + void CudaDeviceInterface::register_hardware_device_with_codec( AVCodecContext* codec_context) { STD_TORCH_CHECK( diff --git a/src/torchcodec/_core/CudaDeviceInterface.h b/src/torchcodec/_core/CudaDeviceInterface.h index d449ccb97..00cc5e817 100644 --- a/src/torchcodec/_core/CudaDeviceInterface.h +++ b/src/torchcodec/_core/CudaDeviceInterface.h @@ -28,14 +28,15 @@ class CudaDeviceInterface : public DeviceInterface { void initialize(const SharedAVCodecContext& codec_context) override; - void initialize_video( + void initialize_video_decoding( const AVStream* av_stream, const UniqueDecodingAVFormatContext& av_format_ctx, + const VideoStreamOptions& video_stream_options) override; + + void initialize_color_conversion( const VideoStreamOptions& video_stream_options, - [[maybe_unused]] const std::vector>& - transforms, - [[maybe_unused]] const std::optional& resized_output_dims) - override; + const std::vector>& transforms, + const std::optional& resized_output_dims) override; void register_hardware_device_with_codec( AVCodecContext* codec_context) override; diff --git a/src/torchcodec/_core/DeviceInterface.h b/src/torchcodec/_core/DeviceInterface.h index a37b292d7..f25efb342 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/src/torchcodec/_core/DeviceInterface.h @@ -56,15 +56,33 @@ class DeviceInterface { // default sendPacket/receiveFrame/flush implementations. virtual void initialize(const SharedAVCodecContext& codec_context) = 0; - // Initialize the device with parameters specific to video decoding. There is - // a default empty implementation. - virtual void initialize_video( + // Initialize state needed to decode packets into raw AVFrames. + virtual void initialize_video_decoding( [[maybe_unused]] const AVStream* av_stream, [[maybe_unused]] const UniqueDecodingAVFormatContext& av_format_ctx, + [[maybe_unused]] const VideoStreamOptions& video_stream_options) {} + + // Initialize state needed to color-convert decoded AVFrames into output + // tensors. + virtual void initialize_color_conversion( [[maybe_unused]] const VideoStreamOptions& video_stream_options, [[maybe_unused]] const std::vector>& - transforms, - [[maybe_unused]] const std::optional& resized_output_dims) {} + transforms = {}, + [[maybe_unused]] const std::optional& resized_output_dims = + std::nullopt) {} + + // Convenience for the combined decode + color-convert path, kept for BC as + // it's used by SingleStreamDecoder and out-of-tree interfaces rely on it. + void initialize_video( + const AVStream* av_stream, + const UniqueDecodingAVFormatContext& av_format_ctx, + const VideoStreamOptions& video_stream_options, + const std::vector>& transforms, + const std::optional& resized_output_dims) { + initialize_video_decoding(av_stream, av_format_ctx, video_stream_options); + initialize_color_conversion( + video_stream_options, transforms, resized_output_dims); + } // Initialize the device with parameters specific to audio decoding. There is // a default empty implementation. diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index 6cecf90d7..b9960e92d 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -75,16 +75,8 @@ PacketDecoder::PacketDecoder( options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet options.device = device; - // TODO_API_BREAKDOWN P0: This isn't right, it's needed only for the NVDEC - // interface. This should probably be initialize_video_only - there's a - // sibling TODO in the ColorConverter code (about color-conversion only.) - std::vector> no_transforms; - device_interface_->initialize_video( - stream, - demuxer.format_context(), - options, - no_transforms, - /*resized_output_dims=*/std::nullopt); + device_interface_->initialize_video_decoding( + stream, demuxer.format_context(), options); } int PacketDecoder::send_packet(AVPacket* packet) { From 3b37c5ad784512d4aecf5e0f3afb3021cd21505a Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 5 Aug 2026 17:08:16 +0100 Subject: [PATCH 14/14] Comments --- .../_core/BetaCudaDeviceInterface.cpp | 33 ++++++++++++------- src/torchcodec/_core/color_conversion.cpp | 2 +- test/test_decoders.py | 5 --- 3 files changed, 22 insertions(+), 18 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 75b7c0998..2d977b9a9 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -733,14 +733,14 @@ int BetaCudaDeviceInterface::receive_frame(UniqueAVFrame& av_frame) { // unmap will cause map to eventually fail. DALI unmaps frames almost // immediately after mapping them: they do the color-conversion in-between, // which involves a copy of the data, so that works. - // We, OTOH, will do the color-conversion later, outside of ReceiveFrame(). So - // we unmap here: just before mapping a new frame. At that point we know that - // the previously-mapped frame is no longer needed: it was either - // color-converted (with a copy), or that's a frame that was discarded in - // SingleStreamDecoder. Either way, the underlying output surface can be - // safely re-used. - // TODO_API_BREAKDOWN P1: We should update this comment slightly to now - // account for the frame copy we do in make_frame_standalone() + // We, OTOH, will do the color-conversion later, outside of receive_frame(). + // So we unmap here: just before mapping a new frame. At that point we know + // that the previously-mapped frame is no longer needed: + // - With SingleStreamDecoder, that frame was either color-converted (with a + // copy), or that's a frame that was discarded in SingleStreamDecoder. + // Either way, the underlying output surface can be safely re-used. + // - With the "Blocks" APIs, the PacketDecoder forces a copy in + // make_frame_standalone(). unmap_previous_frame(); CUresult result = cuvidMapVideoFrame( *decoder_.get(), @@ -858,6 +858,15 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( } void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { + // Make the frame standalone: + // - Crucially, we copy the frame data so that its surface can be unmapped in + // receive_frame() (see comment there). + // - We put the frame in a state such that it can be safely used by a + // ColorConverter (i.e. a *different* instance of this + // BetaCudaDeviceInterface): we attach relevant metadata as the + // StandAloneFrameAttachedData struct, which is then used by the + // ColorConverter in convert_cuda_frame_to_av_frame() to perform the + // color-conversion correctly. CudaContextGuard context_guard(device_.index()); if (!(av_frame->format == AV_PIX_FMT_P016LE || av_frame->format == AV_PIX_FMT_NV12)) { @@ -865,14 +874,14 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { return; } - // Roughly, the number of bytes an NV12 image takes is: + // The amount of bytes an NV12 image takes is: // num_bytes = len(Y) + len(UV) // = num_pixels + num_pixels / 2 // = num_pixels * 3 / 2 // - // To make it correct, we should use num_pixels = pitch * height, not - // num_pixels = width * height. The pitch value also accounts for the data - // size (uint8 vs uint16) so this is also correct for P016. + // where num_pixels = pitch * height, not num_pixels = width * height. The + // pitch value also accounts for the data size (uint8 vs uint16) so this is + // also correct for P016. int64_t even_height = static_cast(round_up_to_even(av_frame->height)); int64_t pitch = static_cast(av_frame->linesize[0]); diff --git a/src/torchcodec/_core/color_conversion.cpp b/src/torchcodec/_core/color_conversion.cpp index 98ffcc35f..b05e8273c 100644 --- a/src/torchcodec/_core/color_conversion.cpp +++ b/src/torchcodec/_core/color_conversion.cpp @@ -221,7 +221,7 @@ torch::stable::Tensor convert_yuv_frame_to_rgb( FrameDims(out_height, out_width), device, out_dtype); } - // TODO_API_BREAKDOWN: This may not be the semantic that we want: this will + // TODO_API_BREAKDOWN P1: This may not be the semantic that we want: this will // wait for all ongoin work on the producer stream to finish. But maybe the // producer stream produced the frame data a long time ago, and lots of // kernels have been launched on it already. We'd be waiting on those to diff --git a/test/test_decoders.py b/test/test_decoders.py index 5530e0c2e..9cbc492fd 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -3275,9 +3275,6 @@ def test_multiple_calls_with_backward_seeks(self): assert wav_samples.pts_seconds == audio_samples.pts_seconds -# The blocks support CPU and the NVDEC CUDA backend, and nothing else. In -# particular they never use the "ffmpeg" CUDA backend, so this is deliberately -# not all_supported_devices(). def _block_devices(): return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda)) @@ -3362,8 +3359,6 @@ def drain(): def _make_blocks(path, device): demuxer = Demuxer(path) decoder = PacketDecoder(demuxer, device=device) - # The converter is standalone on every device: it is not bound to the - # decoder that produced the frames. converter = ColorConverter(device=device) return demuxer, decoder, converter