From 6e589f999f5a3a8659bf02ce284aaebf3ab3bf2c Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Fri, 4 Sep 2026 09:10:57 +0100 Subject: [PATCH] Pass packets as const AVPacket& through the decode seam DeviceInterface::send_packet() only ever reads the packet - the default implementation hands it to avcodec_send_packet(), which takes a const one. Except that BetaCuda's bitstream filter path did not: av_bsf_send_packet() takes ownership of what it is given and resets it, so sending a packet through the seam could empty it, silently, with nothing in the signature saying so. SingleStreamDecoder never noticed, because its packet is dead after the send. PacketDecoder did: the packet it is handed is a user-visible Python object that has to survive being decoded, so it referenced it into a throwaway packet before every single send - a per-packet allocation, on every device, to defend against a contract that was never stated. State the contract instead. The seam takes a const AVPacket&, and the one implementation that needs ownership references the packet itself, where the reason for it is local. --- .../_core/BetaCudaDeviceInterface.cpp | 33 ++++++++++++------- .../_core/BetaCudaDeviceInterface.h | 6 ++-- src/torchcodec/_core/DeviceInterface.h | 6 ++-- src/torchcodec/_core/PacketDecoder.cpp | 13 ++------ src/torchcodec/_core/PacketDecoder.h | 2 +- src/torchcodec/_core/SingleStreamDecoder.cpp | 2 +- src/torchcodec/_core/custom_ops.cpp | 5 ++- 7 files changed, 35 insertions(+), 32 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 4e2ad613a..10cf57f70 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -686,14 +686,14 @@ 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) { +int BetaCudaDeviceInterface::send_packet(const AVPacket& packet) { CudaContextGuard context_guard(device_.index()); if (decoding_on_cpu_) { return cpu_interface_->send_packet(packet); } STD_TORCH_CHECK( - packet.get() && packet->data && packet->size > 0, + packet.data && packet.size > 0, "sendPacket received an empty packet, this is unexpected, please report."); // Apply BSF if needed. We want applyBSF to return a *new* filtered packet, or @@ -703,15 +703,15 @@ int BetaCudaDeviceInterface::send_packet(ReferenceAVPacket& packet) { // reference. AutoAVPacket filtered_auto_packet; ReferenceAVPacket filtered_packet(filtered_auto_packet); - ReferenceAVPacket& packet_to_send = apply_bsf(packet, filtered_packet); + const AVPacket& packet_to_send = apply_bsf(packet, filtered_packet); CUVIDSOURCEDATAPACKET cuvid_packet = {}; - cuvid_packet.payload = packet_to_send->data; - cuvid_packet.payload_size = packet_to_send->size; + cuvid_packet.payload = packet_to_send.data; + cuvid_packet.payload_size = packet_to_send.size; cuvid_packet.flags = CUVID_PKT_TIMESTAMP; - cuvid_packet.timestamp = packet_to_send->pts; + cuvid_packet.timestamp = packet_to_send.pts; - if (packet_to_send->flags & AV_PKT_FLAG_DISCARD) { + if (packet_to_send.flags & AV_PKT_FLAG_DISCARD) { discarded_timestamps_.insert(cuvid_packet.timestamp); } @@ -737,14 +737,25 @@ int BetaCudaDeviceInterface::send_cuvid_packet( return result == CUDA_SUCCESS ? AVSUCCESS : AVERROR_EXTERNAL; } -ReferenceAVPacket& BetaCudaDeviceInterface::apply_bsf( - ReferenceAVPacket& packet, +const AVPacket& BetaCudaDeviceInterface::apply_bsf( + const AVPacket& packet, ReferenceAVPacket& filtered_packet) { if (!bitstream_filter_) { return packet; } - int ret_val = av_bsf_send_packet(bitstream_filter_.get(), packet.get()); + // av_bsf_send_packet() takes ownership of what it is given: it moves the + // reference out of the packet, leaving it empty. Our caller only lends us + // theirs, so send a reference of our own instead. + AutoAVPacket auto_input_packet; + ReferenceAVPacket input_packet(auto_input_packet); + int ret_val = av_packet_ref(input_packet.get(), &packet); + STD_TORCH_CHECK( + ret_val >= AVSUCCESS, + "Failed to reference packet for the bitstream filter: ", + get_ffmpeg_error_string_from_error_code(ret_val)); + + ret_val = av_bsf_send_packet(bitstream_filter_.get(), input_packet.get()); STD_TORCH_CHECK( ret_val >= AVSUCCESS, "Failed to send packet to bitstream filter: ", @@ -761,7 +772,7 @@ ReferenceAVPacket& BetaCudaDeviceInterface::apply_bsf( "Failed to receive packet from bitstream filter: ", get_ffmpeg_error_string_from_error_code(ret_val)); - return filtered_packet; + return *filtered_packet; } // Parser triggers this callback within cuvidParseVideoData when a frame is diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index d9a7b1a1f..44f9b571b 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.h @@ -74,7 +74,7 @@ class BetaCudaDeviceInterface : public DeviceInterface { std::optional pre_allocated_output_tensor) override; - int send_packet(ReferenceAVPacket& packet) override; + int send_packet(const AVPacket& packet) override; int send_eof_packet() override; int receive_frame(UniqueAVFrame& av_frame) override; void flush() override; @@ -99,8 +99,8 @@ class BetaCudaDeviceInterface : public DeviceInterface { const UniqueDecodingAVFormatContext& av_format_ctx); // Apply bitstream filter, returns filtered packet or original if no filter // needed. - ReferenceAVPacket& apply_bsf( - ReferenceAVPacket& packet, + const AVPacket& apply_bsf( + const AVPacket& packet, ReferenceAVPacket& filtered_packet); CUdeviceptr previously_mapped_frame_ = 0; diff --git a/src/torchcodec/_core/DeviceInterface.h b/src/torchcodec/_core/DeviceInterface.h index d1630062a..64ebe592a 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/src/torchcodec/_core/DeviceInterface.h @@ -134,11 +134,13 @@ class DeviceInterface { // Returns AVSUCCESS on success, AVERROR(EAGAIN) if decoder queue full, or // other AVERROR on failure // Default implementation uses FFmpeg directly - virtual int send_packet(ReferenceAVPacket& av_packet) { + // The packet is borrowed: an implementation that hands it to something which + // takes ownership (a bitstream filter, say) must reference it first. + virtual int send_packet(const AVPacket& av_packet) { STD_TORCH_CHECK( codec_context_ != nullptr, "Codec context not available for default packet sending"); - return avcodec_send_packet(codec_context_.get(), av_packet.get()); + return avcodec_send_packet(codec_context_.get(), &av_packet); } // Send an EOF packet to flush the decoder diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index fc7b8ac25..6682ec43a 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -116,17 +116,8 @@ PacketDecoder::PacketDecoder( stream, demuxer.format_context(), options); } -int PacketDecoder::send_packet(AVPacket* packet) { - // The decode seam expects a ReferenceAVPacket. Copy a reference of the - // caller- owned packet into a temporary one (cheap, refcount bump); the - // temporary is unref'd on scope exit while the caller retains ownership of - // `packet`. - AutoAVPacket auto_packet; - ReferenceAVPacket ref(auto_packet); - int status = av_packet_ref(ref.get(), packet); - STD_TORCH_CHECK(status >= AVSUCCESS, "av_packet_ref failed"); - - status = device_interface_->send_packet(ref); +int PacketDecoder::send_packet(const AVPacket& packet) { + int status = device_interface_->send_packet(packet); if (status == AVERROR_INVALIDDATA && packet_data_may_be_misaligned_) { // Seeking in an MPEG program stream lands on a container-level byte offset, diff --git a/src/torchcodec/_core/PacketDecoder.h b/src/torchcodec/_core/PacketDecoder.h index e1dba4eb3..02d127e97 100644 --- a/src/torchcodec/_core/PacketDecoder.h +++ b/src/torchcodec/_core/PacketDecoder.h @@ -41,7 +41,7 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder { std::optional ffmpeg_thread_count = std::nullopt); // Feed one packet to the decoder. Borrows `packet` (does not take ownership). - int send_packet(AVPacket* packet); + int send_packet(const AVPacket& packet); // Signal end-of-stream so the decoder flushes its remaining frames. int send_eof(); // Pull one frame. Returns AVSUCCESS with `av_frame` filled, AVERROR(EAGAIN) diff --git a/src/torchcodec/_core/SingleStreamDecoder.cpp b/src/torchcodec/_core/SingleStreamDecoder.cpp index f9ef433e8..63027a80b 100644 --- a/src/torchcodec/_core/SingleStreamDecoder.cpp +++ b/src/torchcodec/_core/SingleStreamDecoder.cpp @@ -1473,7 +1473,7 @@ UniqueAVFrame SingleStreamDecoder::decode_av_frame( // We got a valid packet. Send it to the decoder, and we'll receive it in // the next iteration. - status = device_interface_->send_packet(packet); + status = device_interface_->send_packet(*packet); if (status == AVERROR_INVALIDDATA && packet_data_may_be_misaligned) { // The MPEG-PS demuxer doesn't return proper packets just after a seek, so diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp index 9c17316e3..25c33bdd3 100644 --- a/src/torchcodec/_core/custom_ops.cpp +++ b/src/torchcodec/_core/custom_ops.cpp @@ -963,9 +963,8 @@ 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 CC 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)); + return static_cast( + decoder_ptr->send_packet(*unwrap_tensor_to_pointer(packet))); } int64_t _blocks_packet_decoder_send_eof(torch::stable::Tensor& decoder) {