Add Cuda Context guard

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -266,6 +266,8 @@ std::optional<cudaVideoSurfaceFormat> get_nvdec_surface_format(
 
 // Callback for freeing CUDA memory associated with AVFrame see where it's used
 // for more details.
+// TODO_API_BREAKDOWN P2: Should we align this with the other free callback
+// below?  Why did we use cudaMalloc? Can we just allocate with torch??
 void cuda_buffer_free_callback(void* opaque, [[maybe_unused]] uint8_t* data) {
   cudaFree(opaque);
 }
@@ -276,6 +278,36 @@ void standalone_frame_free_callback(
   delete reinterpret_cast<StandAloneFrameAttachedData*>(data);
 }
 
+class CudaContextGuard {
+  // There's one CUDA context per process per device. But new threads aren't
+  // bound to a context. The binding often happens automatically when calling
+  // CUDA APIs (like cudaFree), but some APIs like the NVCUVID ones that we use
+  // here aren't automatically binding.
+  // So for a thread to be able to use NVCUVID APIs, it must have a context
+  // bound to it, and we have to enforce that binding manually.
+  // That's what this guard does: it calls cudaFree(nullptr), which is a common
+  // near-free way to force the CUDA runtime to bind the context for the current
+  // thread. And this call must happen within a device guard to make sure we're
+  // binding the context of the device this interface is using.
+  // We must call this guard in every public method of the interface that uses
+  // NVCUVID APIs, because these methods can, in theory, be called from any
+  // thread.
+  // Note that none of this was an issue before when our only entry-point was
+  // the SingleStreamDecoder: all the entry-points were called from the same
+  // thread. Now that we have split the APIs in different blocks (PacketDecoder,
+  // ColorConverter), each of these blocks can be on different threads - and
+  // importantly, they can be created in the main thread (where the context is
+  // bound by our call to initialize_cuda_context_with_pytorch()), but then used
+  // in a different thread that doesn't have the context.
+ public:
+  explicit CudaContextGuard(int device_index) : device_guard_(device_index) {
+    cudaFree(nullptr);
+  }
+
+ private:
+  StableDeviceGuard device_guard_;
+};
+
 } // namespace
 
 BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device)
@@ -284,6 +316,9 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device)
   STD_TORCH_CHECK(
       device_.type() == kStableCUDA, "Unsupported device: must be CUDA");
 
+  // Note: now that we have the CudaContextGuard, we might not need to do that
+  // anymore. The comment says we need pytorch to create the context - maybe
+  // that's true, but that's a very old comment now.
   initialize_cuda_context_with_pytorch(device_);
 
   nvcuvid_available_ = load_nvcuvid_library();
@@ -295,11 +330,12 @@ void BetaCudaDeviceInterface::initialize_video(
     const VideoStreamOptions& video_stream_options,
     const std::vector<std::unique_ptr<Transform>>& transforms,
     const std::optional<FrameDims>& resized_output_dims) {
-  // TODO_API_BREAKDOWN ewwwww
+  // TODO_API_BREAKDOWN P0
   if (!av_stream) {
     return;
   }
   STD_TORCH_CHECK(av_stream != nullptr, "AVStream cannot be null");
+  CudaContextGuard context_guard(device_.index());
   rotation_ = rotation_from_degrees(get_rotation_from_stream(av_stream));
   output_dtype_ = video_stream_options.output_dtype;
 
@@ -403,6 +439,7 @@ void BetaCudaDeviceInterface::send_seqhdr_packet() {
 }
 
 BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
+  CudaContextGuard context_guard(device_.index());
   if (decoder_) {
     // DALI doesn't seem to do any particular cleanup of the decoder before
     // sending it to the cache, so we probably don't need to do anything either.
@@ -554,6 +591,7 @@ int BetaCudaDeviceInterface::stream_property_change(
 // Moral equivalent of avcodec_send_packet(). Here, we pass the AVPacket down to
 // the NVCUVID parser.
 int BetaCudaDeviceInterface::send_packet(ReferenceAVPacket& packet) {
+  CudaContextGuard context_guard(device_.index());
   if (cpu_fallback_) {
     return cpu_fallback_->send_packet(packet);
   }
@@ -581,6 +619,7 @@ int BetaCudaDeviceInterface::send_packet(ReferenceAVPacket& packet) {
 }
 
 int BetaCudaDeviceInterface::send_eof_packet() {
+  CudaContextGuard context_guard(device_.index());
   if (cpu_fallback_) {
     return cpu_fallback_->send_eof_packet();
   }
@@ -656,6 +695,7 @@ int BetaCudaDeviceInterface::frame_ready_in_display_order(
 
 // Moral equivalent of avcodec_receive_frame().
 int BetaCudaDeviceInterface::receive_frame(UniqueAVFrame& av_frame) {
+  CudaContextGuard context_guard(device_.index());
   if (cpu_fallback_) {
     return cpu_fallback_->receive_frame(av_frame);
   }
@@ -698,8 +738,8 @@ int BetaCudaDeviceInterface::receive_frame(UniqueAVFrame& av_frame) {
   // color-converted (with a copy), or that's a frame that was discarded in
   // SingleStreamDecoder. Either way, the underlying output surface can be
   // safely re-used.
-  // TODO_API_BREAKDOWN: We should update this comment slightly to now account
-  // for the frame copy we do in make_frame_standalone()
+  // TODO_API_BREAKDOWN P1: We should update this comment slightly to now
+  // account for the frame copy we do in make_frame_standalone()
   unmap_previous_frame();
   CUresult result = cuvidMapVideoFrame(
       *decoder_.get(),
@@ -807,7 +847,7 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame(
       reinterpret_cast<uint8_t*>(frame_ptr + (pitch * even_height));
   av_frame->data[2] = nullptr;
   av_frame->data[3] = nullptr;
-  // TODO_API_BREAKDOWN_CUDA: Check range before cast?
+  // TODO_API_BREAKDOWN_CUDA P2: Check range before cast?
   av_frame->linesize[0] = static_cast<int>(pitch);
   av_frame->linesize[1] = static_cast<int>(pitch);
   av_frame->linesize[2] = 0;
@@ -816,10 +856,8 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame(
   return av_frame;
 }
 
-// TODO_API_BREAKDOWN_CUDA: Does this even nede to be a method? Maybe it can be
-// a function that just lives in the PacketDecoder so we don't need to expose
-// another API to the DeviceInterface?
 void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
+  CudaContextGuard context_guard(device_.index());
   if (!(av_frame->format == AV_PIX_FMT_P016LE ||
         av_frame->format == AV_PIX_FMT_NV12)) {
     // The CPU frames are already standalone, so we don't need to do anything.
@@ -832,7 +870,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
   //           = num_pixels * 3 / 2
   //
   // To make it correct, we should use num_pixels = pitch * height, not
-  // num_pixels = pitch * height. The pitch value also accounts for the data
+  // num_pixels = width * height. The pitch value also accounts for the data
   // size (uint8 vs uint16) so this is also correct for P016.
   int64_t even_height =
       static_cast<int64_t>(round_up_to_even(av_frame->height));
@@ -842,7 +880,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
   auto storage =
       torch::stable::empty({num_bytes}, kStableUInt8, std::nullopt, device_);
 
-  // TODO_API_BREAKDOWN_CUDA: I suspect we don't need to wait on the nvdec
+  // TODO_API_BREAKDOWN_CUDA P1: I suspect we don't need to wait on the nvdec
   // stream here, because we can only arrive here from a path where the frame
   // has already been mapped so its data is available - worth double checking.
   cudaStream_t current_stream = get_current_cuda_stream(device_.index());
@@ -857,7 +895,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
       "Failed to copy NVDEC surface: ",
       cudaGetErrorString(err));
 
-  // TODO_API_BREAKDOWN_CUDA: Should we unmap here? Or let the next
+  // TODO_API_BREAKDOWN_CUDA P2: Should we unmap here? Or let the next
   // receive_frame() call do it?
   // unmap_previous_frame();
 
@@ -867,7 +905,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
 
   auto attached_data = new StandAloneFrameAttachedData();
   attached_data->producer_stream = current_stream;
-  // TODO_API_BREAKDOWN_CUDA: We don't *really* need to std::move it I guess?
+  // TODO_API_BREAKDOWN_CUDA P2: We don't *really* need to std::move it I guess?
   attached_data->storage = std::move(storage);
   av_frame->opaque_ref = av_buffer_create(
       reinterpret_cast<uint8_t*>(attached_data),
@@ -878,6 +916,7 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
 }
 
 void BetaCudaDeviceInterface::flush() {
+  CudaContextGuard context_guard(device_.index());
   if (cpu_fallback_) {
     cpu_fallback_->flush();
     return;
@@ -1044,9 +1083,10 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
     const AVFrame& av_frame,
     FrameOutput& frame_output,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
-  // TODO_API_BREAKDOWN_CUDA is that accurate and safe? Can there be a CPU NV12
-  // frame in our code? Should we create a helper used in the make_standalone
-  // function too?
+  CudaContextGuard context_guard(device_.index());
+  // TODO_API_BREAKDOWN_CUDA P0 is that accurate and safe? Can there be a CPU
+  // NV12 frame in our code? Should we create a helper used in the
+  // make_standalone function too?
   bool cpu_fallback = av_frame.format != AV_PIX_FMT_NV12 &&
       av_frame.format != AV_PIX_FMT_P016LE;
 
@@ -1062,7 +1102,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
         av_pix_fmt_desc_get(static_cast<AVPixelFormat>(av_frame.format));
     bool is444 = desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0;
     if (is444) {
-      // TODO_API_BREAKDOWN we need to handle this
+      // TODO_API_BREAKDOWN P1: we need to handle this
       FrameOutput cpu_frame_output;
       cpu_fallback_->convert_av_frame_to_frame_output(
           av_frame, cpu_frame_output);
@@ -1098,7 +1138,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
           gpu_frame.format == AV_PIX_FMT_P016LE,
       "Expected NV12 or P016LE format frame");
 
-  // TODO_API_BREAKDOWN: Cleanup how we get the attached data? Make it more
+  // TODO_API_BREAKDOWN P1: Cleanup how we get the attached data? Make it more
   // robust? Should we couple it to a flag on the interface saying "I'm
   // color-conversion only, I absolutely expect frames to be standalone"?
   cudaStream_t producer_stream;
@@ -1111,7 +1151,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
     producer_stream = get_current_cuda_stream(device_.index());
   }
 
-  // TODO_API_BREAKDOWN: we don't suppor output_dtype so some of that is not
+  // TODO_API_BREAKDOWN P1: we don't suppor output_dtype so some of that is not
   // execrcized.
   auto convert_frame = [&](std::optional<torch::stable::Tensor> pre_alloc)
       -> torch::stable::Tensor {

src/torchcodec/_core/ColorConverter.cpp

@@ -25,14 +25,14 @@ ColorConverter::ColorConverter(const StableDevice& device) {
   options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet
   options.device = device;
 
-  // TODO_API_BREAKDOWN It seems unnatural that the color-converter needs its
+  // TODO_API_BREAKDOWN P1 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. Should at the very least write a note about this design that now
   // the DeviceInterface has different modes: decode only, color-convert only,
   // and decode+color-convert (which used to be the only mode).
 
-  // TODO_API_BREAKDOWN: we shouldn't call initialize_video here, this is for
+  // TODO_API_BREAKDOWN P0: we shouldn't call initialize_video here, this is for
   // the decoding+color-convert mode. We should do something cleaner e.g.
   // initialize_color_convertion_only()
   std::vector<std::unique_ptr<Transform>> no_transforms;

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -86,7 +86,7 @@ void CpuDeviceInterface::initialize_video(
     const VideoStreamOptions& video_stream_options,
     const std::vector<std::unique_ptr<Transform>>& transforms,
     const std::optional<FrameDims>& resized_output_dims) {
-  // TODO_API_BREAKDOWN this used to be:
+  // TODO_API_BREAKDOWN P0 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)

src/torchcodec/_core/Demuxer.cpp

@@ -69,7 +69,7 @@ Demuxer::Demuxer(
 }
 
 UniqueAVPacket Demuxer::next_packet() {
-  // TODO_API_BREAKDOWN: Not a fan of the ReferenceAVPacket / AutoAVPacket /
+  // TODO_API_BREAKDOWN P2: Not a fan of the ReferenceAVPacket / AutoAVPacket /
   // UniqueAVPacket dance here. Can we simplify?
   ReferenceAVPacket packet(auto_packet_);
   int status =

src/torchcodec/_core/PacketDecoder.cpp

@@ -8,7 +8,7 @@
 
 namespace facebook::torchcodec {
 
-// TODO_API_BREAKDOWN: we should make sure the block APIs can dispatch to
+// TODO_API_BREAKDOWN P1: we should make sure the block APIs can dispatch to
 // third-party extensions - all of them.
 
 SharedAVCodecContext create_and_open_codec_context(
@@ -75,7 +75,7 @@ PacketDecoder::PacketDecoder(
   options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet
   options.device = device;
 
-  // TODO_API_BREAKDOWN: This isn't right, it's needed only for the NVDEC
+  // TODO_API_BREAKDOWN P0: This isn't right, it's needed only for the NVDEC
   // interface. This should probably be initialize_video_only - there's a
   // sibling TODO in the ColorConverter code (about color-conversion only.)
   std::vector<std::unique_ptr<Transform>> no_transforms;

src/torchcodec/_core/custom_ops.cpp

@@ -827,7 +827,7 @@ int64_t _blocks_packet_decoder_send_packet(
     torch::stable::Tensor& decoder,
     torch::stable::Tensor& packet) {
   PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer<PacketDecoder>(decoder);
-  // TODO_API_BREAKDOWN: Do we really need this to be a raw AVPacket*?
+  // TODO_API_BREAKDOWN P1: Do we really need this to be a raw AVPacket*?
   AVPacket* raw_packet = unwrap_tensor_to_pointer<AVPacket>(packet);
   return static_cast<int64_t>(decoder_ptr->send_packet(raw_packet));
 }

src/torchcodec/decoders/_blocks/_color_converter.py

@@ -11,10 +11,10 @@
 
 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!!
+# TODO_API_BREAKDOWN FEAT support output_dtype?
+# TODO_API_BREAKDOWN FEAT We need to support rotation metadata!!
+# TODO_API_BREAKDOWN FEAT Implement seeking?
+# TODO_API_BREAKDOWN FEAT Implement range-getting (start/end time) for decoding?
 
 
 class ColorConverter:
@@ -30,8 +30,8 @@ class ColorConverter:
     block is intentionally stream-agnostic.
     """
 
-    # TODO_API_BREAKDOWN: device default should be None
-    # TODO_API_BREAKDOWN: add checks for coupling between device param of
+    # TODO_API_BREAKDOWN P1: device default should be None
+    # TODO_API_BREAKDOWN P1: add checks for coupling between device param of
     # PacketDecoder and ColorConverter. What if one is CPU and the other is
     # CUDA? What if they're different CUDA devices? Maybe we should just error.
     def __init__(self, device="cpu"):

src/torchcodec/decoders/_blocks/_packet_decoder.py

@@ -17,7 +17,7 @@
 from ._frame import DecodedFrame, Packet
 
 
-# TODO_API_BREAKDOWN revisit every single docstring / comments at some point.
+# TODO_API_BREAKDOWN P2 revisit every single docstring / comments at some point.
 
 
 class PacketDecoder:
@@ -31,7 +31,7 @@ class PacketDecoder:
     on your own threads.
     """
 
-    # TODO_API_BREAKDOWN: device default should be None
+    # TODO_API_BREAKDOWN P1: device default should be None, here and everywhere else
     def __init__(self, demuxer: Demuxer, device="cpu"):
         self._handle = _blocks_create_packet_decoder(
             demuxer._handle, num_threads=1, device=device
@@ -56,7 +56,7 @@ 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
+    # TODO_API_BREAKDOWN P2 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

test/test_decoders.py

@@ -3381,20 +3381,9 @@ def _decode_sequential(self, path, device):
 
     def _decode_prefetch_frames(self, path, device):
         # [demux + decode] on one thread || [color-convert] on another.
-        demuxer = Demuxer(path)
-        converter = ColorConverter(device=device)
+        demuxer, decoder, converter = self._make_blocks(path, device)
 
-        def demux_and_decode():
-            # Constructed here so the PacketDecoder (and thus the CUDA context it
-            # binds via its device interface) lives on the prefetch worker
-            # thread that actually runs the cuvid* calls.
-            # TODO_API_BREAKDOWN_CUDA: this is a temporary workaroun for a real
-            # issue - needs fixing. Things should work when the objects are
-            # constructed in a different thread than where they're consumed.
-            decoder = PacketDecoder(demuxer, device=device)
-            yield from self._decode(decoder, self._demux(demuxer))
-
-        frames = self.prefetch(demux_and_decode())
+        frames = self.prefetch(self._decode(decoder, self._demux(demuxer)))
         return list(self._convert(converter, frames))
 
     def _decode_prefetch_packets(self, path, device):
@@ -3405,19 +3394,9 @@ def _decode_prefetch_packets(self, path, device):
 
     def _decode_prefetch_packets_and_frames(self, path, device):
         # [demux] || [decode] || [color-convert], each on its own thread.
-        demuxer = Demuxer(path)
-        converter = ColorConverter(device=device)
+        demuxer, decoder, converter = self._make_blocks(path, device)
         packets = self.prefetch(self._demux(demuxer))
-
-        def decode(packets):
-            # Constructed here so the PacketDecoder (and thus the CUDA context it
-            # binds via its device interface) lives on the prefetch worker
-            # thread that actually runs the cuvid* calls.
-            # TODO_API_BREAKDOWN_CUDA: same TODO as above.
-            decoder = PacketDecoder(demuxer, device=device)
-            yield from self._decode(decoder, packets)
-
-        frames = self.prefetch(decode(packets))
+        frames = self.prefetch(self._decode(decoder, packets))
         return list(self._convert(converter, frames))
 
     def _to_frame_batch(self, frames):
@@ -3431,6 +3410,8 @@ def _to_frame_batch(self, frames):
             ),
         )
 
+    # TODO_API_BREAKDOWN P0: We need to test all assets. Generally we need more
+    # tests for all features / edge cases that were eventually fixed.
     @pytest.mark.parametrize(
         "video",
         (
@@ -3499,7 +3480,7 @@ def test_color_converter_reused_across_videos(self, device):
     def test_set_cuda_backend_is_a_noop(self, device):
         # The blocks always use the NVDEC CUDA backend. Asking for the "ffmpeg"
         # one changes nothing, rather than silently producing something else.
-        # TODO_API_BREAKDOWN: let's just error?
+        # TODO_API_BREAKDOWN P2: let's just error?
         with set_cuda_backend("ffmpeg"):
             got = self._to_frame_batch(self._decode_sequential(NASA_VIDEO.path, device))
         ref = self._to_frame_batch(self._decode_sequential(NASA_VIDEO.path, device))