Fix silent frame corruption when a ColorConverter runs on its own CUDA stream (#1651)

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -1019,32 +1019,6 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
     storage = copy_nvdec_surface(av_frame, current_stream);
   }
 
-  // TODO_API_BREAKDOWN CORRECTNESS P0: `storage` comes from the PyTorch
-  // caching allocator on `current_stream`, but a ColorConverter reads it from
-  // whatever stream it runs on. The allocator only tracks the allocating
-  // stream: once the frame is dropped, the block returns to `current_stream`'s
-  // pool with no synchronisation, and the next frame's copy - same size, same
-  // stream - lands right in it while the converter is still reading. Frames
-  // are then silently corrupted whenever the converter lags behind the
-  // decoder, which is the normal state of a two-stream pipeline. Measured on a
-  // 4K clip with the converter backlogged: 117 of 119 frames wrong, and clean
-  // again as soon as the frames are kept alive.
-  // The usual remedy is Tensor::record_stream() on the consumer side, but
-  // neither the stable ABI nor the AOTI shim exposes it, and StableIValue has
-  // no Stream conversion, so it isn't reachable through the dispatcher either.
-  // We don't actually need it though: the allocator recycles the block because
-  // *we* drop our reference too early, so it's enough to hold on to the
-  // storage until the consumer is done. Have the ColorConverter record an
-  // event on its own stream into the attached data, and make
-  // standalone_frame_free_callback() hand (storage, event) to a per-device
-  // pending-release list instead of dropping the tensor. Drain that list
-  // opportunistically with cudaEventQuery. No host stall, and a consumer that
-  // permanently lags shows up as a growing list, i.e. as backpressure rather
-  // than as silent corruption.
-  // Frames that were never converted have no event and can be released
-  // straight away. Raw planes handed to the user via materialize() keep the
-  // storage alive on their own, and once the user drops those, ordering their
-  // own kernels is their responsibility, same as for any other tensor.
   auto attached_data = new StandAloneFrameAttachedData();
   attached_data->producer_stream = current_stream;
   attached_data->storage = std::move(storage);
@@ -1059,6 +1033,62 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
       "Failed to attach standalone frame data");
 }
 
+std::optional<torch::stable::Tensor> BetaCudaDeviceInterface::get_frame_storage(
+    const AVFrame& av_frame) const {
+  STD_TORCH_CHECK(
+      // Only decoder-only should reach here, and this should only be called on
+      // frames that went through make_frame_standalone(), which sets
+      // opaque_ref.
+      mode() == Mode::DecoderOnly && av_frame.opaque_ref != nullptr,
+      "Unexpected call to get_frame_storage(), please report a bug ");
+
+  // Note [Standalone Frame Storage and the need for record_stream]
+  //
+  // A PacketDecoder and a ColorConverter may run on different CUDA streams.
+  // Consider the following:
+  //
+  // ```
+  // with decoder_stream:
+  //   frame = decoder.receive_frame()
+  // with color_converter_stream:
+  //   color_converter.convert(frame)
+  //
+  // del frame
+  //
+  // with decoder_stream:
+  //   frame = decoder.receive_frame()
+  // ```
+  //
+  // The call to convert(frame) is non-blocking and just enqueues the
+  // color-conversion kernel. The CPU moves on immediately to `del frame` while
+  // the kernel is still running (it may also not even have started depending on
+  // how color_converter_stream is congested).
+  //
+  // When the frame is deleted, the torch CUDA allocator reclaims its memory and
+  // it becomes available for reuse for any subsequent allocation on the
+  // decoder_stream. If the next decoder.receive_frame() happens before the
+  // color-conversion kernel has finished (specifically: the new storage
+  // allocation for that next frame in make_frame_standalone()), the memory is
+  // reused, overwritten, and the color-conversion kernel reads garbage (i.e.
+  // the next frame's samples!).
+  //
+  // We're hitting exactly what
+  // https://zdevito.github.io/2022/08/04/cuda-caching-allocator.html describes
+  // in the 'Streams and freeing memory' section, and the solution is to call
+  // record_stream() on the frame's storage within color_conversion_stream just
+  // after the kernel is enqueued: this tells the allocator that it must wait
+  // until this point (on the device side) before reclaiming the memory.
+  //
+  // We call record_stream(color_conversion_stream) on the frame storage in the
+  // ColorConverter, on behalf of the user. But we still must expose the storage
+  // for those users who would like to consume the frame with their own
+  // consumer, i.e. not using the ColorConverter: they need to call
+  // frame.storage.record_stream(color_conversion_stream) themselves.
+  return reinterpret_cast<StandAloneFrameAttachedData*>(
+             av_frame.opaque_ref->data)
+      ->storage;
+}
+
 torch::stable::Tensor BetaCudaDeviceInterface::copy_nvdec_surface(
     UniqueAVFrame& av_frame,
     cudaStream_t current_stream) {

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -121,6 +121,9 @@ class BetaCudaDeviceInterface : public DeviceInterface {
 
   void make_frame_standalone(UniqueAVFrame& av_frame) override;
 
+  std::optional<torch::stable::Tensor> get_frame_storage(
+      const AVFrame& av_frame) const override;
+
   GpuFrameAndStorage upload_cpu_frame_to_gpu(
       const AVFrame& cpu_frame,
       cudaStream_t stream);

src/torchcodec/_core/DeviceInterface.h

@@ -160,6 +160,11 @@ class DeviceInterface {
   virtual void make_frame_standalone([[maybe_unused]] UniqueAVFrame& av_frame) {
   };
 
+  virtual std::optional<torch::stable::Tensor> get_frame_storage(
+      [[maybe_unused]] const AVFrame& av_frame) const {
+    return std::nullopt;
+  }
+
   // Flush remaining frames from decoder
   virtual void flush() {
     STD_TORCH_CHECK(

src/torchcodec/_core/PacketDecoder.h

@@ -46,6 +46,11 @@ 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);
 
+  std::optional<torch::stable::Tensor> get_frame_storage(
+      const AVFrame& av_frame) const {
+    return device_interface_->get_frame_storage(av_frame);
+  }
+
   const StableDevice& device() const {
     return device_interface_->device();
   }

src/torchcodec/_core/custom_ops.cpp

@@ -82,7 +82,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) {
       "_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, str)");
+      "_blocks_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, str, Tensor)");
   m.def(
       "_blocks_create_color_converter(str device=\"cpu\", str output_dtype=\"uint8\") -> Tensor");
   m.def("_blocks_convert_frame(Tensor(a!) converter, Tensor frame) -> Tensor");
@@ -855,9 +855,14 @@ std::string device_to_string(const StableDevice& device) {
   return name;
 }
 
-// (frame_handle, status, pts_seconds, duration_seconds, device).
-using OpsReceiveFrameOutput =
-    std::tuple<torch::stable::Tensor, int64_t, double, double, std::string>;
+// (frame_handle, status, pts_seconds, duration_seconds, device, storage).
+using OpsReceiveFrameOutput = std::tuple<
+    torch::stable::Tensor,
+    int64_t,
+    double,
+    double,
+    std::string,
+    torch::stable::Tensor>;
 
 OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
     torch::stable::Tensor& decoder) {
@@ -871,7 +876,8 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
         static_cast<int64_t>(status),
         0.0,
         0.0,
-        std::string("cpu"));
+        std::string("cpu"),
+        torch::stable::empty({int64_t(0)}, kStableUInt8));
   }
   AVRational time_base = decoder_ptr->time_base();
   double pts_seconds = pts_to_seconds(get_pts_or_dts(*av_frame), time_base);
@@ -881,12 +887,16 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
   // TODO_API_BREAKDOWN DESIGN 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());
+  torch::stable::Tensor storage =
+      decoder_ptr->get_frame_storage(*av_frame).value_or(
+          torch::stable::empty({int64_t(0)}, kStableUInt8));
   return std::make_tuple(
       wrap_pointer_to_tensor(std::move(av_frame)),
       static_cast<int64_t>(0),
       pts_seconds,
       duration_seconds,
-      device);
+      device,
+      storage);
 }
 
 torch::stable::Tensor _blocks_create_color_converter(

src/torchcodec/decoders/_blocks/_color_converter.py

@@ -52,6 +52,9 @@ def __init__(
 
     def convert(self, decoded_frame: DecodedFrame) -> Frame:
         data = _blocks_convert_frame(self._handle, decoded_frame._handle)
+        if decoded_frame.storage is not None:
+            # See [Standalone Frame Storage and the need for record_stream]
+            decoded_frame.storage.record_stream(torch.cuda.current_stream())
         # 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)

src/torchcodec/decoders/_blocks/_frame.py

@@ -35,6 +35,7 @@ def __init__(self, handle: torch.Tensor):
 # and the ColorConverter honors it - but nothing surfaces it in Python, so
 # whoever color-converts the planes themselves has no way of knowing the frame
 # needs rotating. And should reported height and width be pre or post rotation?
+# Same for storage: where should we expose it?
 @dataclass
 class RawFrame:
     planes: tuple[torch.Tensor, ...]
@@ -81,9 +82,11 @@ def __init__(
         pts_seconds: float,
         duration_seconds: float,
         device: str = "cpu",
+        storage: torch.Tensor | None = None,
     ):
         self._handle = handle
         self._device = device
+        self.storage = storage
         self.pts_seconds = pts_seconds
         self.duration_seconds = duration_seconds
 

src/torchcodec/decoders/_blocks/_packet_decoder.py

@@ -40,7 +40,7 @@ def __init__(self, demuxer: Demuxer, device="cpu"):
     def _drain(self) -> list[DecodedFrame]:
         frames = []
         while True:
-            handle, status, pts_seconds, duration_seconds, device = (
+            handle, status, pts_seconds, duration_seconds, device, storage = (
                 _blocks_packet_decoder_receive_frame(self._handle)
             )
             if status != 0:  # EAGAIN (need more packets) or EOF: nothing ready
@@ -51,6 +51,7 @@ def _drain(self) -> list[DecodedFrame]:
                     pts_seconds,
                     duration_seconds,
                     device=device,
+                    storage=storage if storage.numel() > 0 else None,
                 )
             )
         return frames

test/test_decoders.py

@@ -7,6 +7,7 @@
 import concurrent.futures
 import contextlib
 import gc
+import itertools
 import math
 import queue
 import threading
@@ -3974,6 +3975,75 @@ def test_materialize_planes_are_not_rotated_but_color_conversion_rotates(
         assert Y.shape == (width, height)
         assert converter.convert(frame).data.shape == (3, height, width)
 
+    @pytest.mark.needs_cuda
+    @pytest.mark.parametrize("record_stream", (True, False))
+    def test_storage_record_stream(self, record_stream):
+        # Using PacketDecoder on one stream and consuming the frames on a
+        # different stream requires the user to call record_stream() on the
+        # frame storage.
+        # Without the record_stream() call the decoder's next frame may be
+        # handed the same buffer and overwrites it while the read is still
+        # queued.
+        # See [Standalone Frame Storage and the need for record_stream]
+        video = NASA_VIDEO.path
+        decode_stream = torch.cuda.Stream()
+        read_stream = torch.cuda.Stream()
+
+        def run(separate_stream):
+            demuxer, decoder, _ = self._make_blocks(video, "cuda")
+            reads = []
+            for packet in itertools.chain(demuxer, [None]):
+                with torch.cuda.stream(decode_stream):
+                    decoded = (
+                        decoder.flush() if packet is None else decoder.decode(packet)
+                    )
+                with torch.cuda.stream(
+                    read_stream if separate_stream else decode_stream
+                ):
+                    while decoded:
+                        frame = decoded.pop(0)
+                        torch.cuda._sleep(20_000_000)  # ~10ms, fall behind
+                        reads.append(frame.materialize().planes[0].clone())
+                        if separate_stream and record_stream:
+                            frame.storage.record_stream(read_stream)
+            torch.cuda.synchronize()
+            return reads
+
+        ref = run(separate_stream=False)
+        got = run(separate_stream=True)
+        wrong = sum(1 for a, b in zip(ref, got) if not torch.equal(a, b))
+        if record_stream:
+            assert wrong == 0
+        else:
+            # Guards the test itself: without the call this really does corrupt,
+            # so the assertion above is meaningful.
+            assert wrong > 0
+
+    @pytest.mark.needs_cuda
+    def test_backlogged_converter_on_separate_stream(self):
+        # Similar test to test_storage_record_stream(), but with the ColorConverter on a
+        # separate stream. In this case, *we* call record_stream() on behalf of
+        # the user.
+        # See [Standalone Frame Storage and the need for record_stream]
+        video = NASA_VIDEO.path
+        demuxer, decoder, converter = self._make_blocks(video, "cuda")
+        decode_stream = torch.cuda.Stream()
+        convert_stream = torch.cuda.Stream()
+
+        frames = []
+        for packet in itertools.chain(demuxer, [None]):
+            with torch.cuda.stream(decode_stream):
+                decoded = decoder.flush() if packet is None else decoder.decode(packet)
+            with torch.cuda.stream(convert_stream):
+                while decoded:
+                    torch.cuda._sleep(20_000_000)  # ~10ms
+                    frames.append(converter.convert(decoded.pop(0)))
+        torch.cuda.synchronize()
+
+        got = self._to_frame_batch(frames)
+        ref = VideoDecoder(video, device="cuda").get_all_frames()
+        torch.testing.assert_close(got.data, ref.data, atol=0, rtol=0)
+
     @pytest.mark.needs_cuda
     @pytest.mark.parametrize("case", _MATERIALIZE_VIDEOS, ids=_materialize_ids)
     def test_materialize_cuda_planes_match_cpu(self, case):