Skip to content

Commit 6899cb5

Browse files
committed
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.
1 parent 3f7692b commit 6899cb5

25 files changed

Lines changed: 201 additions & 254 deletions

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -752,7 +752,8 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame(
752752
// Note that we used to rely on videoFormat_.frame_rate for this, but that
753753
// proved less accurate than FFmpeg.
754754
set_duration(
755-
av_frame, compute_safe_duration(frame_rate_avg_from_ffmpeg_, time_base_));
755+
*av_frame,
756+
compute_safe_duration(frame_rate_avg_from_ffmpeg_, time_base_));
756757

757758
// We need to assign the frame colorspace. This is crucial for proper color
758759
// conversion. NVCUVID stores that in the matrix_coefficients field, but
@@ -824,7 +825,7 @@ void BetaCudaDeviceInterface::flush() {
824825
}
825826

826827
UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
827-
UniqueAVFrame& cpu_frame,
828+
const AVFrame& cpu_frame,
828829
AVPixelFormat target_pix_fmt) {
829830
// This is called in the context of the CPU fallback: the frame was decoded on
830831
// 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(
838839
// (rounded up) width and height, even if the original CPU frame had odd
839840
// dimensions.
840841

841-
STD_TORCH_CHECK(cpu_frame != nullptr, "CPU frame cannot be null");
842842
// NV12 = 1 byte per sample, P016 = 2 bytes per sample
843843
STD_TORCH_CHECK(
844844
target_pix_fmt == AV_PIX_FMT_NV12 || target_pix_fmt == AV_PIX_FMT_P016LE,
845845
"targetPixFmt must be NV12 or P016LE");
846846
int bytes_per_sample = (target_pix_fmt == AV_PIX_FMT_P016LE) ? 2 : 1;
847847

848-
int width = cpu_frame->width;
849-
int height = cpu_frame->height;
848+
int width = cpu_frame.width;
849+
int height = cpu_frame.height;
850850
int even_width = round_up_to_even(width);
851851
int even_height = round_up_to_even(height);
852852

@@ -868,8 +868,8 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
868868
SwsConfig sws_config(
869869
width,
870870
height,
871-
static_cast<AVPixelFormat>(cpu_frame->format),
872-
cpu_frame->colorspace,
871+
static_cast<AVPixelFormat>(cpu_frame.format),
872+
cpu_frame.colorspace,
873873
even_width,
874874
even_height,
875875
target_pix_fmt);
@@ -881,8 +881,8 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
881881

882882
int converted_height = sws_scale(
883883
sws_context_.get(),
884-
cpu_frame->data,
885-
cpu_frame->linesize,
884+
cpu_frame.data,
885+
cpu_frame.linesize,
886886
0,
887887
height,
888888
intermediate_cpu_frame->data,
@@ -944,7 +944,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
944944
"Failed to copy UV plane to GPU: ",
945945
cudaGetErrorString(err));
946946

947-
ret = av_frame_copy_props(gpu_frame.get(), cpu_frame.get());
947+
ret = av_frame_copy_props(gpu_frame.get(), &cpu_frame);
948948
STD_TORCH_CHECK(
949949
ret >= 0,
950950
"Failed to copy frame properties: ",
@@ -967,7 +967,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
967967
}
968968

969969
void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
970-
UniqueAVFrame& av_frame,
970+
const AVFrame& av_frame,
971971
FrameOutput& frame_output,
972972
std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
973973
if (cpu_fallback_) {
@@ -979,7 +979,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
979979
// do the color conversion on the CPU and then send the full RGB frame to
980980
// the GPU.
981981
const AVPixFmtDescriptor* desc =
982-
av_pix_fmt_desc_get(static_cast<AVPixelFormat>(av_frame->format));
982+
av_pix_fmt_desc_get(static_cast<AVPixelFormat>(av_frame.format));
983983
bool is444 = desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0;
984984
if (is444) {
985985
FrameOutput cpu_frame_output;
@@ -1001,28 +1001,29 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
10011001

10021002
// Capture original dimensions before transferCpuFrameToGpu()
10031003
// may round them up to even.
1004-
FrameDims original_dims(av_frame->height, av_frame->width);
1004+
FrameDims original_dims(av_frame.height, av_frame.width);
10051005

1006-
UniqueAVFrame gpu_frame;
1006+
// On the CPU fallback we own the GPU frame we just created; otherwise the
1007+
// input frame is already what we need, and we only observe it.
1008+
UniqueAVFrame transferred_frame;
10071009
if (cpu_fallback_) {
10081010
AVPixelFormat target_pix_fmt = (output_dtype_ == OutputDtype::FLOAT32)
10091011
? AV_PIX_FMT_P016LE
10101012
: AV_PIX_FMT_NV12;
1011-
gpu_frame = transfer_cpu_frame_to_gpu(av_frame, target_pix_fmt);
1012-
} else {
1013-
gpu_frame = std::move(av_frame);
1013+
transferred_frame = transfer_cpu_frame_to_gpu(av_frame, target_pix_fmt);
10141014
}
1015+
const AVFrame& gpu_frame = cpu_fallback_ ? *transferred_frame : av_frame;
10151016

10161017
STD_TORCH_CHECK(
1017-
gpu_frame->format == AV_PIX_FMT_NV12 ||
1018-
gpu_frame->format == AV_PIX_FMT_P016LE,
1018+
gpu_frame.format == AV_PIX_FMT_NV12 ||
1019+
gpu_frame.format == AV_PIX_FMT_P016LE,
10191020
"Expected NV12 or P016LE format frame");
10201021

10211022
cudaStream_t nvdec_stream = get_current_cuda_stream(device_.index());
10221023

10231024
auto convert_frame = [&](std::optional<torch::stable::Tensor> pre_alloc)
10241025
-> torch::stable::Tensor {
1025-
bool is_p016 = (gpu_frame->format == AV_PIX_FMT_P016LE);
1026+
bool is_p016 = (gpu_frame.format == AV_PIX_FMT_P016LE);
10261027
int bit_depth = 8;
10271028
if (is_p016) {
10281029
bit_depth = cpu_fallback_

src/torchcodec/_core/BetaCudaDeviceInterface.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
5252
OutputDtype requested_dtype) const override;
5353

5454
void convert_av_frame_to_frame_output(
55-
UniqueAVFrame& av_frame,
55+
const AVFrame& av_frame,
5656
FrameOutput& frame_output,
5757
std::optional<torch::stable::Tensor> pre_allocated_output_tensor)
5858
override;
@@ -92,7 +92,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
9292
const CUVIDPARSERDISPINFO& disp_info);
9393

9494
UniqueAVFrame transfer_cpu_frame_to_gpu(
95-
UniqueAVFrame& cpu_frame,
95+
const AVFrame& cpu_frame,
9696
AVPixelFormat target_pix_fmt);
9797

9898
void apply_rotation(

src/torchcodec/_core/ColorConverter.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ ColorConverter::ColorConverter(
4545
/*resized_output_dims=*/std::nullopt);
4646
}
4747

48-
torch::stable::Tensor ColorConverter::convert(UniqueAVFrame& av_frame) {
48+
torch::stable::Tensor ColorConverter::convert(const AVFrame& av_frame) {
4949
FrameOutput frame_output;
5050
device_interface_->convert_av_frame_to_frame_output(
5151
av_frame, frame_output, std::nullopt);

src/torchcodec/_core/ColorConverter.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ class FORCE_PUBLIC_VISIBILITY ColorConverter {
2121
const StableDevice& device = StableDevice(kStableCPU),
2222
std::string_view device_variant = "default");
2323

24-
torch::stable::Tensor convert(UniqueAVFrame& av_frame);
24+
torch::stable::Tensor convert(const AVFrame& av_frame);
2525

2626
private:
2727
std::unique_ptr<DeviceInterface> device_interface_;

src/torchcodec/_core/CpuDeviceInterface.cpp

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ ColorConversionLibrary CpuDeviceInterface::get_color_conversion_library(
200200
}
201201

202202
void CpuDeviceInterface::convert_av_frame_to_frame_output(
203-
UniqueAVFrame& av_frame,
203+
const AVFrame& av_frame,
204204
FrameOutput& frame_output,
205205
std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
206206
STD_TORCH_CHECK(initialized_, "CpuDeviceInterface was not initialized.");
@@ -223,7 +223,7 @@ void CpuDeviceInterface::convert_av_frame_to_frame_output(
223223
// Dimension order of the preAllocatedOutputTensor must be HWC, regardless of
224224
// `dimension_order` parameter. It's up to callers to re-shape it if needed.
225225
void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
226-
UniqueAVFrame& av_frame,
226+
const AVFrame& av_frame,
227227
FrameOutput& frame_output,
228228
std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
229229
// 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(
239239
// Both cases cause problems for our batch APIs, as we allocate
240240
// FrameBatchOutputs based on the the stream metadata. But single-frame APIs
241241
// can still work in such situations, so they should.
242-
auto input_dims = FrameDims(av_frame->height, av_frame->width);
242+
auto input_dims = FrameDims(av_frame.height, av_frame.width);
243243
auto output_dims = resized_output_dims_.value_or(input_dims);
244244

245245
if (pre_allocated_output_tensor.has_value()) {
@@ -264,12 +264,12 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
264264
pre_allocated_output_tensor.value_or(allocate_empty_hwc_tensor(
265265
output_dims, kStableCPU, video_stream_options_.output_dtype));
266266

267-
auto av_frame_format = static_cast<AVPixelFormat>(av_frame->format);
267+
auto av_frame_format = static_cast<AVPixelFormat>(av_frame.format);
268268
SwsConfig sws_config(
269-
av_frame->width,
270-
av_frame->height,
269+
av_frame.width,
270+
av_frame.height,
271271
av_frame_format,
272-
av_frame->colorspace,
272+
av_frame.colorspace,
273273
output_dims.width,
274274
output_dims.height,
275275
output_pixel_format_);
@@ -326,15 +326,15 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
326326

327327
torch::stable::Tensor
328328
CpuDeviceInterface::convert_av_frame_to_tensor_using_filter_graph(
329-
const UniqueAVFrame& av_frame,
329+
const AVFrame& av_frame,
330330
const FrameDims& output_dims) {
331-
auto av_frame_format = static_cast<AVPixelFormat>(av_frame->format);
331+
auto av_frame_format = static_cast<AVPixelFormat>(av_frame.format);
332332

333333
FiltersConfig filters_config(
334-
av_frame->width,
335-
av_frame->height,
334+
av_frame.width,
335+
av_frame.height,
336336
av_frame_format,
337-
av_frame->sample_aspect_ratio,
337+
av_frame.sample_aspect_ratio,
338338
output_dims.width,
339339
output_dims.height,
340340
output_pixel_format_,
@@ -346,17 +346,17 @@ CpuDeviceInterface::convert_av_frame_to_tensor_using_filter_graph(
346346
std::make_unique<FilterGraph>(filters_config, video_stream_options_);
347347
prev_filters_config_ = std::move(filters_config);
348348
}
349-
return rgb_av_frame_to_tensor(filter_graph_->convert(av_frame));
349+
return rgb_av_frame_to_tensor(*filter_graph_->convert(av_frame));
350350
}
351351

352352
void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
353-
UniqueAVFrame& src_av_frame,
353+
const AVFrame& src_av_frame,
354354
FrameOutput& frame_output) {
355355
AVSampleFormat src_sample_format =
356-
static_cast<AVSampleFormat>(src_av_frame->format);
356+
static_cast<AVSampleFormat>(src_av_frame.format);
357357
AVSampleFormat out_sample_format = AV_SAMPLE_FMT_FLTP;
358358

359-
int src_sample_rate = src_av_frame->sample_rate;
359+
int src_sample_rate = src_av_frame.sample_rate;
360360
int out_sample_rate =
361361
audio_stream_options_.sample_rate.value_or(src_sample_rate);
362362

@@ -397,10 +397,9 @@ void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
397397
out_sample_rate,
398398
out_num_channels);
399399
}
400-
const UniqueAVFrame& av_frame =
401-
must_convert ? converted_av_frame : src_av_frame;
400+
const AVFrame& av_frame = must_convert ? *converted_av_frame : src_av_frame;
402401

403-
AVSampleFormat format = static_cast<AVSampleFormat>(av_frame->format);
402+
AVSampleFormat format = static_cast<AVSampleFormat>(av_frame.format);
404403
STD_TORCH_CHECK(
405404
format == out_sample_format,
406405
"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(
419418
num_channels,
420419
" instead.");
421420

422-
auto num_samples = av_frame->nb_samples;
421+
auto num_samples = av_frame.nb_samples;
423422

424423
frame_output.data = torch::stable::empty({num_channels, num_samples});
425424

@@ -431,7 +430,7 @@ void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
431430
++channel, output_channel_data += num_bytes_per_channel) {
432431
std::memcpy(
433432
output_channel_data,
434-
av_frame->extended_data[channel],
433+
av_frame.extended_data[channel],
435434
num_bytes_per_channel);
436435
}
437436
}

src/torchcodec/_core/CpuDeviceInterface.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ class CpuDeviceInterface : public DeviceInterface {
4141
override;
4242

4343
void convert_av_frame_to_frame_output(
44-
UniqueAVFrame& av_frame,
44+
const AVFrame& av_frame,
4545
FrameOutput& frame_output,
4646
std::optional<torch::stable::Tensor> pre_allocated_output_tensor)
4747
override;
@@ -59,16 +59,16 @@ class CpuDeviceInterface : public DeviceInterface {
5959

6060
private:
6161
void convert_audio_av_frame_to_frame_output(
62-
UniqueAVFrame& src_av_frame,
62+
const AVFrame& src_av_frame,
6363
FrameOutput& frame_output);
6464

6565
void convert_video_av_frame_to_frame_output(
66-
UniqueAVFrame& av_frame,
66+
const AVFrame& av_frame,
6767
FrameOutput& frame_output,
6868
std::optional<torch::stable::Tensor> pre_allocated_output_tensor);
6969

7070
torch::stable::Tensor convert_av_frame_to_tensor_using_filter_graph(
71-
const UniqueAVFrame& av_frame,
71+
const AVFrame& av_frame,
7272
const FrameDims& output_dims);
7373

7474
ColorConversionLibrary get_color_conversion_library(

0 commit comments

Comments
 (0)