diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 45efcb643..ce37ff4f9 100644 --- a/CHANGELOG.rst +++ b/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: diff --git a/av/video/frame.pxd b/av/video/frame.pxd index 0a96311e0..66926058d 100644 --- a/av/video/frame.pxd +++ b/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) diff --git a/av/video/frame.py b/av/video/frame.py index 4a85d5b11..c6c2f26a1 100644 --- a/av/video/frame.py +++ b/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,6 +84,14 @@ 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__( @@ -36,14 +99,27 @@ def __cinit__( 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 diff --git a/av/video/frame.pyi b/av/video/frame.pyi index dd6f481d1..f07a8079f 100644 --- a/av/video/frame.pyi +++ b/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): diff --git a/include/avutil.pxd b/include/avutil.pxd index 8d671c98b..8818fdd14 100644 --- a/include/avutil.pxd +++ b/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) diff --git a/include/hwcontext_cuda.pxd b/include/hwcontext_cuda.pxd new file mode 100644 index 000000000..b97962967 --- /dev/null +++ b/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 + #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 diff --git a/tests/test_dlpack.py b/tests/test_dlpack.py index 15a513e6d..b02675525 100644 --- a/tests/test_dlpack.py +++ b/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}")