feat(nvenc): add support for explicit cuda streams

CHANGELOG.rst

@@ -38,6 +38,7 @@ Features:
 - Add ``Packet.rescale_ts()`` to rescale packet PTS, DTS, and duration to a new ``AVRational`` time base by :gh-user:`WyattBlue`.
 - Support reusing the thread's current CUDA context via a ``current_ctx`` flag on ``CudaContext`` and ``VideoFrame.from_dlpack``, for interop with libraries like PyTorch that initialize CUDA first by :gh-user:`Yozer` (:pr:`2339`).
 - ``VideoFrame.from_dlpack`` no longer requires restating ``primary_ctx``/``current_ctx`` when passing an explicit ``cuda_context``; the flags are only validated when explicitly given by :gh-user:`WyattBlue`.
+- Support passing an explicit CUDA stream to FFmpeg CUDA operations, including NVENC input and output, via a ``cuda_stream`` parameter on ``CudaContext``; currently limited to logical CUDA device 0 by :gh-user:`Yozer` (:pr:`2360`).
 
 Fixes:
 

av/video/frame.pxd

@@ -1,5 +1,5 @@
 cimport libav as lib
-from libc.stdint cimport uint8_t
+from libc.stdint cimport uint8_t, uintptr_t
 
 from av.frame cimport Frame
 from av.video.format cimport VideoFormat
@@ -10,6 +10,7 @@ cdef class CudaContext:
     cdef readonly int device_id
     cdef readonly bint primary_ctx
     cdef readonly bint current_ctx
+    cdef readonly uintptr_t cuda_stream
     cdef lib.AVBufferRef* _device_ref
     cdef dict _frames_cache
     cdef lib.AVBufferRef* _get_device_ref(self)

av/video/frame.py

@@ -16,7 +16,62 @@
     PyCapsule_SetName,
 )
 from cython.cimports.cpython.ref import Py_DECREF, Py_INCREF
-from cython.cimports.libc.stdint import int64_t, uint8_t
+from cython.cimports.hwcontext_cuda import AVCUDADeviceContext, CUstream
+from cython.cimports.libc.stdint import int64_t, uint8_t, uintptr_t
+
+
+@cython.cfunc
+@cython.nogil
+@cython.exceptval(check=False)
+def _cuda_device_ctx_free(
+    device_ctx: cython.pointer[lib.AVHWDeviceContext],
+) -> cython.void:
+    owner_ref: cython.pointer[lib.AVBufferRef] = cython.cast(
+        cython.pointer[lib.AVBufferRef], device_ctx.user_opaque
+    )
+    lib.av_buffer_unref(cython.address(owner_ref))
+    device_ctx.user_opaque = cython.NULL
+
+
+@cython.cfunc
+def _check_current_ctx_is_device0_primary(
+    owner_ref: cython.pointer[lib.AVBufferRef],
+) -> cython.void:
+    # The manually initialized device context below cannot populate FFmpeg's
+    # private cuda_device field, which stays 0. Requiring the wrapped context
+    # to be the primary context of device 0 keeps that metadata correct.
+    probe_ref: cython.pointer[lib.AVBufferRef] = cython.NULL
+    probe_device: cython.p_char = b"0"
+    options: Dictionary = Dictionary({"primary_ctx": "1"})
+    err_check(
+        lib.av_hwdevice_ctx_create(
+            cython.address(probe_ref),
+            lib.AV_HWDEVICE_TYPE_CUDA,
+            probe_device,
+            options.ptr,
+            0,
+        )
+    )
+    owner_device_ctx: cython.pointer[lib.AVHWDeviceContext] = cython.cast(
+        cython.pointer[lib.AVHWDeviceContext], owner_ref.data
+    )
+    owner_cuda_ctx: cython.pointer[AVCUDADeviceContext] = cython.cast(
+        cython.pointer[AVCUDADeviceContext], owner_device_ctx.hwctx
+    )
+    probe_device_ctx: cython.pointer[lib.AVHWDeviceContext] = cython.cast(
+        cython.pointer[lib.AVHWDeviceContext], probe_ref.data
+    )
+    probe_cuda_ctx: cython.pointer[AVCUDADeviceContext] = cython.cast(
+        cython.pointer[AVCUDADeviceContext], probe_device_ctx.hwctx
+    )
+    matches: cython.bint = probe_cuda_ctx.cuda_ctx == owner_cuda_ctx.cuda_ctx
+    lib.av_buffer_unref(cython.address(probe_ref))
+    if not matches:
+        raise ValueError(
+            "cuda_stream with current_ctx requires the calling thread's current "
+            "CUDA context to be the primary context of CUDA device 0; expose the "
+            "desired physical GPU as logical device 0 via CUDA_VISIBLE_DEVICES"
+        )
 
 
 @cython.final
@@ -29,21 +84,42 @@ class CudaContext:
     :param bool current_ctx: Wrap the CUDA context current on the calling thread
         when this object is first used. This is mutually exclusive with
         ``primary_ctx``.
+    :param int cuda_stream: Optional raw ``CUstream`` pointer for FFmpeg CUDA
+        operations, including NVENC input and output. The caller owns the
+        stream and must keep it alive as long as this context can be used.
+        Nonzero stream pointers are currently supported only for logical CUDA
+        device 0; with ``current_ctx``, the calling thread's current CUDA
+        context must be the primary context of device 0 (the context used by
+        runtime-API libraries such as PyTorch or CuPy). To use another physical
+        GPU, expose it as logical device 0 via ``CUDA_VISIBLE_DEVICES``.
     """
 
     def __cinit__(
         self,
         device_id: cython.int = 0,
         primary_ctx: cython.bint = True,
         current_ctx: cython.bint = False,
+        cuda_stream: object = None,
     ):
         self.device_id = device_id
         self.primary_ctx = primary_ctx
         self.current_ctx = current_ctx
+        self.cuda_stream = 0
         self._device_ref = cython.NULL
         self._frames_cache = {}
         if primary_ctx and current_ctx:
             raise ValueError("primary_ctx and current_ctx are mutually exclusive")
+        if cuda_stream is not None:
+            if not isinstance(cuda_stream, int):
+                raise TypeError("cuda_stream must be an integer or None")
+            if cuda_stream < 0:
+                raise ValueError("cuda_stream must be non-negative")
+            self.cuda_stream = cython.cast(uintptr_t, cuda_stream)
+        if self.cuda_stream and self.device_id != 0:
+            raise ValueError(
+                "cuda_stream currently requires device_id 0; expose the desired "
+                "physical GPU as logical device 0 via CUDA_VISIBLE_DEVICES"
+            )
 
     def __dealloc__(self):
         ref: cython.pointer[lib.AVBufferRef]
@@ -64,10 +140,11 @@ def __dealloc__(self):
     @cython.cfunc
     def _get_device_ref(self) -> cython.pointer[lib.AVBufferRef]:
         device_ref: cython.pointer[lib.AVBufferRef] = self._device_ref
+        owner_ref: cython.pointer[lib.AVBufferRef]
         if device_ref != cython.NULL:
             return device_ref
 
-        device_ref = cython.NULL
+        owner_ref = cython.NULL
         device_bytes = f"{self.device_id}".encode()
         c_device: cython.p_char = device_bytes
         options_dict = {"primary_ctx": "1" if self.primary_ctx else "0"}
@@ -76,13 +153,53 @@ def _get_device_ref(self) -> cython.pointer[lib.AVBufferRef]:
         options: Dictionary = Dictionary(options_dict)
         err_check(
             lib.av_hwdevice_ctx_create(
-                cython.address(device_ref),
+                cython.address(owner_ref),
                 lib.AV_HWDEVICE_TYPE_CUDA,
                 c_device,
                 options.ptr,
                 0,
             )
         )
+        if not self.cuda_stream:
+            self._device_ref = owner_ref
+            return owner_ref
+
+        if self.current_ctx:
+            try:
+                _check_current_ctx_is_device0_primary(owner_ref)
+            except Exception:
+                lib.av_buffer_unref(cython.address(owner_ref))
+                raise
+
+        device_ref = lib.av_hwdevice_ctx_alloc(lib.AV_HWDEVICE_TYPE_CUDA)
+        if device_ref == cython.NULL:
+            lib.av_buffer_unref(cython.address(owner_ref))
+            raise MemoryError("av_hwdevice_ctx_alloc() failed")
+
+        try:
+            owner_device_ctx: cython.pointer[lib.AVHWDeviceContext] = cython.cast(
+                cython.pointer[lib.AVHWDeviceContext], owner_ref.data
+            )
+            owner_cuda_ctx: cython.pointer[AVCUDADeviceContext] = cython.cast(
+                cython.pointer[AVCUDADeviceContext], owner_device_ctx.hwctx
+            )
+            device_ctx: cython.pointer[lib.AVHWDeviceContext] = cython.cast(
+                cython.pointer[lib.AVHWDeviceContext], device_ref.data
+            )
+            cuda_ctx: cython.pointer[AVCUDADeviceContext] = cython.cast(
+                cython.pointer[AVCUDADeviceContext], device_ctx.hwctx
+            )
+            cuda_ctx.cuda_ctx = owner_cuda_ctx.cuda_ctx
+            cuda_ctx.stream = cython.cast(CUstream, self.cuda_stream)
+            device_ctx.free = _cuda_device_ctx_free
+            device_ctx.user_opaque = cython.cast(cython.p_void, owner_ref)
+            owner_ref = cython.NULL
+            err_check(lib.av_hwdevice_ctx_init(device_ref))
+        except Exception:
+            lib.av_buffer_unref(cython.address(device_ref))
+            lib.av_buffer_unref(cython.address(owner_ref))
+            raise
+
         self._device_ref = device_ref
         return device_ref
 

av/video/frame.pyi

@@ -36,11 +36,14 @@ class CudaContext:
     def primary_ctx(self) -> bool: ...
     @property
     def current_ctx(self) -> bool: ...
+    @property
+    def cuda_stream(self) -> int: ...
     def __init__(
         self,
         device_id: int = 0,
         primary_ctx: bool = True,
         current_ctx: bool = False,
+        cuda_stream: int | None = None,
     ) -> None: ...
 
 class VideoFrame(Frame):

include/avutil.pxd

@@ -192,7 +192,9 @@ cdef extern from "libavutil/frame.h" nogil:
 
 cdef extern from "libavutil/hwcontext.h" nogil:
     cdef struct AVHWDeviceContext:
-        pass
+        void *hwctx
+        void (*free)(AVHWDeviceContext *ctx)
+        void *user_opaque
 
     enum AVHWDeviceType:
         AV_HWDEVICE_TYPE_NONE
@@ -220,6 +222,8 @@ cdef extern from "libavutil/hwcontext.h" nogil:
         int width
         int height
 
+    cdef AVBufferRef *av_hwdevice_ctx_alloc(AVHWDeviceType type)
+    cdef int av_hwdevice_ctx_init(AVBufferRef *ref)
     cdef int av_hwdevice_ctx_create(AVBufferRef **device_ctx, AVHWDeviceType type, const char *device, AVDictionary *opts, int flags)
     cdef AVHWDeviceType av_hwdevice_find_type_by_name(const char *name)
     cdef const char *av_hwdevice_get_type_name(AVHWDeviceType type)

include/hwcontext_cuda.pxd

@@ -0,0 +1,19 @@
+cdef extern from *:
+    """
+    #ifndef CUDA_VERSION
+    #define PYAV_CUDA_TYPES
+    #define CUDA_VERSION 1
+    typedef struct CUctx_st *CUcontext;
+    typedef struct CUstream_st *CUstream;
+    #endif
+    #include <libavutil/hwcontext_cuda.h>
+    #ifdef PYAV_CUDA_TYPES
+    #undef CUDA_VERSION
+    #undef PYAV_CUDA_TYPES
+    #endif
+    """
+    ctypedef void *CUcontext
+    ctypedef void *CUstream
+    ctypedef struct AVCUDADeviceContext:
+        CUcontext cuda_ctx
+        CUstream stream

tests/test_dlpack.py

@@ -533,13 +533,33 @@ def test_cuda_context_flags() -> None:
     default_ctx = av.video.frame.CudaContext()
     assert default_ctx.primary_ctx is True
     assert default_ctx.current_ctx is False
+    assert default_ctx.cuda_stream == 0
 
     with pytest.raises(ValueError, match="mutually exclusive"):
         av.video.frame.CudaContext(primary_ctx=True, current_ctx=True)
 
-    ctx = av.video.frame.CudaContext(primary_ctx=False, current_ctx=True)
+    ctx = av.video.frame.CudaContext(
+        primary_ctx=False, current_ctx=True, cuda_stream=1234
+    )
     assert ctx.primary_ctx is False
     assert ctx.current_ctx is True
+    assert ctx.cuda_stream == 1234
+
+    with pytest.raises(TypeError, match="integer or None"):
+        av.video.frame.CudaContext(cuda_stream="1234")  # type: ignore[arg-type]
+
+    with pytest.raises(ValueError, match="non-negative"):
+        av.video.frame.CudaContext(cuda_stream=-1)
+
+    with pytest.raises(ValueError, match="requires device_id 0.*CUDA_VISIBLE_DEVICES"):
+        av.video.frame.CudaContext(device_id=1, cuda_stream=1234)
+
+    for default_stream in (None, 0):
+        nonzero_device_ctx = av.video.frame.CudaContext(
+            device_id=1, cuda_stream=default_stream
+        )
+        assert nonzero_device_ctx.device_id == 1
+        assert nonzero_device_ctx.cuda_stream == 0
 
 
 def test_video_frame_from_dlpack_cuda_hw_frame_behavior_if_available() -> None:
@@ -679,3 +699,53 @@ def test_encode_cuda_frame_with_nvenc_if_available(use_current_ctx: bool) -> Non
         assert any(p.size for p in packets)
     except av.FFmpegError as e:
         pytest.skip(f"nvenc/CUDA not available in this build/runtime: {e}")
+
+
+def test_encode_cuda_frame_with_nvenc_external_stream_if_available() -> None:
+    try:
+        import torch  # type: ignore
+    except Exception:
+        pytest.skip("PyTorch is not available.")
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available.")
+
+    width, height = 256, 256
+    with torch.cuda.device(0):
+        producer_stream = torch.cuda.Stream()
+        encoder_stream = torch.cuda.Stream()
+    encoder_stream_ptr = int(encoder_stream.cuda_stream)
+
+    try:
+        with torch.cuda.device(0), torch.cuda.stream(producer_stream):
+            y = torch.zeros((height, width), dtype=torch.uint8, device="cuda:0")
+            uv = torch.zeros(
+                (height // 2, width // 2, 2),
+                dtype=torch.uint8,
+                device="cuda:0",
+            )
+            cuda_context = av.video.frame.CudaContext(
+                device_id=int(y.__dlpack_device__()[1]),
+                primary_ctx=False,
+                current_ctx=True,
+                cuda_stream=encoder_stream_ptr,
+            )
+            frame = VideoFrame.from_dlpack(
+                (y, uv),
+                format="nv12",
+                stream=encoder_stream_ptr,
+                cuda_context=cuda_context,
+            )
+
+        assert cuda_context.cuda_stream == encoder_stream_ptr
+        cc = av.CodecContext.create("h264_nvenc", "w")
+        cc.width = width
+        cc.height = height
+        cc.time_base = Fraction(1, 24)
+        cc.framerate = Fraction(24, 1)
+        cc.pix_fmt = "cuda"
+
+        packets = cc.encode(frame)
+        packets += cc.encode(None)
+        assert any(p.size for p in packets)
+    except av.FFmpegError as e:
+        pytest.skip(f"nvenc/CUDA not available in this build/runtime: {e}")