Skip to content

Store a UniqueAVFrame in the blocks frame handle #1596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
NicolasHug wants to merge 1 commit into from
Closed

Store a UniqueAVFrame in the blocks frame handle #1596

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 8 additions & 39 deletions src/torchcodec/_core/custom_ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ SingleStreamDecoder* unwrap_tensor_to_get_decoder(
}

// Generic pointer<->tensor laundering for the building-block handle types
// (Demuxer / PacketDecoder / ColorConverter). Same trick as
// (Demuxer / PacketDecoder / ColorConverter / UniqueAVFrame). 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 <typename T>
Expand All @@ -206,9 +206,9 @@ T* unwrap_tensor_to_pointer(torch::stable::Tensor& tensor) {
return static_cast<T*>(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.
// Opaque packet handle: launder a raw AVPacket* 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;
Expand All @@ -230,32 +230,6 @@ AVPacket* unwrap_tensor_to_packet(torch::stable::Tensor& tensor) {
return static_cast<AVPacket*>(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<AVFrame*>(tensor.mutable_data_ptr());
}

torch::stable::Tensor wrap_multi_stream_encoder_pointer_to_tensor(
std::unique_ptr<MultiStreamEncoder> unique_encoder) {
MultiStreamEncoder* encoder = unique_encoder.release();
Expand Down Expand Up @@ -910,9 +884,9 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
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();
return std::make_tuple(
wrap_frame_pointer_to_tensor(raw_frame),
wrap_pointer_to_tensor<UniqueAVFrame>(
std::make_unique<UniqueAVFrame>(std::move(av_frame))),
static_cast<int64_t>(0),
pts_seconds,
duration_seconds);
Expand All @@ -932,13 +906,8 @@ torch::stable::Tensor _blocks_convert_frame(
torch::stable::Tensor& frame) {
ColorConverter* converter_ptr =
unwrap_tensor_to_pointer<ColorConverter>(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;
UniqueAVFrame* frame_ptr = unwrap_tensor_to_pointer<UniqueAVFrame>(frame);
return converter_ptr->convert(*frame_ptr);
}

// For testing only. We need to implement this operation as a core library
Expand Down
Loading