From 5fee250cff7eda6132977908c98da4802022d75a Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 7 Jul 2026 11:03:50 +0100 Subject: [PATCH 1/3] API breakdown - claude plan --- API_breakdown_claude_plan.md | 368 +++++++++++++++++++++++++++++++++++ 1 file changed, 368 insertions(+) create mode 100644 API_breakdown_claude_plan.md diff --git a/API_breakdown_claude_plan.md b/API_breakdown_claude_plan.md new file mode 100644 index 000000000..d3027902d --- /dev/null +++ b/API_breakdown_claude_plan.md @@ -0,0 +1,368 @@ +# Low-level decoding building blocks (Demuxer / Decoder / ColorConverter) + +## Context + +**Why:** `VideoDecoder` is a single coarse box: demux, decode, and color-conversion run +back-to-back on one thread, so frame N's decode can't start until frame N-1's +color-conversion finished. We want to let a user (or an agent doing hill-climbing over +pipeline parameters — thread counts, which stage on which thread, queue depths, batch +sizes, number of parallel videos) build their *own* overlapped pipeline out of +lower-level pieces. Because our PyTorch custom ops release the GIL, exposing the three +stages as separate GIL-free ops lets plain Python threads achieve real parallelism — +**we write no threading in C++**; we hand out building blocks and the user wires them. + +Optimal overlap differs by device (motivating the breakdown): +- **CPU:** color-conversion (swscale) is expensive → overlap `[demux+decode]` ∥ `[color-convert]`. +- **GPU:** color-conversion is a cheap CUDA kernel → overlap `[demux]` ∥ `[decode+color-convert]`. + +**Second bird:** users have asked for YUV output instead of RGB. Decode already produces +YUV; RGB only appears after color-conversion. Exposing color-conversion as its own step +lets us offer a `ColorConverter(output="yuv420p")` mode. + +**Non-goals:** `VideoDecoder` and `SingleStreamDecoder`'s public op surface stay +behavior-identical. No duplication of demux/decode/convert logic — we *refactor* the +monolith into shared building blocks. First cut is **CPU, single video**; GPU and +multi-video are later phases (design accommodates them but they are out of scope here). + +## Decisions (confirmed with user) + +1. **Intermediate data = both forms.** Default: zero-copy **opaque handles** (raw + `AVPacket`/`AVFrame` pointer laundered through a `[1]` int64 tensor, exactly like the + existing `wrap_decoder_pointer_to_tensor`). Opt-in: **materialized** serializable + forms (`Packet.to_tensor()` → bytes+meta; `Frame.to_tensor()` → native YUV + planes+meta) for crossing a process boundary. Handles are thread-only (raw pointers + are meaningless in another process). +2. **CPU single-video first.** +3. **YUV feature = ColorConverter output mode** (`output="yuv420p"`), producing a + normalized planar YUV via swscale/filtergraph. This is distinct from the raw-frame + `Frame.to_tensor()` materialization (which is native NV12/whatever, for portability). + +## What exists today (grounding) + +- **Demux** is inline `av_read_frame(...)` in `SingleStreamDecoder::decode_av_frame` + (`SingleStreamDecoder.cpp:1532`); it is NOT abstracted. Seek/keyframe/index policy also + lives in `SingleStreamDecoder`: `maybe_seek_to_before_desired_pts` (:1436), + `can_we_avoid_seeking` (:1333), `scan_file_and_update_metadata_and_index` (:280), + `get_pts`/`seconds_to_index_lower_bound`/`_upper_bound` (:1770/:1704/:1738), keyframe + lookup helpers, and the `all_frames`/`key_frames` tables in the private `StreamInfo`. +- **Decode** already has a clean seam: `DeviceInterface::send_packet` / `send_eof_packet` + / `receive_frame` / `flush` (`DeviceInterface.h:114-157`). +- **Color-convert** already has a clean seam: + `DeviceInterface::convert_av_frame_to_frame_output` (`DeviceInterface.h:101`), wrapped + by `SingleStreamDecoder::convert_av_frame_to_frame_output` (:1600, stamps pts/duration). +- Intermediate types: packets via `AutoAVPacket`/`ReferenceAVPacket` + (`FFMPEGCommon.h:157-181`); frames via `UniqueAVFrame` (`FFMPEGCommon.h:84`). +- Boundary convention (`custom_ops.cpp`): C++ object → `[1]` int64 CPU tensor via + `from_blob` + capturing deleter (`wrap_decoder_pointer_to_tensor` :109, unwrap :156). + Ops in `STABLE_TORCH_LIBRARY(torchcodec_ns)` take only Tensor/int/float/bool/str; + frames cross as `std::tuple` (`OpsFrameOutput` :202) with pts/ + duration as 0-D float64 tensors; every tensor-returning op has an `@register_fake` in + `ops.py` (:250+). Decoder-mutating ops are annotated `Tensor(a!)`. + +## Target Python API (new namespace: `torchcodec.pipeline`) + +Building blocks (compose them on your own threads): + +```python +from torchcodec.pipeline import Demuxer, Decoder, ColorConverter + +demuxer = Demuxer("v.mp4", stream_index=0, seek_mode="exact") +decoder = Decoder(demuxer, num_ffmpeg_threads=4) # bound to demuxer (needs codec params, stateful) +converter = ColorConverter(decoder, output="rgb") # or output="yuv420p" <-- YUV feature + +# --- fully manual loop (thread it however you like) --- +for packet in demuxer: # demux -> Packet (opaque handle) + for frame in decoder.decode(packet): # decode -> Frame (opaque YUV handle) + out = converter.convert(frame) # color-convert -> tensor (RGB or YUV) + +# --- opt-in materialization to cross a process boundary --- +blob = packet.to_tensor() # (uint8 bytes, int64 meta[pts,dts,flags,...]) +pkt2 = Packet.from_tensor(*blob) # reconstruct in another process +yuv = frame.to_tensor() # native YUV planes + JSON meta +``` + +Composed convenience (sequential, no threading) so `get_frames_at` / `get_frames_played_at` +keep working — the seek/index intelligence lives in the Demuxer: + +```python +pipe = Pipeline(demuxer, decoder, converter) # thin wire-up of the 3 blocks +fb = pipe.get_frames_at([10, 20, 30]) # reuses argsort/dedup/seek logic +fb = pipe.get_frames_played_at([1.0, 2.0]) +``` + +Binding rules (match the C++ constraints): +- `Decoder` is **hard-bound** to one `Demuxer` (needs codec params; stateful; not + thread-safe within a stream). One demux thread per video. +- `ColorConverter` on **CPU** may be standalone/shared; the `Decoder` arg supplies the + transform/dtype config. (On GPU later it must share the decoder's DeviceInterface.) + +## Design contract & completeness checklist + +**We are the building-block layer, not a scheduler (hard non-goal).** We deliver GIL-free, +individually-callable stages over movable handles; we do NOT ship a thread pool, async +executor, pipeline runner, per-stage concurrency/ordering controls, or IPC machinery. The +user (or an agent hill-climbing over its own threading) writes the scheduling with plain +`threading` + `queue.Queue`. Every stage must therefore be *scheduler-friendly*: +- a stage is a plain callable `handle -> handle` (thin methods, no hidden global state); +- opaque handles are freely movable across threads within a process; +- packets (only) are serializable across a process boundary via `to_tensor()`. + +To make sure the block set is complete for a caller building their own overlapped pipeline, +the design must also cover (fold into Phase 1 unless noted): +- **Handles carry `pts`.** The opaque `Packet`/`Frame` handles must expose `pts`/`duration` + (as `FrameBatch` already does) so a caller running stages out-of-order for throughput can + reorder results by presentation time. Document this reordering responsibility as the + caller's, not ours. +- **Packet reuse.** Allow feeding the same packet(s) to two decoders/converters (e.g. two + transforms) without re-demuxing — clone/refcount semantics on the packet handle. +- **Optional (later phase):** an explicit CPU→GPU transfer building block (pinned memory + + dedicated stream) so H2D copy overlaps compute in a CPU-decode→GPU-train pattern; and a + raw FFmpeg `filter_desc` escape hatch alongside the curated `transform_specs`. + +## User code examples + +### Use-case 1 — single video, overlapped demux+decode ∥ color-convert (CPU strategy) + +The whole point: while thread B color-converts frame N-1, thread A already decodes frame N. +Because the ops release the GIL, these two Python threads run on two cores. + +```python +import threading, queue +from torchcodec.pipeline import Demuxer, Decoder, ColorConverter + +demuxer = Demuxer("video.mp4", stream_index=0) +decoder = Decoder(demuxer, num_ffmpeg_threads=4) # bound to this demuxer +converter = ColorConverter(decoder, output="rgb") + +frame_q = queue.Queue(maxsize=8) # bounded -> backpressure +_SENTINEL = object() +frames_out = [] + +def demux_and_decode(): # thread A: demux + decode + for packet in demuxer: # -> Packet (opaque handle) + for frame in decoder.decode(packet): # -> Frame (opaque YUV handle) + frame_q.put(frame) + for frame in decoder.flush(): # drain the codec's buffered frames + frame_q.put(frame) + frame_q.put(_SENTINEL) + +def color_convert(): # thread B: YUV -> RGB tensor + while (frame := frame_q.get()) is not _SENTINEL: + out = converter.convert(frame) # Frame dataclass (data, pts, duration) + frames_out.append(out) + +a = threading.Thread(target=demux_and_decode) +b = threading.Thread(target=color_convert) +a.start(); b.start(); a.join(); b.join() +``` + +An agent hill-climbs by varying: `num_ffmpeg_threads`, `frame_q` depth, and *which* stages share a +thread. E.g. the **GPU** strategy is one line different — put decode+convert together and +peel off only demux (color-convert is a cheap kernel there): + +```python +def demux_only(): # thread A: demux + for packet in demuxer: + packet_q.put(packet) + packet_q.put(_SENTINEL) + +def decode_and_convert(): # thread B: decode + color-convert + while (packet := packet_q.get()) is not _SENTINEL: + for frame in decoder.decode(packet): + frames_out.append(converter.convert(frame)) +``` + +Index/time access stays a one-liner via the sequential convenience wrapper (seek/index +logic lives in the Demuxer): + +```python +from torchcodec.pipeline import Pipeline +pipe = Pipeline(demuxer, decoder, converter) +fb = pipe.get_frames_at([10, 20, 30]) # FrameBatch, same output as VideoDecoder +fb = pipe.get_frames_played_at([1.0, 2.5]) +``` + +YUV output (the second bird): + +```python +yuv_conv = ColorConverter(decoder, output="yuv420p") +frame = next(iter(decoder.decode(next(iter(demuxer))))) +yuv = yuv_conv.convert(frame) # planar YUV instead of RGB +``` + +### Use-case 2 — multiple videos: N demuxers, N decoders, a shared pool of converters + +One demux+decode thread **per video** (an `AVFormatContext` / decoder is not thread-safe, +and each `Decoder` is bound to its own `Demuxer`). Color-converters are **not** bound to a +video, so a small shared pool drains a common queue; each frame is tagged with its source +`video_id` so results are recoverable per video. + +```python +import threading, queue +from torchcodec.pipeline import Demuxer, Decoder, ColorConverter + +videos = ["a.mp4", "b.mp4", "c.mp4", "d.mp4"] +frame_q = queue.Queue(maxsize=64) # (video_id, Frame) from all decoders +_SENTINEL = object() +results = {i: [] for i in range(len(videos))} # recover frames per source video +results_lock = threading.Lock() + +def demux_and_decode(video_id, path): # one thread per video + demuxer = Demuxer(path) + decoder = Decoder(demuxer, num_ffmpeg_threads=1) # 1 each: many run concurrently + for packet in demuxer: + for frame in decoder.decode(packet): + frame_q.put((video_id, frame)) + for frame in decoder.flush(): + frame_q.put((video_id, frame)) + +def convert_worker(): # pool: not bound to any video + converter = ColorConverter(output="rgb") # standalone; rebuilds config per input format + while True: + item = frame_q.get() + if item is _SENTINEL: + frame_q.put(_SENTINEL) # re-post so peers also stop + return + video_id, frame = item + rgb = converter.convert(frame) + with results_lock: + results[video_id].append(rgb) + +producers = [threading.Thread(target=demux_and_decode, args=(i, p)) + for i, p in enumerate(videos)] +converters = [threading.Thread(target=convert_worker) for _ in range(4)] # tunable pool size + +for t in producers + converters: t.start() +for t in producers: t.join() +frame_q.put(_SENTINEL) # signal converter pool to drain +for t in converters: t.join() +# results[video_id] -> list of RGB frames for that video +``` + +Hill-climbing surface here: number of concurrent videos, `num_ffmpeg_threads` per decoder, +converter-pool size, and queue depth — an agent searches these for peak throughput. +(A standalone CPU `ColorConverter` recreates its swscale/filtergraph config when the input +frame format changes, so one pool can serve videos with differing formats. On GPU a +converter must instead be created per-decoder — deferred to the GPU phase.) + +## C++ refactor (phased, no logic duplication) + +### Phase 0 — extract `StreamDemuxer` (behavior-preserving, no new ops) +New `Demuxer.{h,cpp}` class `StreamDemuxer` (mechanism, not policy). Move — verbatim, not +rewritten — out of `SingleStreamDecoder`: +- ownership of `format_context_` + `avio_context_holder_`, the file/tensor/file-like + constructors + `initialize_decoder`; +- `select_stream(optional)` = the `av_find_best_stream` + `discard=AVDISCARD_ALL` + half of `add_stream` (:471,:541-545); exposes selected `AVStream*`/`AVCodecParameters*`; +- `scan_and_build_index()` / `read_custom_frame_mappings(...)` (:280,:371) and the + `all_frames`/`key_frames` tables (move `FrameInfo`); +- metadata + pts↔index conversions (`get_pts`, `seconds_to_index_lower_bound`/ + `_upper_bound`, keyframe lookup, key-frame identifier); +- `bool next_packet(ReferenceAVPacket&)` = the `av_read_frame` + stream-filter loop + (:1531-1552); +- `seek_to_keyframe_before_pts(int64_t pts)` = keyframe-corrected `avformat_seek_file` + (:1453-1473) **without** the flush (flush is the decoder's job); +- seek-avoidance as a **pure query** taking decode feedback as args: + `bool can_avoid_seek(cursor, last_decoded_pts, reorder_buffer_size, in_flight_frames)` + — index half stays in demuxer; codec half (`has_b_frames`, `thread_count`) passed in. +Also move `get_pts_or_dts` (anon ns :29-33) into `FFMPEGCommon.h`. +`SingleStreamDecoder` now delegates to `StreamDemuxer`. **All existing tests stay green** — +this is the de-risking step. + +### Phase 1 — extract `PacketDecoder` + `ColorConverter`, then opaque handles + CPU thread API +- New `StreamDecoder.{h,cpp}` class `PacketDecoder`: absorbs codec-context construction + (:504-532), holds the `DeviceInterface` + `SharedAVCodecContext`; API = the existing + seam (`send_packet`/`send_eof`/`receive_frame`/`flush`) plus `reorder_buffer_size()` / + `in_flight_frames()` accessors to feed `can_avoid_seek`. +- New `ColorConverter.{h,cpp}`: wraps `convert_av_frame_to_frame_output` + pts/duration + stamping (:1600); gains a YUV output path (`output="yuv420p"`). +- `SingleStreamDecoder::decode_av_frame` collapses to orchestration over the three + sub-objects (public methods keep their bodies → VideoDecoder unchanged). +- New `DecodeHandles.{h,cpp}`: opaque `Packet`/`Frame` handles via the existing launder + pattern — heap `UniqueAVPacket`/`UniqueAVFrame` → `from_blob` `[1]` int64 tensor with a + deleter running `av_packet_free`/`av_frame_free`. Each emitted packet is its own + `av_packet_ref` (the `AutoAVPacket` reuse optimization can't cross a thread boundary). +- Materialization: `packet_to_tensors`/`packet_from_tensors`; `frame_to_yuv` (packed YUV + tensor + JSON meta — prefer a single packed tensor over `Tensor[]`; verify stable-ABI + `Tensor[]` support before relying on it). + +### New custom ops (all in `STABLE_TORCH_LIBRARY(torchcodec_ns)`, `@register_fake` where they return tensors) +- Demuxer: `create_demuxer_from_file/_from_tensor/_from_file_like`, `demuxer_select_stream`, + `demuxer_scan`, `demuxer_get_container_json_metadata`/`_stream_json_metadata`/ + `_key_frame_indices`, `demuxer_next_packet -> (Tensor, bool)`, + `demuxer_seek_to_keyframe_before_pts`, `demuxer_seek_for_index`, + `demuxer_index_for_pts`/`demuxer_pts_for_index`. + Packet/packet-group handles expose `pts`/`dts`/`duration` accessors so a caller can + reorder out-of-order results. +- Packet: `packet_to_tensors -> (Tensor,Tensor)`, `packet_from_tensors -> Tensor`. +- Decoder: `create_decoder_for_stream(demuxer, *, num_threads, transform_specs, + output_dtype)` — the op keeps the existing `num_threads` name (as `_add_video_stream` + does); the Python `Decoder` wrapper exposes it as `num_ffmpeg_threads` (default `1`) to + match `VideoDecoder`. `decoder_send_packet -> int`, `decoder_send_eof -> int`, + `decoder_receive_frame -> (Tensor, int)` (handle, status), `decoder_flush`. +- Color: `create_color_converter(decoder, *, transform_specs, output_dtype, output)`, + `convert_frame_to_rgb -> (Tensor,Tensor,Tensor)`, `frame_to_yuv -> (Tensor, str)`. + +### Hard problems (flagged; only #2–#4 in scope for CPU-first) +1. **GPU frame-surface lifetime (deferred to GPU phase).** NVDEC keeps one surface mapped, + unmapping on the next `receive_frame` (`BetaCudaDeviceInterface.cpp:679`) → an opaque + GPU frame handle dies on the next decode. Fix later via forced D2D YUV copy, or keep + decode+convert coupled and split only demux, or a surface pool. +2. **Split seek-avoidance feedback.** `can_avoid_seek` needs the async "last decoded pts" + from the decode thread. Kept as a pure demuxer query; the Python scheduler feeds back + the latest decoded pts. Random access drives explicit `demuxer_seek_for_index` + + `decoder_flush`; the optimization is for sequential runs. +3. **Flush-on-seek ordering across threads.** Demux-thread seek invalidates decode-thread + state. `seek_to_keyframe_before_pts` deliberately does NOT flush; use a seek-epoch tag + on packets so the decode thread flushes when the epoch changes. +4. **`AVFormatContext` not thread-safe** → exactly one demux thread per `Demuxer`; + parallelism is demux(1)→decode(1)→convert(N) within a video, or one `Demuxer` per video. + +## Files + +- Change: `SingleStreamDecoder.{h,cpp}` (delegate to new classes), `custom_ops.cpp` (new + ops + wrap/unwrap for Packet/Frame), `ops.py` (aliases + `@register_fake`), + `FFMPEGCommon.h` (move `get_pts_or_dts`), `_core/CMakeLists`+BUCK sources (new files). +- Add: `_core/Demuxer.{h,cpp}`, `_core/StreamDecoder.{h,cpp}`, `_core/ColorConverter.{h,cpp}`, + `_core/DecodeHandles.{h,cpp}`, `pipeline/{__init__,_demuxer,_decoder,_color_converter, + _pipeline,_frame}.py`. +- Untouched: `decoders/_video_decoder.py`, all `VideoDecoder` ops. + +## Verification + +- **Phase 0/1 safety net:** full `pytest` (incl. `-m ""`) and the C++ gtest suite stay + green after each extraction — proves behavior preservation for `VideoDecoder`. Pay + special attention to seek-heuristic edge cases (`INT64_MIN` cursor init, mkv/webm index + quirks noted at `SingleStreamDecoder.h:365-376`). +- **New API correctness:** new tests assert the building-block pipeline returns + frame-identical output to `VideoDecoder.get_frames_at`/`get_frames_played_at` on the + same assets (`assert_frames_equal`); round-trip `Packet.to_tensor()/from_tensor()` and + `frame_to_yuv`; `ColorConverter(output="yuv420p")` shape/format checks. +- **Overlap win (the point):** a benchmark script decoding a video (a) monolithically vs + (b) with demux+decode on one thread and color-convert on another (Python + `threading` + `queue.Queue`), showing wall-clock speedup on CPU — confirms the GIL is + released and the breakdown pays off. +- Run `pre-commit run --all-files` and `mypy` before finishing. + +## Follow-ups (additive, non-breaking — NOT in scope now) + +These extend the design without changing anything above: new demuxer methods and a new +one-shot decode op, layered on top of the per-packet building blocks once they exist. Do +NOT design for them now. + +- **Packet-group / clip granularity.** Today's units are one packet in / one-or-more frames + out (fine-grained, ideal for demux∥decode overlap within a single video). As a follow-up, + add a multi-packet "packet group" handle plus a one-shot + `decoder_decode_packets(decoder, packet_group) -> frame-group-handle` that drains a whole + group in a single GIL-free call — the natural "decode this clip" unit for a + many-videos/one-clip-each scheduling pattern. Purely additive: the per-packet + `send_packet`/`receive_frame` path is unchanged. +- **Windowed / chunked demux.** Demuxer helpers that assemble those groups: + `demuxer_next_packets(demuxer, num_packets)` and + `demuxer_packets_for_time_range(demuxer, start, stop)` (seek + collect the covering + packets into one group handle). Decomposes what `get_frames_played_in_range` does + internally into a reusable "collect packets" block. +- **Optional building blocks** (also later): an explicit CPU→GPU transfer block (pinned + memory + dedicated stream) so H2D copy overlaps compute; a raw FFmpeg `filter_desc` + escape hatch alongside the curated `transform_specs`; packet clone/reuse so the same + packets can feed two decoders/converters without re-demuxing. From faaceb0102eb6b33a9e22fef842584a7bc433caf Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 8 Jul 2026 10:27:15 +0100 Subject: [PATCH 2/3] WIP --- API_breakdown_claude_plan.md | 368 ------------------ benchmarks/bench_blocks.py | 211 ++++++++++ src/torchcodec/_core/ColorConverter.cpp | 55 +++ src/torchcodec/_core/ColorConverter.h | 30 ++ src/torchcodec/_core/CpuDeviceInterface.cpp | 11 +- src/torchcodec/_core/Demuxer.cpp | 94 +++++ src/torchcodec/_core/Demuxer.h | 58 +++ src/torchcodec/_core/FFMPEGCommon.cpp | 13 + src/torchcodec/_core/FFMPEGCommon.h | 5 + src/torchcodec/_core/PacketDecoder.cpp | 93 +++++ src/torchcodec/_core/PacketDecoder.h | 58 +++ src/torchcodec/_core/SingleStreamDecoder.cpp | 78 ++-- src/torchcodec/_core/custom_ops.cpp | 222 +++++++++++ src/torchcodec/_core/ops.py | 24 ++ src/torchcodec/_core/sources.bzl | 3 + src/torchcodec/decoders/_blocks/__init__.py | 23 ++ .../decoders/_blocks/_color_converter.py | 45 +++ src/torchcodec/decoders/_blocks/_demuxer.py | 48 +++ src/torchcodec/decoders/_blocks/_frame.py | 42 ++ .../decoders/_blocks/_packet_decoder.py | 60 +++ test/test_decoders.py | 186 +++++++++ 21 files changed, 1302 insertions(+), 425 deletions(-) delete mode 100644 API_breakdown_claude_plan.md create mode 100644 benchmarks/bench_blocks.py create mode 100644 src/torchcodec/_core/ColorConverter.cpp create mode 100644 src/torchcodec/_core/ColorConverter.h create mode 100644 src/torchcodec/_core/Demuxer.cpp create mode 100644 src/torchcodec/_core/Demuxer.h create mode 100644 src/torchcodec/_core/PacketDecoder.cpp create mode 100644 src/torchcodec/_core/PacketDecoder.h create mode 100644 src/torchcodec/decoders/_blocks/__init__.py create mode 100644 src/torchcodec/decoders/_blocks/_color_converter.py create mode 100644 src/torchcodec/decoders/_blocks/_demuxer.py create mode 100644 src/torchcodec/decoders/_blocks/_frame.py create mode 100644 src/torchcodec/decoders/_blocks/_packet_decoder.py diff --git a/API_breakdown_claude_plan.md b/API_breakdown_claude_plan.md deleted file mode 100644 index d3027902d..000000000 --- a/API_breakdown_claude_plan.md +++ /dev/null @@ -1,368 +0,0 @@ -# Low-level decoding building blocks (Demuxer / Decoder / ColorConverter) - -## Context - -**Why:** `VideoDecoder` is a single coarse box: demux, decode, and color-conversion run -back-to-back on one thread, so frame N's decode can't start until frame N-1's -color-conversion finished. We want to let a user (or an agent doing hill-climbing over -pipeline parameters — thread counts, which stage on which thread, queue depths, batch -sizes, number of parallel videos) build their *own* overlapped pipeline out of -lower-level pieces. Because our PyTorch custom ops release the GIL, exposing the three -stages as separate GIL-free ops lets plain Python threads achieve real parallelism — -**we write no threading in C++**; we hand out building blocks and the user wires them. - -Optimal overlap differs by device (motivating the breakdown): -- **CPU:** color-conversion (swscale) is expensive → overlap `[demux+decode]` ∥ `[color-convert]`. -- **GPU:** color-conversion is a cheap CUDA kernel → overlap `[demux]` ∥ `[decode+color-convert]`. - -**Second bird:** users have asked for YUV output instead of RGB. Decode already produces -YUV; RGB only appears after color-conversion. Exposing color-conversion as its own step -lets us offer a `ColorConverter(output="yuv420p")` mode. - -**Non-goals:** `VideoDecoder` and `SingleStreamDecoder`'s public op surface stay -behavior-identical. No duplication of demux/decode/convert logic — we *refactor* the -monolith into shared building blocks. First cut is **CPU, single video**; GPU and -multi-video are later phases (design accommodates them but they are out of scope here). - -## Decisions (confirmed with user) - -1. **Intermediate data = both forms.** Default: zero-copy **opaque handles** (raw - `AVPacket`/`AVFrame` pointer laundered through a `[1]` int64 tensor, exactly like the - existing `wrap_decoder_pointer_to_tensor`). Opt-in: **materialized** serializable - forms (`Packet.to_tensor()` → bytes+meta; `Frame.to_tensor()` → native YUV - planes+meta) for crossing a process boundary. Handles are thread-only (raw pointers - are meaningless in another process). -2. **CPU single-video first.** -3. **YUV feature = ColorConverter output mode** (`output="yuv420p"`), producing a - normalized planar YUV via swscale/filtergraph. This is distinct from the raw-frame - `Frame.to_tensor()` materialization (which is native NV12/whatever, for portability). - -## What exists today (grounding) - -- **Demux** is inline `av_read_frame(...)` in `SingleStreamDecoder::decode_av_frame` - (`SingleStreamDecoder.cpp:1532`); it is NOT abstracted. Seek/keyframe/index policy also - lives in `SingleStreamDecoder`: `maybe_seek_to_before_desired_pts` (:1436), - `can_we_avoid_seeking` (:1333), `scan_file_and_update_metadata_and_index` (:280), - `get_pts`/`seconds_to_index_lower_bound`/`_upper_bound` (:1770/:1704/:1738), keyframe - lookup helpers, and the `all_frames`/`key_frames` tables in the private `StreamInfo`. -- **Decode** already has a clean seam: `DeviceInterface::send_packet` / `send_eof_packet` - / `receive_frame` / `flush` (`DeviceInterface.h:114-157`). -- **Color-convert** already has a clean seam: - `DeviceInterface::convert_av_frame_to_frame_output` (`DeviceInterface.h:101`), wrapped - by `SingleStreamDecoder::convert_av_frame_to_frame_output` (:1600, stamps pts/duration). -- Intermediate types: packets via `AutoAVPacket`/`ReferenceAVPacket` - (`FFMPEGCommon.h:157-181`); frames via `UniqueAVFrame` (`FFMPEGCommon.h:84`). -- Boundary convention (`custom_ops.cpp`): C++ object → `[1]` int64 CPU tensor via - `from_blob` + capturing deleter (`wrap_decoder_pointer_to_tensor` :109, unwrap :156). - Ops in `STABLE_TORCH_LIBRARY(torchcodec_ns)` take only Tensor/int/float/bool/str; - frames cross as `std::tuple` (`OpsFrameOutput` :202) with pts/ - duration as 0-D float64 tensors; every tensor-returning op has an `@register_fake` in - `ops.py` (:250+). Decoder-mutating ops are annotated `Tensor(a!)`. - -## Target Python API (new namespace: `torchcodec.pipeline`) - -Building blocks (compose them on your own threads): - -```python -from torchcodec.pipeline import Demuxer, Decoder, ColorConverter - -demuxer = Demuxer("v.mp4", stream_index=0, seek_mode="exact") -decoder = Decoder(demuxer, num_ffmpeg_threads=4) # bound to demuxer (needs codec params, stateful) -converter = ColorConverter(decoder, output="rgb") # or output="yuv420p" <-- YUV feature - -# --- fully manual loop (thread it however you like) --- -for packet in demuxer: # demux -> Packet (opaque handle) - for frame in decoder.decode(packet): # decode -> Frame (opaque YUV handle) - out = converter.convert(frame) # color-convert -> tensor (RGB or YUV) - -# --- opt-in materialization to cross a process boundary --- -blob = packet.to_tensor() # (uint8 bytes, int64 meta[pts,dts,flags,...]) -pkt2 = Packet.from_tensor(*blob) # reconstruct in another process -yuv = frame.to_tensor() # native YUV planes + JSON meta -``` - -Composed convenience (sequential, no threading) so `get_frames_at` / `get_frames_played_at` -keep working — the seek/index intelligence lives in the Demuxer: - -```python -pipe = Pipeline(demuxer, decoder, converter) # thin wire-up of the 3 blocks -fb = pipe.get_frames_at([10, 20, 30]) # reuses argsort/dedup/seek logic -fb = pipe.get_frames_played_at([1.0, 2.0]) -``` - -Binding rules (match the C++ constraints): -- `Decoder` is **hard-bound** to one `Demuxer` (needs codec params; stateful; not - thread-safe within a stream). One demux thread per video. -- `ColorConverter` on **CPU** may be standalone/shared; the `Decoder` arg supplies the - transform/dtype config. (On GPU later it must share the decoder's DeviceInterface.) - -## Design contract & completeness checklist - -**We are the building-block layer, not a scheduler (hard non-goal).** We deliver GIL-free, -individually-callable stages over movable handles; we do NOT ship a thread pool, async -executor, pipeline runner, per-stage concurrency/ordering controls, or IPC machinery. The -user (or an agent hill-climbing over its own threading) writes the scheduling with plain -`threading` + `queue.Queue`. Every stage must therefore be *scheduler-friendly*: -- a stage is a plain callable `handle -> handle` (thin methods, no hidden global state); -- opaque handles are freely movable across threads within a process; -- packets (only) are serializable across a process boundary via `to_tensor()`. - -To make sure the block set is complete for a caller building their own overlapped pipeline, -the design must also cover (fold into Phase 1 unless noted): -- **Handles carry `pts`.** The opaque `Packet`/`Frame` handles must expose `pts`/`duration` - (as `FrameBatch` already does) so a caller running stages out-of-order for throughput can - reorder results by presentation time. Document this reordering responsibility as the - caller's, not ours. -- **Packet reuse.** Allow feeding the same packet(s) to two decoders/converters (e.g. two - transforms) without re-demuxing — clone/refcount semantics on the packet handle. -- **Optional (later phase):** an explicit CPU→GPU transfer building block (pinned memory + - dedicated stream) so H2D copy overlaps compute in a CPU-decode→GPU-train pattern; and a - raw FFmpeg `filter_desc` escape hatch alongside the curated `transform_specs`. - -## User code examples - -### Use-case 1 — single video, overlapped demux+decode ∥ color-convert (CPU strategy) - -The whole point: while thread B color-converts frame N-1, thread A already decodes frame N. -Because the ops release the GIL, these two Python threads run on two cores. - -```python -import threading, queue -from torchcodec.pipeline import Demuxer, Decoder, ColorConverter - -demuxer = Demuxer("video.mp4", stream_index=0) -decoder = Decoder(demuxer, num_ffmpeg_threads=4) # bound to this demuxer -converter = ColorConverter(decoder, output="rgb") - -frame_q = queue.Queue(maxsize=8) # bounded -> backpressure -_SENTINEL = object() -frames_out = [] - -def demux_and_decode(): # thread A: demux + decode - for packet in demuxer: # -> Packet (opaque handle) - for frame in decoder.decode(packet): # -> Frame (opaque YUV handle) - frame_q.put(frame) - for frame in decoder.flush(): # drain the codec's buffered frames - frame_q.put(frame) - frame_q.put(_SENTINEL) - -def color_convert(): # thread B: YUV -> RGB tensor - while (frame := frame_q.get()) is not _SENTINEL: - out = converter.convert(frame) # Frame dataclass (data, pts, duration) - frames_out.append(out) - -a = threading.Thread(target=demux_and_decode) -b = threading.Thread(target=color_convert) -a.start(); b.start(); a.join(); b.join() -``` - -An agent hill-climbs by varying: `num_ffmpeg_threads`, `frame_q` depth, and *which* stages share a -thread. E.g. the **GPU** strategy is one line different — put decode+convert together and -peel off only demux (color-convert is a cheap kernel there): - -```python -def demux_only(): # thread A: demux - for packet in demuxer: - packet_q.put(packet) - packet_q.put(_SENTINEL) - -def decode_and_convert(): # thread B: decode + color-convert - while (packet := packet_q.get()) is not _SENTINEL: - for frame in decoder.decode(packet): - frames_out.append(converter.convert(frame)) -``` - -Index/time access stays a one-liner via the sequential convenience wrapper (seek/index -logic lives in the Demuxer): - -```python -from torchcodec.pipeline import Pipeline -pipe = Pipeline(demuxer, decoder, converter) -fb = pipe.get_frames_at([10, 20, 30]) # FrameBatch, same output as VideoDecoder -fb = pipe.get_frames_played_at([1.0, 2.5]) -``` - -YUV output (the second bird): - -```python -yuv_conv = ColorConverter(decoder, output="yuv420p") -frame = next(iter(decoder.decode(next(iter(demuxer))))) -yuv = yuv_conv.convert(frame) # planar YUV instead of RGB -``` - -### Use-case 2 — multiple videos: N demuxers, N decoders, a shared pool of converters - -One demux+decode thread **per video** (an `AVFormatContext` / decoder is not thread-safe, -and each `Decoder` is bound to its own `Demuxer`). Color-converters are **not** bound to a -video, so a small shared pool drains a common queue; each frame is tagged with its source -`video_id` so results are recoverable per video. - -```python -import threading, queue -from torchcodec.pipeline import Demuxer, Decoder, ColorConverter - -videos = ["a.mp4", "b.mp4", "c.mp4", "d.mp4"] -frame_q = queue.Queue(maxsize=64) # (video_id, Frame) from all decoders -_SENTINEL = object() -results = {i: [] for i in range(len(videos))} # recover frames per source video -results_lock = threading.Lock() - -def demux_and_decode(video_id, path): # one thread per video - demuxer = Demuxer(path) - decoder = Decoder(demuxer, num_ffmpeg_threads=1) # 1 each: many run concurrently - for packet in demuxer: - for frame in decoder.decode(packet): - frame_q.put((video_id, frame)) - for frame in decoder.flush(): - frame_q.put((video_id, frame)) - -def convert_worker(): # pool: not bound to any video - converter = ColorConverter(output="rgb") # standalone; rebuilds config per input format - while True: - item = frame_q.get() - if item is _SENTINEL: - frame_q.put(_SENTINEL) # re-post so peers also stop - return - video_id, frame = item - rgb = converter.convert(frame) - with results_lock: - results[video_id].append(rgb) - -producers = [threading.Thread(target=demux_and_decode, args=(i, p)) - for i, p in enumerate(videos)] -converters = [threading.Thread(target=convert_worker) for _ in range(4)] # tunable pool size - -for t in producers + converters: t.start() -for t in producers: t.join() -frame_q.put(_SENTINEL) # signal converter pool to drain -for t in converters: t.join() -# results[video_id] -> list of RGB frames for that video -``` - -Hill-climbing surface here: number of concurrent videos, `num_ffmpeg_threads` per decoder, -converter-pool size, and queue depth — an agent searches these for peak throughput. -(A standalone CPU `ColorConverter` recreates its swscale/filtergraph config when the input -frame format changes, so one pool can serve videos with differing formats. On GPU a -converter must instead be created per-decoder — deferred to the GPU phase.) - -## C++ refactor (phased, no logic duplication) - -### Phase 0 — extract `StreamDemuxer` (behavior-preserving, no new ops) -New `Demuxer.{h,cpp}` class `StreamDemuxer` (mechanism, not policy). Move — verbatim, not -rewritten — out of `SingleStreamDecoder`: -- ownership of `format_context_` + `avio_context_holder_`, the file/tensor/file-like - constructors + `initialize_decoder`; -- `select_stream(optional)` = the `av_find_best_stream` + `discard=AVDISCARD_ALL` - half of `add_stream` (:471,:541-545); exposes selected `AVStream*`/`AVCodecParameters*`; -- `scan_and_build_index()` / `read_custom_frame_mappings(...)` (:280,:371) and the - `all_frames`/`key_frames` tables (move `FrameInfo`); -- metadata + pts↔index conversions (`get_pts`, `seconds_to_index_lower_bound`/ - `_upper_bound`, keyframe lookup, key-frame identifier); -- `bool next_packet(ReferenceAVPacket&)` = the `av_read_frame` + stream-filter loop - (:1531-1552); -- `seek_to_keyframe_before_pts(int64_t pts)` = keyframe-corrected `avformat_seek_file` - (:1453-1473) **without** the flush (flush is the decoder's job); -- seek-avoidance as a **pure query** taking decode feedback as args: - `bool can_avoid_seek(cursor, last_decoded_pts, reorder_buffer_size, in_flight_frames)` - — index half stays in demuxer; codec half (`has_b_frames`, `thread_count`) passed in. -Also move `get_pts_or_dts` (anon ns :29-33) into `FFMPEGCommon.h`. -`SingleStreamDecoder` now delegates to `StreamDemuxer`. **All existing tests stay green** — -this is the de-risking step. - -### Phase 1 — extract `PacketDecoder` + `ColorConverter`, then opaque handles + CPU thread API -- New `StreamDecoder.{h,cpp}` class `PacketDecoder`: absorbs codec-context construction - (:504-532), holds the `DeviceInterface` + `SharedAVCodecContext`; API = the existing - seam (`send_packet`/`send_eof`/`receive_frame`/`flush`) plus `reorder_buffer_size()` / - `in_flight_frames()` accessors to feed `can_avoid_seek`. -- New `ColorConverter.{h,cpp}`: wraps `convert_av_frame_to_frame_output` + pts/duration - stamping (:1600); gains a YUV output path (`output="yuv420p"`). -- `SingleStreamDecoder::decode_av_frame` collapses to orchestration over the three - sub-objects (public methods keep their bodies → VideoDecoder unchanged). -- New `DecodeHandles.{h,cpp}`: opaque `Packet`/`Frame` handles via the existing launder - pattern — heap `UniqueAVPacket`/`UniqueAVFrame` → `from_blob` `[1]` int64 tensor with a - deleter running `av_packet_free`/`av_frame_free`. Each emitted packet is its own - `av_packet_ref` (the `AutoAVPacket` reuse optimization can't cross a thread boundary). -- Materialization: `packet_to_tensors`/`packet_from_tensors`; `frame_to_yuv` (packed YUV - tensor + JSON meta — prefer a single packed tensor over `Tensor[]`; verify stable-ABI - `Tensor[]` support before relying on it). - -### New custom ops (all in `STABLE_TORCH_LIBRARY(torchcodec_ns)`, `@register_fake` where they return tensors) -- Demuxer: `create_demuxer_from_file/_from_tensor/_from_file_like`, `demuxer_select_stream`, - `demuxer_scan`, `demuxer_get_container_json_metadata`/`_stream_json_metadata`/ - `_key_frame_indices`, `demuxer_next_packet -> (Tensor, bool)`, - `demuxer_seek_to_keyframe_before_pts`, `demuxer_seek_for_index`, - `demuxer_index_for_pts`/`demuxer_pts_for_index`. - Packet/packet-group handles expose `pts`/`dts`/`duration` accessors so a caller can - reorder out-of-order results. -- Packet: `packet_to_tensors -> (Tensor,Tensor)`, `packet_from_tensors -> Tensor`. -- Decoder: `create_decoder_for_stream(demuxer, *, num_threads, transform_specs, - output_dtype)` — the op keeps the existing `num_threads` name (as `_add_video_stream` - does); the Python `Decoder` wrapper exposes it as `num_ffmpeg_threads` (default `1`) to - match `VideoDecoder`. `decoder_send_packet -> int`, `decoder_send_eof -> int`, - `decoder_receive_frame -> (Tensor, int)` (handle, status), `decoder_flush`. -- Color: `create_color_converter(decoder, *, transform_specs, output_dtype, output)`, - `convert_frame_to_rgb -> (Tensor,Tensor,Tensor)`, `frame_to_yuv -> (Tensor, str)`. - -### Hard problems (flagged; only #2–#4 in scope for CPU-first) -1. **GPU frame-surface lifetime (deferred to GPU phase).** NVDEC keeps one surface mapped, - unmapping on the next `receive_frame` (`BetaCudaDeviceInterface.cpp:679`) → an opaque - GPU frame handle dies on the next decode. Fix later via forced D2D YUV copy, or keep - decode+convert coupled and split only demux, or a surface pool. -2. **Split seek-avoidance feedback.** `can_avoid_seek` needs the async "last decoded pts" - from the decode thread. Kept as a pure demuxer query; the Python scheduler feeds back - the latest decoded pts. Random access drives explicit `demuxer_seek_for_index` + - `decoder_flush`; the optimization is for sequential runs. -3. **Flush-on-seek ordering across threads.** Demux-thread seek invalidates decode-thread - state. `seek_to_keyframe_before_pts` deliberately does NOT flush; use a seek-epoch tag - on packets so the decode thread flushes when the epoch changes. -4. **`AVFormatContext` not thread-safe** → exactly one demux thread per `Demuxer`; - parallelism is demux(1)→decode(1)→convert(N) within a video, or one `Demuxer` per video. - -## Files - -- Change: `SingleStreamDecoder.{h,cpp}` (delegate to new classes), `custom_ops.cpp` (new - ops + wrap/unwrap for Packet/Frame), `ops.py` (aliases + `@register_fake`), - `FFMPEGCommon.h` (move `get_pts_or_dts`), `_core/CMakeLists`+BUCK sources (new files). -- Add: `_core/Demuxer.{h,cpp}`, `_core/StreamDecoder.{h,cpp}`, `_core/ColorConverter.{h,cpp}`, - `_core/DecodeHandles.{h,cpp}`, `pipeline/{__init__,_demuxer,_decoder,_color_converter, - _pipeline,_frame}.py`. -- Untouched: `decoders/_video_decoder.py`, all `VideoDecoder` ops. - -## Verification - -- **Phase 0/1 safety net:** full `pytest` (incl. `-m ""`) and the C++ gtest suite stay - green after each extraction — proves behavior preservation for `VideoDecoder`. Pay - special attention to seek-heuristic edge cases (`INT64_MIN` cursor init, mkv/webm index - quirks noted at `SingleStreamDecoder.h:365-376`). -- **New API correctness:** new tests assert the building-block pipeline returns - frame-identical output to `VideoDecoder.get_frames_at`/`get_frames_played_at` on the - same assets (`assert_frames_equal`); round-trip `Packet.to_tensor()/from_tensor()` and - `frame_to_yuv`; `ColorConverter(output="yuv420p")` shape/format checks. -- **Overlap win (the point):** a benchmark script decoding a video (a) monolithically vs - (b) with demux+decode on one thread and color-convert on another (Python - `threading` + `queue.Queue`), showing wall-clock speedup on CPU — confirms the GIL is - released and the breakdown pays off. -- Run `pre-commit run --all-files` and `mypy` before finishing. - -## Follow-ups (additive, non-breaking — NOT in scope now) - -These extend the design without changing anything above: new demuxer methods and a new -one-shot decode op, layered on top of the per-packet building blocks once they exist. Do -NOT design for them now. - -- **Packet-group / clip granularity.** Today's units are one packet in / one-or-more frames - out (fine-grained, ideal for demux∥decode overlap within a single video). As a follow-up, - add a multi-packet "packet group" handle plus a one-shot - `decoder_decode_packets(decoder, packet_group) -> frame-group-handle` that drains a whole - group in a single GIL-free call — the natural "decode this clip" unit for a - many-videos/one-clip-each scheduling pattern. Purely additive: the per-packet - `send_packet`/`receive_frame` path is unchanged. -- **Windowed / chunked demux.** Demuxer helpers that assemble those groups: - `demuxer_next_packets(demuxer, num_packets)` and - `demuxer_packets_for_time_range(demuxer, start, stop)` (seek + collect the covering - packets into one group handle). Decomposes what `get_frames_played_in_range` does - internally into a reusable "collect packets" block. -- **Optional building blocks** (also later): an explicit CPU→GPU transfer block (pinned - memory + dedicated stream) so H2D copy overlaps compute; a raw FFmpeg `filter_desc` - escape hatch alongside the curated `transform_specs`; packet clone/reuse so the same - packets can feed two decoders/converters without re-demuxing. diff --git a/benchmarks/bench_blocks.py b/benchmarks/bench_blocks.py new file mode 100644 index 000000000..e5c48e336 --- /dev/null +++ b/benchmarks/bench_blocks.py @@ -0,0 +1,211 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import queue +import subprocess +import threading +from pathlib import Path +from time import perf_counter_ns + +import psutil +import torch + +from torchcodec.decoders import VideoDecoder +from torchcodec.decoders._blocks import ColorConverter, Demuxer, PacketDecoder + +# Kept minimal on purpose; the filename is derived from exactly these. +_DURATION_S = 10 +_HEIGHT = 720 +_WIDTH = 1280 +_FPS = 30 +_SOURCE = "testsrc2" + + +def make_video() -> str: + """Generate (once) a 720p/10s test clip in /tmp, keyed by its generation + parameters, and return its path. Reused if it already exists.""" + key = f"{_SOURCE}_{_WIDTH}x{_HEIGHT}_{_FPS}fps_{_DURATION_S}s" + path = Path("/tmp") / f"bench_blocks_{key}.mp4" + if not path.exists(): + lavfi = f"{_SOURCE}=size={_WIDTH}x{_HEIGHT}:rate={_FPS}:duration={_DURATION_S}" + subprocess.run( + [ + "ffmpeg", + "-y", + "-f", + "lavfi", + "-i", + lavfi, + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-g", + str(_FPS), + str(path), + ], + check=True, + capture_output=True, + ) + return str(path) + + +# The three decode stages, each a generator transforming an iterator of inputs +# into an iterator of outputs. They compose directly (sequential is just +# convert(decode(demux()))); a thread boundary between any two stages is inserted +# with prefetch(). Where you insert it decides which stages overlap. + + +def _demux(demuxer): + yield from demuxer + + +def _decode(decoder, packets): + for packet in packets: + yield from decoder.decode(packet) + yield from decoder.flush() + + +def _convert(converter, frames): + for frame in frames: + yield converter.convert(frame) + + +def prefetch(upstream, buffer_size=8): + # Run `upstream` (a generator chaining one or more stages) on a background + # thread, yielding its items through a bounded queue. The queue applies + # backpressure: the worker blocks in q.put() when the buffer is full, so it + # only runs ~buffer_size items ahead of the consumer. + q: queue.Queue = queue.Queue(maxsize=buffer_size) + eof = object() + error = [] + + def worker(): + try: + for item in upstream: + q.put(item) + except Exception as e: # surface failures instead of hanging + error.append(e) + finally: + q.put(eof) + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + + def drain(): + while (item := q.get()) is not eof: + yield item + thread.join() # worker enqueued eof and is finishing; make it explicit + if error: + raise error[0] + + return drain() + + +def _consume(frames): + for _ in frames: + pass + + +def _decode_sequential(path): + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + converter = ColorConverter() + _consume(_convert(converter, _decode(decoder, _demux(demuxer)))) + + +def _decode_prefetch_frames(path): + # [demux + decode] on one thread || [color-convert] on another. + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + converter = ColorConverter() + frames = prefetch(_decode(decoder, _demux(demuxer))) + _consume(_convert(converter, frames)) + + +def _decode_prefetch_packets(path): + # [demux] on one thread || [decode + color-convert] on another. + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + converter = ColorConverter() + packets = prefetch(_demux(demuxer)) + _consume(_convert(converter, _decode(decoder, packets))) + + +def _decode_prefetch_packets_and_frames(path): + # [demux] || [decode] || [color-convert], each on its own thread. + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + converter = ColorConverter() + packets = prefetch(_demux(demuxer)) + frames = prefetch(_decode(decoder, packets)) + _consume(_convert(converter, frames)) + + +def _decode_video_decoder(path): + # approximate seek mode to match the blocks + VideoDecoder(path, seek_mode="approximate").get_all_frames() + + +def get_num_frames(path): + return VideoDecoder(path).metadata.num_frames + + +# --------------------------------------------------------------------------- +# Benchmark harness +# --------------------------------------------------------------------------- + + +def bench(f, *args, num_exp=10, warmup=2, **kwargs): + process = psutil.Process() + for _ in range(warmup): + f(*args, **kwargs) + times = [] + cpu_utils = [] + for _ in range(num_exp): + process.cpu_percent(interval=None) # reset the measurement window + start = perf_counter_ns() + f(*args, **kwargs) + end = perf_counter_ns() + cpu_utils.append(process.cpu_percent(interval=None)) # since reset + times.append(end - start) + return torch.tensor(times).float(), torch.tensor(cpu_utils).float() + + +def main(): + path = make_video() + print(f"Video: {path} ({_SOURCE} {_WIDTH}x{_HEIGHT} {_FPS}fps {_DURATION_S}s)\n") + + methods = { + "VideoDecoder": _decode_video_decoder, + "sequential": _decode_sequential, + "demux+decode || cc": _decode_prefetch_frames, + "demux || decode+cc": _decode_prefetch_packets, + "demux || decode || cc": _decode_prefetch_packets_and_frames, + } + + results = {} + for name, fn in methods.items(): + times_ns, cpu = bench(fn, path) + results[name] = { + "mean_ms": (times_ns / 1e6).mean().item(), + "std_ms": (times_ns / 1e6).std().item(), + "cpu": cpu.mean().item(), + } + + baseline = results["VideoDecoder"]["mean_ms"] + print(f"{'Method':<24}{'Time (ms)':>20}{'CPU %':>10}{'vs VideoDecoder':>18}") + print("-" * 72) + for name, r in results.items(): + time_str = f"{r['mean_ms']:.1f} +/- {r['std_ms']:.1f}" + speedup = baseline / r["mean_ms"] + print(f"{name:<24}{time_str:>20}{r['cpu']:>10.0f}{speedup:>17.2f}x") + + print(f"\nFrames per run: {get_num_frames(path)}") + + +if __name__ == "__main__": + main() diff --git a/src/torchcodec/_core/ColorConverter.cpp b/src/torchcodec/_core/ColorConverter.cpp new file mode 100644 index 000000000..6263af1cd --- /dev/null +++ b/src/torchcodec/_core/ColorConverter.cpp @@ -0,0 +1,55 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "ColorConverter.h" + +#include +#include + +#include "Frame.h" +#include "StreamOptions.h" +#include "Transform.h" + +namespace facebook::torchcodec { + +ColorConverter::ColorConverter( + const StableDevice& device, + std::string_view device_variant) { + device_interface_ = create_device_interface(device, device_variant); + STD_TORCH_CHECK( + device_interface_ != nullptr, + "Failed to create device interface. This should never happen, please report."); + + VideoStreamOptions options; + 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. + std::vector> no_transforms; + device_interface_->initialize_video( + /*av_stream=*/nullptr, + UniqueDecodingAVFormatContext{}, + options, + no_transforms, + /*resized_output_dims=*/std::nullopt); +} + +torch::stable::Tensor ColorConverter::convert(UniqueAVFrame& av_frame) { + FrameOutput frame_output; + device_interface_->convert_av_frame_to_frame_output( + av_frame, frame_output, std::nullopt); + return frame_output.data; +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/ColorConverter.h b/src/torchcodec/_core/ColorConverter.h new file mode 100644 index 000000000..cf6368735 --- /dev/null +++ b/src/torchcodec/_core/ColorConverter.h @@ -0,0 +1,30 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include + +#include "DeviceInterface.h" +#include "FFMPEGCommon.h" +#include "StableABICompat.h" + +namespace facebook::torchcodec { + +class FORCE_PUBLIC_VISIBILITY ColorConverter { + public: + explicit ColorConverter( + const StableDevice& device = StableDevice(kStableCPU), + std::string_view device_variant = "default"); + + torch::stable::Tensor convert(UniqueAVFrame& av_frame); + + private: + std::unique_ptr device_interface_; +}; + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/CpuDeviceInterface.cpp b/src/torchcodec/_core/CpuDeviceInterface.cpp index c8edccab3..2d51a8be7 100644 --- a/src/torchcodec/_core/CpuDeviceInterface.cpp +++ b/src/torchcodec/_core/CpuDeviceInterface.cpp @@ -44,8 +44,15 @@ void CpuDeviceInterface::initialize_video( const VideoStreamOptions& video_stream_options, const std::vector>& transforms, const std::optional& resized_output_dims) { - STD_TORCH_CHECK(av_stream != nullptr, "avStream is null"); - time_base_ = av_stream->time_base; + + // TODO_API_BREAKDOWN 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. + 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/Demuxer.cpp b/src/torchcodec/_core/Demuxer.cpp new file mode 100644 index 000000000..106e074a3 --- /dev/null +++ b/src/torchcodec/_core/Demuxer.cpp @@ -0,0 +1,94 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "Demuxer.h" +#include "StableABICompat.h" + +namespace facebook::torchcodec { + +int read_next_packet( + AVFormatContext* format_context, + int active_stream_index, + ReferenceAVPacket& packet) { + int status = AVSUCCESS; + do { + status = av_read_frame(format_context, packet.get()); + if (status == AVERROR_EOF) { + return AVERROR_EOF; + } + if (status < AVSUCCESS) { + return status; + } + } while (packet->stream_index != active_stream_index); + return AVSUCCESS; +} + +Demuxer::Demuxer( + const std::string& file_path, + std::optional stream_index) { + set_ffmpeg_log_level(); + + AVFormatContext* raw_context = nullptr; + int status = + avformat_open_input(&raw_context, file_path.c_str(), nullptr, nullptr); + STD_TORCH_CHECK( + status == 0, + "Could not open input file: " + file_path + " " + + get_ffmpeg_error_string_from_error_code(status)); + STD_TORCH_CHECK(raw_context != nullptr, "Failed to allocate AVFormatContext"); + format_context_.reset(raw_context); + + status = avformat_find_stream_info(format_context_.get(), nullptr); + STD_TORCH_CHECK( + status >= 0, + "Failed to find stream info: ", + get_ffmpeg_error_string_from_error_code(status)); + + active_stream_index_ = av_find_best_stream( + format_context_.get(), + AVMEDIA_TYPE_VIDEO, + stream_index.value_or(-1), + /*related_stream=*/-1, + /*decoder_ret=*/nullptr, + /*flags=*/0); + STD_TORCH_CHECK( + active_stream_index_ >= 0, + "No valid video stream found in input file (requested index ", + stream_index.value_or(-1), + ")."); + stream_ = format_context_->streams[active_stream_index_]; + + // We only need packets from the active stream, so tell FFmpeg to discard the + // others. Note av_read_frame() may still return some of them under certain + // conditions, which is why read_next_packet() also filters by stream index. + for (unsigned int i = 0; i < format_context_->nb_streams; ++i) { + if (i != static_cast(active_stream_index_)) { + format_context_->streams[i]->discard = AVDISCARD_ALL; + } + } +} + +AVPacket* 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; + } + STD_TORCH_CHECK( + status >= AVSUCCESS, + "Could not read frame from input file: ", + get_ffmpeg_error_string_from_error_code(status)); + + // 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(); + STD_TORCH_CHECK(owned != nullptr, "Failed to allocate AVPacket"); + av_packet_move_ref(owned, packet.get()); + return owned; +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/Demuxer.h b/src/torchcodec/_core/Demuxer.h new file mode 100644 index 000000000..e0cfbed3e --- /dev/null +++ b/src/torchcodec/_core/Demuxer.h @@ -0,0 +1,58 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include + +#include "FFMPEGCommon.h" +#include "StableABICompat.h" + +namespace facebook::torchcodec { + +// Reads the next packet belonging to active_stream_index from format_context +// into `packet`. Returns AVSUCCESS with `packet` filled, AVERROR_EOF at end of +// stream, or a negative error code otherwise. Shared by Demuxer and +// SingleStreamDecoder so the demux + stream-filter loop lives in one place. +int read_next_packet( + AVFormatContext* format_context, + int active_stream_index, + ReferenceAVPacket& packet); + +// Demux building block: owns an AVFormatContext, selects one video stream, and +// yields its (compressed) packets. Does no decoding. Not thread-safe. +class FORCE_PUBLIC_VISIBILITY Demuxer { + public: + explicit 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(); + + AVStream* active_stream() const { + return stream_; + } + + int active_stream_index() const { + return active_stream_index_; + } + + const UniqueDecodingAVFormatContext& format_context() const { + return format_context_; + } + + private: + UniqueDecodingAVFormatContext format_context_; + int active_stream_index_ = -1; + AVStream* stream_ = nullptr; + AutoAVPacket auto_packet_; +}; + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/FFMPEGCommon.cpp b/src/torchcodec/_core/FFMPEGCommon.cpp index ec7b423ce..b7b6c1550 100644 --- a/src/torchcodec/_core/FFMPEGCommon.cpp +++ b/src/torchcodec/_core/FFMPEGCommon.cpp @@ -62,6 +62,19 @@ int64_t get_duration(const UniqueAVFrame& av_frame) { #endif } +// Some videos aren't properly encoded and do not specify pts values for +// packets, and thus for frames. Unset values correspond to INT64_MIN. When that +// happens, we fall back to the dts value which hopefully exists and is correct. +// Accessing AVFrames' and AVPackets' pts values should **always** go through +// these helpers. +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; +} + void set_duration(const UniqueAVFrame& av_frame, int64_t duration) { #if LIBAVUTIL_VERSION_MAJOR < 58 av_frame->pkt_duration = duration; diff --git a/src/torchcodec/_core/FFMPEGCommon.h b/src/torchcodec/_core/FFMPEGCommon.h index 3bcaae45f..88bb04ceb 100644 --- a/src/torchcodec/_core/FFMPEGCommon.h +++ b/src/torchcodec/_core/FFMPEGCommon.h @@ -209,6 +209,11 @@ std::string get_ffmpeg_error_string_from_error_code(int error_code); int64_t get_duration(const UniqueAVFrame& frame); void set_duration(const UniqueAVFrame& 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); + const int* get_supported_sample_rates(const AVCodec& av_codec); const AVSampleFormat* get_supported_output_sample_formats( const AVCodec& av_codec); diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp new file mode 100644 index 000000000..4cdab49a8 --- /dev/null +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -0,0 +1,93 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "PacketDecoder.h" + +namespace facebook::torchcodec { + +SharedAVCodecContext create_and_open_codec_context( + AVStream* stream, + const AVCodec* av_codec, + DeviceInterface* device_interface, + std::optional thread_count) { + AVCodecContext* raw_codec_context = avcodec_alloc_context3(av_codec); + STD_TORCH_CHECK( + raw_codec_context != nullptr, "Failed to allocate codec context"); + SharedAVCodecContext codec_context = + make_shared_av_codec_context(raw_codec_context); + + int ret = + avcodec_parameters_to_context(codec_context.get(), stream->codecpar); + STD_TORCH_CHECK(ret == AVSUCCESS, "avcodec_parameters_to_context failed"); + + codec_context->thread_count = thread_count.value_or(0); + codec_context->pkt_timebase = stream->time_base; + + // We must register the hardware device context with the codec context before + // calling avcodec_open2(). Otherwise, decoding will happen on the CPU and not + // the hardware device. + device_interface->register_hardware_device_with_codec(codec_context.get()); + ret = avcodec_open2(codec_context.get(), av_codec, nullptr); + STD_TORCH_CHECK( + ret >= AVSUCCESS, get_ffmpeg_error_string_from_error_code(ret)); + + codec_context->time_base = stream->time_base; + return codec_context; +} + +namespace { +const AVCodec* find_decoder( + AVStream* stream, + DeviceInterface* device_interface) { + const AVCodec* av_codec = avcodec_find_decoder(stream->codecpar->codec_id); + STD_TORCH_CHECK(av_codec != nullptr, "Codec not found"); + if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) { + av_codec = device_interface->find_codec(stream->codecpar->codec_id) + .value_or(av_codec); + } + return av_codec; +} +} // namespace + +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); + STD_TORCH_CHECK( + device_interface_ != nullptr, + "Failed to create device interface. This should never happen, please report."); + + AVStream* stream = demuxer.active_stream(); + time_base_ = stream->time_base; + const AVCodec* av_codec = find_decoder(stream, device_interface_.get()); + codec_context_ = create_and_open_codec_context( + stream, av_codec, device_interface_.get(), ffmpeg_thread_count); + device_interface_->initialize(codec_context_); +} + +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"); + return device_interface_->send_packet(ref); +} + +int PacketDecoder::send_eof() { + return device_interface_->send_eof_packet(); +} + +int PacketDecoder::receive_frame(UniqueAVFrame& av_frame) { + return device_interface_->receive_frame(av_frame); +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/PacketDecoder.h b/src/torchcodec/_core/PacketDecoder.h new file mode 100644 index 000000000..929698ac6 --- /dev/null +++ b/src/torchcodec/_core/PacketDecoder.h @@ -0,0 +1,58 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include +#include + +#include "Demuxer.h" +#include "DeviceInterface.h" +#include "FFMPEGCommon.h" +#include "StableABICompat.h" + +namespace facebook::torchcodec { + +// Creates, configures and opens a codec context for `stream` using `av_codec`, +// registering the hardware device (if any) via `device_interface`. Shared by +// SingleStreamDecoder and PacketDecoder to avoid duplicating codec setup. +SharedAVCodecContext create_and_open_codec_context( + AVStream* stream, + const AVCodec* av_codec, + DeviceInterface* device_interface, + std::optional thread_count); + +// Decode building block: turns compressed packets into decoded (YUV) frames. +// Configured from a Demuxer's active stream; stateful. Not thread-safe. +class FORCE_PUBLIC_VISIBILITY PacketDecoder { + public: + 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). + int send_packet(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) + // if more input is needed, AVERROR_EOF at end, or a negative error code. + int receive_frame(UniqueAVFrame& av_frame); + + // The stream time base, used to convert frame pts/duration to seconds. + AVRational time_base() const { + return time_base_; + } + + private: + std::unique_ptr device_interface_; + SharedAVCodecContext codec_context_; + AVRational time_base_ = {}; +}; + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/SingleStreamDecoder.cpp b/src/torchcodec/_core/SingleStreamDecoder.cpp index bcf49d4ae..b99a25cc3 100644 --- a/src/torchcodec/_core/SingleStreamDecoder.cpp +++ b/src/torchcodec/_core/SingleStreamDecoder.cpp @@ -10,7 +10,9 @@ #include #include #include +#include "Demuxer.h" #include "Metadata.h" +#include "PacketDecoder.h" #include "StableABICompat.h" extern "C" { @@ -18,23 +20,6 @@ extern "C" { } namespace facebook::torchcodec { -namespace { - -// Some videos aren't properly encoded and do not specify pts values for -// packets, and thus for frames. Unset values correspond to INT64_MIN. When that -// happens, we fallback to the dts value which hopefully exists and is correct. -// Accessing AVFrames and AVPackets's pts values should **always** go through -// the helpers below. Then, the "pts" fields in our structs like FrameInfo.pts -// should be interpreted as "pts if it exists, dts otherwise". -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; -} - -} // namespace // -------------------------------------------------------------------------- // CONSTRUCTORS, INITIALIZATION, DESTRUCTORS @@ -506,27 +491,12 @@ void SingleStreamDecoder::add_stream( .value_or(av_codec)); } - AVCodecContext* codec_context = avcodec_alloc_context3(av_codec); - STD_TORCH_CHECK(codec_context != nullptr, "Failed to allocate codec context"); - stream_info.codec_context = make_shared_av_codec_context(codec_context); - - int ret_val = avcodec_parameters_to_context( - stream_info.codec_context.get(), stream_info.stream->codecpar); - STD_TORCH_CHECK(ret_val == AVSUCCESS, "avcodec_parameters_to_context failed"); - - stream_info.codec_context->thread_count = ffmpeg_thread_count.value_or(0); - stream_info.codec_context->pkt_timebase = stream_info.stream->time_base; - - // Note that we must make sure to register the harware device context - // with the codec context before calling avcodec_open2(). Otherwise, decoding - // will happen on the CPU and not the hardware device. - device_interface_->register_hardware_device_with_codec( - stream_info.codec_context.get()); - ret_val = avcodec_open2(stream_info.codec_context.get(), av_codec, nullptr); - STD_TORCH_CHECK( - ret_val >= AVSUCCESS, get_ffmpeg_error_string_from_error_code(ret_val)); - - stream_info.codec_context->time_base = stream_info.stream->time_base; + // Create + configure + open the codec context (shared with PacketDecoder). + stream_info.codec_context = create_and_open_codec_context( + stream_info.stream, + av_codec, + device_interface_.get(), + ffmpeg_thread_count); // Initialize the device interface with the codec context device_interface_->initialize(stream_info.codec_context); @@ -1526,30 +1496,28 @@ UniqueAVFrame SingleStreamDecoder::decode_av_frame( } // We still haven't found the frame we're looking for. So let's read more - // packets and send them to the decoder. + // packets and send them to the decoder. (read_next_packet is shared with + // the Demuxer building block.) ReferenceAVPacket packet(auto_av_packet); - do { - status = av_read_frame(format_context_.get(), packet.get()); - decode_stats_.num_packets_read++; + status = + read_next_packet(format_context_.get(), active_stream_index_, packet); + decode_stats_.num_packets_read++; - if (status == AVERROR_EOF) { - // End of file reached. We must drain the decoder - status = device_interface_->send_eof_packet(); - STD_TORCH_CHECK( - status >= AVSUCCESS, - "Could not flush decoder: ", - get_ffmpeg_error_string_from_error_code(status)); - - reached_eof = true; - break; - } + if (status == AVERROR_EOF) { + // End of file reached. We must drain the decoder + status = device_interface_->send_eof_packet(); + STD_TORCH_CHECK( + status >= AVSUCCESS, + "Could not flush decoder: ", + get_ffmpeg_error_string_from_error_code(status)); + reached_eof = true; + } else { STD_TORCH_CHECK( status >= AVSUCCESS, "Could not read frame from input file: ", get_ffmpeg_error_string_from_error_code(status)); - - } while (packet->stream_index != active_stream_index_); + } if (reached_eof) { // We don't have any more packets to send to the decoder. So keep on diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp index 9a81687b6..06f356d85 100644 --- a/src/torchcodec/_core/custom_ops.cpp +++ b/src/torchcodec/_core/custom_ops.cpp @@ -22,9 +22,12 @@ extern "C" { #include "AVIOFileContext.h" #include "AVIOFileLikeContext.h" #include "AVIOTensorContext.h" +#include "ColorConverter.h" +#include "Demuxer.h" #include "Encoder.h" #include "Logging.h" #include "NVDECCacheConfig.h" +#include "PacketDecoder.h" #include "SingleStreamDecoder.h" #include "StableABICompat.h" #include "ValidationUtils.h" @@ -68,6 +71,19 @@ STABLE_TORCH_LIBRARY(torchcodec_ns, m) { "get_frames_by_pts_in_range_audio(Tensor(a!) decoder, *, float start_seconds, float? stop_seconds) -> (Tensor, Tensor)"); m.def( "get_frames_by_pts(Tensor(a!) decoder, *, Tensor timestamps) -> (Tensor, Tensor, Tensor)"); + m.def( + "_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"); + 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_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"); m.def("get_container_json_metadata(Tensor(a!) decoder) -> str"); @@ -163,6 +179,81 @@ SingleStreamDecoder* unwrap_tensor_to_get_decoder( return 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) { + T* raw = ptr.release(); + auto deleter = [raw](void*) { delete raw; }; + int64_t sizes[] = {static_cast(sizeof(T*))}; + int64_t strides[] = {1}; + return torch::stable::from_blob( + raw, + {sizes, 1}, + {strides, 1}, + StableDevice(kStableCPU), + kStableInt64, + deleter); +} + +template +T* unwrap_tensor_to_pointer(torch::stable::Tensor& tensor) { + STD_TORCH_CHECK(tensor.is_contiguous(), "handle tensor must be contiguous"); + 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(); @@ -735,6 +826,118 @@ OpsAudioFramesOutput get_frames_by_pts_in_range_audio( return make_ops_audio_frames_output(result); } +// ============================== +// Building-block ops (torchcodec.decoders._blocks) +// ============================== + +torch::stable::Tensor _blocks_create_demuxer( + std::string filename, + std::optional stream_index) { + std::optional stream_index_int; + if (stream_index.has_value()) { + stream_index_int = static_cast(stream_index.value()); + } + auto demuxer = std::make_unique(filename, stream_index_int); + return wrap_pointer_to_tensor(std::move(demuxer)); +} + +// (packet_handle, is_eof). On EOF the packet_handle is a dummy tensor that must +// not be used. Native bool avoids per-frame .item() overhead in Python. +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(); + 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); +} + +torch::stable::Tensor _blocks_create_packet_decoder( + torch::stable::Tensor& demuxer, + std::optional num_threads, + std::string device, + std::string device_variant) { + Demuxer* demuxer_ptr = unwrap_tensor_to_pointer(demuxer); + validate_device_interface(device, device_variant); + 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); + return wrap_pointer_to_tensor(std::move(decoder)); +} + +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); + return static_cast(decoder_ptr->send_packet(raw_packet)); +} + +int64_t _blocks_packet_decoder_send_eof(torch::stable::Tensor& decoder) { + PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer(decoder); + return static_cast(decoder_ptr->send_eof()); +} + +// (frame_handle, status, pts_seconds, duration_seconds). status == 0 means a +// frame was produced; nonzero (EAGAIN/EOF) means no frame (dummy handle, +// zeros). pts/duration are stamped here (the decoder knows the stream time +// base) so the ColorConverter doesn't need to be bound to a stream. Native +// scalars avoid per-frame .item() overhead. +using OpsReceiveFrameOutput = + std::tuple; + +OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame( + torch::stable::Tensor& decoder) { + PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer(decoder); + UniqueAVFrame av_frame(av_frame_alloc()); + STD_TORCH_CHECK(av_frame != nullptr, "Failed to allocate AVFrame"); + int status = decoder_ptr->receive_frame(av_frame); + if (status != AVSUCCESS) { + return std::make_tuple( + torch::stable::full({1}, 0, kStableInt64), + static_cast(status), + 0.0, + 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(); + return std::make_tuple( + wrap_frame_pointer_to_tensor(raw_frame), + static_cast(0), + pts_seconds, + 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); + return wrap_pointer_to_tensor(std::move(converter)); +} + +torch::stable::Tensor _blocks_convert_frame( + torch::stable::Tensor& converter, + 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; +} + // For testing only. We need to implement this operation as a core library // function because what we're testing is round-tripping pts values as // double-precision floating point numbers from C++ to Python and back to C++. @@ -1219,6 +1422,10 @@ STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, BackendSelect, m) { m.impl("create_from_file", TORCH_BOX(&create_from_file)); m.impl("create_from_tensor", TORCH_BOX(&create_from_tensor)); m.impl("_create_from_file_like", TORCH_BOX(&_create_from_file_like)); + m.impl("_blocks_create_demuxer", TORCH_BOX(&_blocks_create_demuxer)); + m.impl( + "_blocks_create_color_converter", + TORCH_BOX(&_blocks_create_color_converter)); m.impl( "_get_json_ffmpeg_library_versions", TORCH_BOX(&_get_json_ffmpeg_library_versions)); @@ -1278,6 +1485,21 @@ STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) { "get_frames_by_pts_in_range_audio", TORCH_BOX(&get_frames_by_pts_in_range_audio)); m.impl("get_frames_by_pts", TORCH_BOX(&get_frames_by_pts)); + m.impl( + "_blocks_demuxer_next_packet", TORCH_BOX(&_blocks_demuxer_next_packet)); + m.impl( + "_blocks_create_packet_decoder", + TORCH_BOX(&_blocks_create_packet_decoder)); + m.impl( + "_blocks_packet_decoder_send_packet", + TORCH_BOX(&_blocks_packet_decoder_send_packet)); + m.impl( + "_blocks_packet_decoder_send_eof", + TORCH_BOX(&_blocks_packet_decoder_send_eof)); + m.impl( + "_blocks_packet_decoder_receive_frame", + TORCH_BOX(&_blocks_packet_decoder_receive_frame)); + m.impl("_blocks_convert_frame", TORCH_BOX(&_blocks_convert_frame)); m.impl("_test_frame_pts_equality", TORCH_BOX(&_test_frame_pts_equality)); m.impl( "scan_all_streams_to_update_metadata", diff --git a/src/torchcodec/_core/ops.py b/src/torchcodec/_core/ops.py index 6ccc75b31..e23fca040 100644 --- a/src/torchcodec/_core/ops.py +++ b/src/torchcodec/_core/ops.py @@ -107,6 +107,30 @@ def add_video_stream( torch.ops.torchcodec_ns.get_frames_by_pts_in_range_audio.default ) get_json_metadata = torch.ops.torchcodec_ns.get_json_metadata.default + +# Building-block ops for torchcodec.decoders._blocks (CPU). See +# API_breakdown_claude_plan.md. +_blocks_create_demuxer = torch.ops.torchcodec_ns._blocks_create_demuxer.default +_blocks_demuxer_next_packet = ( + torch.ops.torchcodec_ns._blocks_demuxer_next_packet.default +) +_blocks_create_packet_decoder = ( + torch.ops.torchcodec_ns._blocks_create_packet_decoder.default +) +_blocks_packet_decoder_send_packet = ( + torch.ops.torchcodec_ns._blocks_packet_decoder_send_packet.default +) +_blocks_packet_decoder_send_eof = ( + torch.ops.torchcodec_ns._blocks_packet_decoder_send_eof.default +) +_blocks_packet_decoder_receive_frame = ( + torch.ops.torchcodec_ns._blocks_packet_decoder_receive_frame.default +) +_blocks_create_color_converter = ( + torch.ops.torchcodec_ns._blocks_create_color_converter.default +) +_blocks_convert_frame = torch.ops.torchcodec_ns._blocks_convert_frame.default + _test_frame_pts_equality = torch.ops.torchcodec_ns._test_frame_pts_equality.default _get_container_json_metadata = ( torch.ops.torchcodec_ns.get_container_json_metadata.default diff --git a/src/torchcodec/_core/sources.bzl b/src/torchcodec/_core/sources.bzl index 44c588159..4f4e3f775 100644 --- a/src/torchcodec/_core/sources.bzl +++ b/src/torchcodec/_core/sources.bzl @@ -32,6 +32,9 @@ decoder_core_sources = [ "Frame.cpp", "DeviceInterface.cpp", "CpuDeviceInterface.cpp", + "Demuxer.cpp", + "PacketDecoder.cpp", + "ColorConverter.cpp", "SingleStreamDecoder.cpp", "Encoder.cpp", "ValidationUtils.cpp", diff --git a/src/torchcodec/decoders/_blocks/__init__.py b/src/torchcodec/decoders/_blocks/__init__.py new file mode 100644 index 000000000..9bf9074e3 --- /dev/null +++ b/src/torchcodec/decoders/_blocks/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Private, experimental building-block decode API (CPU only, for now). + +Exposes the three decode stages -- :class:`Demuxer`, :class:`PacketDecoder`, +:class:`ColorConverter` -- as passive, composable, GIL-releasing units, so a +caller can build its own (threaded) decode pipeline and tune how the stages +overlap. The blocks do no threading themselves. + +This is experimental and private; the API may change. See +API_breakdown_claude_plan.md for the design and rationale. +""" + +from ._color_converter import ColorConverter +from ._demuxer import Demuxer +from ._frame import DecodedFrame, Packet +from ._packet_decoder import PacketDecoder + +__all__ = ["Demuxer", "PacketDecoder", "ColorConverter", "Packet", "DecodedFrame"] diff --git a/src/torchcodec/decoders/_blocks/_color_converter.py b/src/torchcodec/decoders/_blocks/_color_converter.py new file mode 100644 index 000000000..69e3bf3c7 --- /dev/null +++ b/src/torchcodec/decoders/_blocks/_color_converter.py @@ -0,0 +1,45 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from torchcodec._core.ops import _blocks_convert_frame, _blocks_create_color_converter +from torchcodec._frame import Frame + +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!! + + +class ColorConverter: + """Color-conversion building block: turns a decoded (YUV) + :class:`DecodedFrame` into an RGB :class:`~torchcodec._frame.Frame` + (CHW, uint8 -- matching ``VideoDecoder``'s default output). + + Not bound to anything: everything it needs (dims, pixel format, colorspace) + comes from the frame itself, so one converter can process frames from any + video. Passive and *not* thread-safe: use one ``ColorConverter`` per thread. + + Note: automatic rotation (from stream side data) is not applied, since this + block is intentionally stream-agnostic. + """ + + def __init__(self): + self._handle = _blocks_create_color_converter() + + def convert(self, decoded_frame: DecodedFrame) -> Frame: + data = _blocks_convert_frame(self._handle, decoded_frame._handle) + # The core op produces HWC; permute to CHW to match VideoDecoder (which + # also returns a non-contiguous permuted view). + data = data.permute(2, 0, 1) + return Frame( + data=data, + pts_seconds=decoded_frame.pts_seconds, + duration_seconds=decoded_frame.duration_seconds, + ) diff --git a/src/torchcodec/decoders/_blocks/_demuxer.py b/src/torchcodec/decoders/_blocks/_demuxer.py new file mode 100644 index 000000000..1da766f2e --- /dev/null +++ b/src/torchcodec/decoders/_blocks/_demuxer.py @@ -0,0 +1,48 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from pathlib import Path + +from torchcodec._core.ops import _blocks_create_demuxer, _blocks_demuxer_next_packet + +from ._frame import Packet + + +class Demuxer: + """Demux building block: opens a container and yields the compressed + :class:`Packet`\\ s for one (video) stream. Does no decoding. + + This block is passive (it does no threading of its own) and is *not* + thread-safe: use one ``Demuxer`` per thread. It streams from the start of + the file; seeking is not supported yet. + + A :class:`Demuxer` also carries the stream configuration used to build a + :class:`PacketDecoder` and :class:`ColorConverter`, so those are constructed from + a demuxer and no extra container is opened. + """ + + def __init__(self, source: str | Path, *, stream_index: int | None = None): + if isinstance(source, Path): + source = str(source) + if not isinstance(source, str): + raise TypeError( + f"source must be a path (str or pathlib.Path), got {type(source)}" + ) + self._handle = _blocks_create_demuxer(source, stream_index) + + def next_packet(self) -> Packet | None: + """Return the next :class:`Packet`, or ``None`` at end of stream.""" + handle, is_eof = _blocks_demuxer_next_packet(self._handle) + return None if is_eof else Packet(handle) + + def __iter__(self): + while True: + packet = self.next_packet() + if packet is None: + return + yield packet diff --git a/src/torchcodec/decoders/_blocks/_frame.py b/src/torchcodec/decoders/_blocks/_frame.py new file mode 100644 index 000000000..9ad6fc4ef --- /dev/null +++ b/src/torchcodec/decoders/_blocks/_frame.py @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import torch + + +class Packet: + """Opaque, thread-movable handle to a demuxed (compressed) packet. + + Produced by :class:`Demuxer`, consumed by :class:`PacketDecoder`. It wraps a raw + pointer, so it is only valid within the process that created it (it cannot + cross a process boundary). + """ + + def __init__(self, handle: torch.Tensor): + self._handle = handle + + +class DecodedFrame: + """A decoded (YUV) frame: an opaque, thread-movable handle to the raw frame + plus its presentation timestamp and duration (in seconds). + + Produced by :class:`PacketDecoder`, consumed by :class:`ColorConverter`. The + handle wraps a raw pointer and is process-local. pts/duration are stamped by + the decoder (which knows the stream time base) and carried here so the + :class:`ColorConverter` need not be bound to any stream. + """ + + def __init__( + self, + handle: torch.Tensor, + pts_seconds: float, + duration_seconds: float, + ): + self._handle = handle + self.pts_seconds = pts_seconds + self.duration_seconds = duration_seconds diff --git a/src/torchcodec/decoders/_blocks/_packet_decoder.py b/src/torchcodec/decoders/_blocks/_packet_decoder.py new file mode 100644 index 000000000..dd892e060 --- /dev/null +++ b/src/torchcodec/decoders/_blocks/_packet_decoder.py @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from torchcodec._core.ops import ( + _blocks_create_packet_decoder, + _blocks_packet_decoder_receive_frame, + _blocks_packet_decoder_send_eof, + _blocks_packet_decoder_send_packet, +) + +from ._demuxer import Demuxer +from ._frame import DecodedFrame, Packet + + +# TODO_API_BREAKDOWN revisit every single docstring / comments at some point. + + +class PacketDecoder: + """Decode building block: turns compressed :class:`Packet`\\ s into decoded + (YUV) :class:`DecodedFrame`\\ s. + + Built from a :class:`Demuxer` (for its codec parameters) and stateful (it + holds the codec's reference-frame buffer). Passive and *not* thread-safe: + use one ``PacketDecoder`` per thread. FFmpeg's internal codec thread count + is kept at 1 for now (not exposed); parallelism comes from composing blocks + on your own threads. + """ + + def __init__(self, demuxer: Demuxer): + self._handle = _blocks_create_packet_decoder(demuxer._handle, num_threads=1) + + def _drain(self) -> list[DecodedFrame]: + frames = [] + while True: + handle, status, pts_seconds, duration_seconds = ( + _blocks_packet_decoder_receive_frame(self._handle) + ) + if status != 0: # EAGAIN (need more packets) or EOF: nothing ready + break + frames.append(DecodedFrame(handle, pts_seconds, duration_seconds)) + return frames + + def decode(self, packet: Packet) -> list[DecodedFrame]: + """Send one packet and return whatever frames are now ready (possibly + empty, e.g. while the codec buffers B-frames).""" + status = _blocks_packet_decoder_send_packet(self._handle, packet._handle) + if status < 0: + raise RuntimeError(f"Failed to send packet to decoder (status {status})") + return self._drain() + + def flush(self) -> list[DecodedFrame]: + """Signal end-of-stream and return all remaining buffered frames. Call + once, after the last packet.""" + _blocks_packet_decoder_send_eof(self._handle) + return self._drain() diff --git a/test/test_decoders.py b/test/test_decoders.py index 81805eff2..f8b5bf94a 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -6,12 +6,15 @@ import contextlib import gc +import queue +import threading from functools import partial import numpy import pytest import torch from torchcodec import _core, ffmpeg_major_version, FrameBatch +from torchcodec._frame import Frame from torchcodec.decoders import ( AudioDecoder, AudioStreamMetadata, @@ -22,6 +25,13 @@ VideoStreamMetadata, WavDecoder, ) +from torchcodec.decoders._blocks import ( + ColorConverter, + DecodedFrame, + Demuxer, + Packet, + PacketDecoder, +) from torchcodec.decoders._decoder_utils import _get_cuda_backend from torchcodec.encoders import VideoEncoder from torchcodec.transforms import CenterCrop, RandomCrop, Resize @@ -3129,3 +3139,179 @@ def test_multiple_calls_with_backward_seeks(self): wav_samples.data, audio_samples.data, rtol=0, atol=0 ) assert wav_samples.pts_seconds == audio_samples.pts_seconds + + +class TestBlocks: + + def test_block_output_types(self): + # 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() + + num_packets = 0 + for packet in demuxer: + assert isinstance(packet, Packet) + num_packets += 1 + for decoded in decoder.decode(packet): + assert isinstance(decoded, DecodedFrame) + frame = converter.convert(decoded) + assert isinstance(frame, Frame) + assert frame.data.ndim == 3 # CHW + assert frame.data.shape[0] == 3 # channels first + assert frame.data.dtype == torch.uint8 + assert frame.duration_seconds >= 0 + + assert num_packets > 0 + + # The three decode stages, each expressed as a generator that transforms an + # iterator of inputs into an iterator of outputs. They compose directly (the + # sequential pipeline is just convert(decode(demux()))), and a thread + # boundary between any two stages is inserted with prefetch() below. + + @staticmethod + def _demux(demuxer): + yield from demuxer + + @staticmethod + def _decode(decoder, packets): + for packet in packets: + yield from decoder.decode(packet) + yield from decoder.flush() + + @staticmethod + def _convert(converter, frames): + for frame in frames: + yield converter.convert(frame) + + @staticmethod + def prefetch(upstream, buffer_size=8): + # Run `upstream` (a generator chaining one or more stages) on a + # background thread, yielding its items through a bounded queue. This is + # the only threading primitive: where you insert it decides which stages + # overlap. The queue applies backpressure (the worker blocks in q.put() + # when the buffer is full), so it runs ~buffer_size items ahead. + q: queue.Queue = queue.Queue(maxsize=buffer_size) + eof = object() + error = [] + + def worker(): + try: + for item in upstream: + q.put(item) + except Exception as e: # surface failures instead of hanging + error.append(e) + finally: + q.put(eof) + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + + def drain(): + while (item := q.get()) is not eof: + yield item + thread.join() # worker enqueued eof and is finishing; make it explicit + if error: + raise error[0] + + return drain() + + def _decoded_frames(self, path): + # demux + decode, as a single generator of DecodedFrames (pts order). + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + return self._decode(decoder, self._demux(demuxer)) + + def _decode_sequential(self, path): + # demux -> decode -> color-convert, all on the calling thread. + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + converter = ColorConverter() + return list( + self._convert(converter, self._decode(decoder, self._demux(demuxer))) + ) + + def _decode_prefetch_frames(self, path): + # [demux + decode] on one thread || [color-convert] on another. + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + converter = ColorConverter() + frames = self.prefetch(self._decode(decoder, self._demux(demuxer))) + return list(self._convert(converter, frames)) + + def _decode_prefetch_packets(self, path): + # [demux] on one thread || [decode + color-convert] on another. + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + converter = ColorConverter() + packets = self.prefetch(self._demux(demuxer)) + return list(self._convert(converter, self._decode(decoder, packets))) + + def _decode_prefetch_packets_and_frames(self, path): + # [demux] || [decode] || [color-convert], each on its own thread. + demuxer = Demuxer(path) + decoder = PacketDecoder(demuxer) + converter = ColorConverter() + packets = self.prefetch(self._demux(demuxer)) + frames = self.prefetch(self._decode(decoder, packets)) + return list(self._convert(converter, frames)) + + def _to_frame_batch(self, frames): + return FrameBatch( + data=torch.stack([f.data for f in frames]), + pts_seconds=torch.tensor( + [f.pts_seconds for f in frames], dtype=torch.float64 + ), + duration_seconds=torch.tensor( + [f.duration_seconds for f in frames], dtype=torch.float64 + ), + ) + + @pytest.mark.parametrize("video", (NASA_VIDEO, BT709_FULL_RANGE)) + @pytest.mark.parametrize( + "decode_method", + ( + _decode_sequential, + _decode_prefetch_frames, + _decode_prefetch_packets, + _decode_prefetch_packets_and_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() + + 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) + torch.testing.assert_close( + got.duration_seconds, ref.duration_seconds, atol=0, rtol=0 + ) + + def test_color_converter_reused_across_videos(self): + # 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() + videos = [NASA_VIDEO, BT709_FULL_RANGE] + generators = [self._decoded_frames(v.path) for v in videos] + outputs = [[] for _ in videos] + + done = [False] * len(videos) + while not all(done): + for i, gen in enumerate(generators): + if done[i]: + continue + decoded = next(gen, None) + if decoded is None: + done[i] = True + continue + outputs[i].append(converter.convert(decoded)) + + for video, frames in zip(videos, outputs): + got = self._to_frame_batch(frames) + ref = VideoDecoder(video.path).get_all_frames() + assert got.data.shape == ref.data.shape + torch.testing.assert_close(got.data, ref.data, atol=0, rtol=0) From c623b89091a5784bc9357475c2c46ba3d463f185 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Thu, 9 Jul 2026 14:39:23 +0100 Subject: [PATCH 3/3] WIP --- src/torchcodec/_core/CpuDeviceInterface.cpp | 1 - src/torchcodec/_core/Demuxer.cpp | 5 +---- src/torchcodec/_core/PacketDecoder.cpp | 3 +++ src/torchcodec/decoders/_blocks/_packet_decoder.py | 2 ++ 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/torchcodec/_core/CpuDeviceInterface.cpp b/src/torchcodec/_core/CpuDeviceInterface.cpp index 45d336f5d..46731ea0a 100644 --- a/src/torchcodec/_core/CpuDeviceInterface.cpp +++ b/src/torchcodec/_core/CpuDeviceInterface.cpp @@ -86,7 +86,6 @@ 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: // STD_TORCH_CHECK(av_stream != nullptr, "avStream is null"); // time_base_ = av_stream->time_base; diff --git a/src/torchcodec/_core/Demuxer.cpp b/src/torchcodec/_core/Demuxer.cpp index 106e074a3..4f4ff4c6e 100644 --- a/src/torchcodec/_core/Demuxer.cpp +++ b/src/torchcodec/_core/Demuxer.cpp @@ -16,10 +16,7 @@ int read_next_packet( int status = AVSUCCESS; do { status = av_read_frame(format_context, packet.get()); - if (status == AVERROR_EOF) { - return AVERROR_EOF; - } - if (status < AVSUCCESS) { + if (status == AVERROR_EOF || status < AVSUCCESS) { return status; } } while (packet->stream_index != active_stream_index); diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index 4cdab49a8..23e726f8e 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -8,6 +8,9 @@ namespace facebook::torchcodec { +// TODO_API_BREAKDOWN: we should make sure the block APIs can dispatch to +// third-party extensions - all of them. + SharedAVCodecContext create_and_open_codec_context( AVStream* stream, const AVCodec* av_codec, diff --git a/src/torchcodec/decoders/_blocks/_packet_decoder.py b/src/torchcodec/decoders/_blocks/_packet_decoder.py index dd892e060..506096062 100644 --- a/src/torchcodec/decoders/_blocks/_packet_decoder.py +++ b/src/torchcodec/decoders/_blocks/_packet_decoder.py @@ -53,6 +53,8 @@ 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 + # 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 once, after the last packet."""