diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index b0f732c15..d916afdc4 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -676,14 +676,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 @@ -693,15 +693,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); } @@ -727,14 +727,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: ", @@ -751,7 +762,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 fc2b82cbd..b68010c99 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.h @@ -73,7 +73,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; @@ -98,8 +98,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 e7036199e..6a644dee8 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/src/torchcodec/_core/DeviceInterface.h @@ -130,11 +130,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 c920e3376..49bc0dd97 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -121,17 +121,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 1c2b944c5..93adfcaf0 100644 --- a/src/torchcodec/_core/SingleStreamDecoder.cpp +++ b/src/torchcodec/_core/SingleStreamDecoder.cpp @@ -1479,7 +1479,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) {