Hand out CUDA frames from a CUDA PacketDecoder, always (#1636)

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -33,14 +33,16 @@
 #include "nvcuvid_include/nvcuvid.h"
 
 namespace facebook::torchcodec {
+// TODO_API_BREAKDOWN P2: the name says "standalone", but this is really about
+// owning a GPU buffer. Find one that covers both.
 struct StandAloneFrameAttachedData {
   cudaStream_t producer_stream = nullptr;
   torch::stable::Tensor storage;
-  // Whether the frame's samples are on the GPU. False for the CPU-fallback
-  // frames a PacketDecoder hands out for streams NVDEC can't decode. The
-  // pixel format can't tell the two apart: a 4:4:4 stream yields yuv444p
-  // either way, natively from NVDEC or from the CPU fallback.
-  bool is_device_frame = false;
+};
+
+struct GpuFrameAndStorage {
+  UniqueAVFrame av_frame;
+  torch::stable::Tensor storage;
 };
 
 class BetaCudaDeviceInterface : public DeviceInterface {
@@ -108,12 +110,12 @@ class BetaCudaDeviceInterface : public DeviceInterface {
 
   void make_frame_standalone(UniqueAVFrame& av_frame) override;
 
-  bool is_device_frame(
-      [[maybe_unused]] const UniqueAVFrame& av_frame) const override;
+  GpuFrameAndStorage upload_cpu_frame_to_gpu_on_current_stream(
+      const AVFrame& cpu_frame);
 
-  UniqueAVFrame transfer_cpu_frame_to_gpu(
-      const AVFrame& cpu_frame,
-      AVPixelFormat target_pix_fmt);
+  torch::stable::Tensor copy_nvdec_surface(
+      UniqueAVFrame& av_frame,
+      cudaStream_t stream);
 
   void apply_rotation(
       FrameOutput& frame_output,

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -160,11 +160,6 @@ class DeviceInterface {
   virtual void make_frame_standalone([[maybe_unused]] UniqueAVFrame& av_frame) {
   };
 
-  virtual bool is_device_frame(
-      [[maybe_unused]] const UniqueAVFrame& av_frame) const {
-    return false;
-  }
-
   // Flush remaining frames from decoder
   virtual void flush() {
     STD_TORCH_CHECK(

src/torchcodec/_core/DeviceInterface.h

@@ -45,10 +45,6 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder {
   // if more input is needed, AVERROR_EOF at end, or a negative error code.
   int receive_frame(UniqueAVFrame& av_frame);
 
-  bool is_device_frame(const UniqueAVFrame& av_frame) const {
-    return device_interface_->is_device_frame(av_frame);
-  }
-
   const StableDevice& device() const {
     return device_interface_->device();
   }

src/torchcodec/_core/PacketDecoder.h

@@ -876,9 +876,11 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
   AVRational time_base = decoder_ptr->time_base();
   double pts_seconds = pts_to_seconds(get_pts_or_dts(*av_frame), time_base);
   double duration_seconds = pts_to_seconds(get_duration(*av_frame), time_base);
-  std::string device = decoder_ptr->is_device_frame(av_frame)
-      ? device_to_string(decoder_ptr->device())
-      : std::string("cpu");
+  // A decoder's frames always live on the decoder's own device, including the
+  // ones a CUDA decoder had to decode on the CPU and upload.
+  // TODO_API_BREAKDOWN_P1: Not sure we need to return the device at all, the
+  // device should always match the device parameter now.
+  std::string device = device_to_string(decoder_ptr->device());
   return std::make_tuple(
       wrap_pointer_to_tensor(std::move(av_frame)),
       static_cast<int64_t>(0),

src/torchcodec/_core/custom_ops.cpp

@@ -4,7 +4,7 @@
 # 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).
+"""Private, experimental building-block decode API.
 
 Exposes the three decode stages -- :class:`Demuxer`, :class:`PacketDecoder`,
 :class:`ColorConverter` -- as passive, composable, GIL-releasing units, so a

src/torchcodec/decoders/_blocks/__init__.py

@@ -33,22 +33,23 @@ def __init__(self, handle: torch.Tensor):
 @dataclass
 class RawFrame:
     planes: tuple[torch.Tensor, ...]
-    pix_fmt: str  # FFmpeg pixel-format name, e.g. "yuv420p"
+    # FFmpeg pixel-format name. On CPU this is the source's own format, e.g.
+    # "yuv420p". On CUDA it is always an NVDEC surface format: "nv12",
+    # "p010le", "p012le", "p016le", "yuv444p" or "yuv444p16le".
+    pix_fmt: str
     colorspace: str  # e.g. "bt709"
     color_range: str  # "tv" (limited) or "pc" (full)
-    # The depth of pix_fmt (always).
-    # This is also the source's bit depth, except for 12b-bit sources CUDA
-    # frames: those are technically P012, but P012 was only introduced in FFmpeg
-    # 6. So For FFmpeg < 6, we must report those as P016, and so the bit_depth
-    # field here reports 16 (on CPU, it'd still be 12).
+    # The depth of pix_fmt (always), which is the source's bit depth except
+    # where a CUDA surface format's container is wider than the source samples:
+    # a 10-bit 4:4:4 source is uploaded as yuv444p16le, and a 12-bit source is
+    # tagged p016le on FFmpeg < 6 (which lacks p012le). Both report 16 here (on
+    # CPU they'd report 10 and 12).
     # Everything downstream still reads right, because those samples are
-    # msb-aligned and are therefore genuinely valid 16-bit ones, with 4 zeroed
+    # msb-aligned and are therefore genuinely valid 16-bit ones, with zeroed
     # low bits.
     #
-    # TODO_API_BREAKDOWN P2: We can't do anything about the P016 report, but
-    # should this actually report the depth of the source instead of the depth
-    # of the pixel format? Again the only discrepency arises for 12-bit sources
-    # on CUDA for FFmpeg < 6.
+    # TODO_API_BREAKDOWN P2: should this report the depth of the source instead
+    # of the depth of the pixel format?
     bit_depth: int
 
 
@@ -63,9 +64,10 @@ class DecodedFrame:
     the decoder (which knows the stream time base) and carried here so the
     :class:`ColorConverter` need not be bound to any stream.
 
-    ``device`` is where this frame's samples actually are, which is not
-    necessarily the device its decoder was created with: a CUDA decoder falls
-    back to CPU decoding for streams NVDEC can't handle.
+    ``device`` is where this frame's samples are, and it is always the device
+    its decoder was created with. A CUDA decoder falls back to CPU decoding for
+    streams NVDEC can't handle, but it uploads those frames before handing them
+    out, so they are indistinguishable from NVDEC ones here.
     """
 
     def __init__(

src/torchcodec/decoders/_blocks/_frame.py

@@ -134,6 +134,7 @@
     TESTSRC2_AV1_10BIT,
     TESTSRC2_ODD_HEIGHT_444,
     TESTSRC2_ODD_HEIGHT_AND_WIDTH_444,
+    TESTSRC2_ODD_HEIGHT_AND_WIDTH_444_10BIT,
     TESTSRC2_ODD_HEIGHT_AND_WIDTH_VP9,
     TESTSRC2_ODD_HEIGHT_AND_WIDTH_VP9_10BIT,
     TESTSRC2_ODD_HEIGHT_VP9,
@@ -2300,7 +2301,9 @@ def test_nvdec_cpu_fallback_yuv444(self, tmp_path):
         cpu_frames = cpu_decoder.get_frames_in_range(start=0, stop=num_frames).data
         cuda_frames = cuda_decoder.get_frames_in_range(start=0, stop=num_frames).data
 
-        torch.testing.assert_close(cpu_frames, cuda_frames.cpu(), rtol=0, atol=0)
+        # The CUDA path uploads these as yuv444p and color-converts them with
+        # our kernel, which truncates where swscale rounds.
+        torch.testing.assert_close(cpu_frames, cuda_frames.cpu(), rtol=0, atol=1)
 
     @needs_cuda
     def test_nvdec_cuda_interface_error(self):
@@ -3500,6 +3503,7 @@ def pix_fmt(self, device):
     TESTSRC2_ODD_WIDTH_VP9_10BIT,
     TESTSRC2_ODD_HEIGHT_VP9_10BIT,
     TESTSRC2_ODD_HEIGHT_AND_WIDTH_VP9_10BIT,
+    TESTSRC2_ODD_HEIGHT_AND_WIDTH_444_10BIT,
 )
 
 
@@ -3509,9 +3513,12 @@ def pix_fmt(self, device):
 _MATERIALIZE_VIDEOS = (
     _MaterializeCase(NASA_VIDEO, 8, "yuv420p", "nv12"),  # even dims
     _MaterializeCase(TESTSRC2_ODD_HEIGHT_AND_WIDTH_VP9, 8, "yuv420p", "nv12"),  # odd
-    # 4:4:4 (full-res chroma). NVDEC can't decode H264 4:4:4, so this one falls
-    # back to the CPU and keeps its native format.
+    # 4:4:4 (full-res chroma). NVDEC can't decode H264 4:4:4, so these fall back
+    # to the CPU and are uploaded as 4:4:4 rather than have their chroma halved.
     _MaterializeCase(TESTSRC2_ODD_HEIGHT_AND_WIDTH_444, 8, "yuv444p", "yuv444p"),
+    _MaterializeCase(
+        TESTSRC2_ODD_HEIGHT_AND_WIDTH_444_10BIT, 10, "yuv444p10le", "yuv444p16le"
+    ),
     # HEVC 4:4:4, which NVDEC decodes natively into its YUV444 surfaces. The
     # 16-bit one is the only 4:4:4 surface above 8 bits, so 10- and 12-bit
     # sources both land in yuv444p16le.
@@ -3696,6 +3703,7 @@ def _assert_matches_video_decoder(got, ref, video):
             TESTSRC2_ODD_WIDTH_444,
             TESTSRC2_ODD_HEIGHT_444,
             TESTSRC2_ODD_HEIGHT_AND_WIDTH_444,
+            TESTSRC2_ODD_HEIGHT_AND_WIDTH_444_10BIT,
             # HEVC 4:4:4: NVDEC decodes these natively instead.
             TESTSRC2_444_8BIT_HEVC,
             TESTSRC2_444_10BIT_HEVC,
@@ -3867,6 +3875,15 @@ def test_materialize_structure(self, case, device):
         expected_dtype = torch.uint16 if raw.bit_depth > 8 else torch.uint8
         assert all(plane.dtype == expected_dtype for plane in planes)
 
+        if device == "cuda" and expected_dtype == torch.uint16:
+            # Whatever depth the format claims, a 16-bit CUDA surface holds the
+            # source's samples msb-aligned, with the unused low bits zeroed.
+            # That's what lets test_materialize_cuda_planes_match_cpu shift them
+            # back down by 16 - bit_depth and compare against the CPU planes.
+            unused_low_bits = (1 << (16 - case.bit_depth)) - 1
+            for plane in planes:
+                assert (plane.to(torch.int32) & unused_low_bits).count_nonzero() == 0
+
         Y, U, V = planes
         height, width = converter.convert(frame).data.shape[1:]
         assert Y.shape == (height, width)
@@ -4005,23 +4022,29 @@ def test_materialize_neutral_chroma_is_grayscale(self, video, device):
         torch.testing.assert_close(g, b, atol=1, rtol=0)
 
     @pytest.mark.needs_cuda
-    @pytest.mark.parametrize("video", (H265_VIDEO, TESTSRC2_ODD_HEIGHT_AND_WIDTH_444))
-    def test_materialize_cpu_fallback_stays_on_cpu(self, video):
-        # TODO_NOW: This may not be what we want. We probalby want
-        # to output CUDA data.  But how? Do we put the YUV420 on CUDA? Then we
-        # need a specialized kernel to color-convert? Or we put those frames we
-        # can have on NV12 - but for 444 it's a problem becaus ewe can't convert
-        # them to NV12, so we'd need a 444 color-conversion kernel anyway. So
-        # all frames would be NV12 except *some* (the 444 ones)?
-        # Unclear what to do here honestly. Maybe outputting CPU frames is
-        # actually justifiable?
+    @pytest.mark.parametrize(
+        "video, expected_pix_fmt",
+        (
+            # Too small for NVDEC.
+            (H265_VIDEO, "nv12"),
+            # H264 4:4:4, which NVDEC can't decode. Uploading it as NV12 would
+            # halve its chroma resolution, so it stays 4:4:4.
+            (TESTSRC2_ODD_HEIGHT_AND_WIDTH_444, "yuv444p"),
+            (TESTSRC2_ODD_HEIGHT_AND_WIDTH_444_10BIT, "yuv444p16le"),
+        ),
+    )
+    def test_materialize_cpu_fallback_is_on_cuda(self, video, expected_pix_fmt):
+        # A CUDA PacketDecoder hands out CUDA frames even for the streams it has
+        # to decode on the CPU, and they're in an NVDEC surface format like any
+        # other CUDA frame.
+        assert VideoDecoder(video.path, device="cuda").cpu_fallback
+
         frame, _ = self._first_frame(video.path, "cuda")
-        assert frame.device == "cpu"
+        assert frame.device == "cuda"
 
         raw = frame.materialize()
-        planes, pix_fmt = raw.planes, raw.pix_fmt
-        assert pix_fmt != "nv12"
-        assert all(plane.device.type == "cpu" for plane in planes)
+        assert raw.pix_fmt == expected_pix_fmt
+        assert all(plane.device.type == "cuda" for plane in raw.planes)
 
 
 # Small helpers to avoid having to always specify the same skip marks and decode_fn

test/resources/testsrc2_odd_height_and_width_444_10bit.mp4

@@ -1229,6 +1229,18 @@ def get_empty_chw_tensor(self, *, stream_index: int) -> torch.Tensor:
     frames={0: {}},
 )
 
+# ffmpeg -f lavfi -i "testsrc2=size=321x241:rate=25:duration=1,format=rgb24" \
+#  -c:v libx264 -pix_fmt yuv444p10le -profile:v high444 \
+#  testsrc2_odd_height_and_width_444_10bit.mp4
+TESTSRC2_ODD_HEIGHT_AND_WIDTH_444_10BIT = TestVideo(
+    filename="testsrc2_odd_height_and_width_444_10bit.mp4",
+    default_stream_index=0,
+    stream_infos={
+        0: TestVideoStreamInfo(width=321, height=241, num_color_channels=3),
+    },
+    frames={0: {}},
+)
+
 # HEVC 4:4:4, which NVDEC *can* decode natively (unlike H264 4:4:4 above), at
 # 8, 10 and 12 bits. Odd dimensions, so they also cover the cropping NVDEC's
 # even-aligned surfaces need. Encoded with, for DEPTH in 8/10/12: