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..41fa04b9e 100644 --- a/av/video/frame.py +++ b/av/video/frame.py @@ -16,7 +16,21 @@ 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.final @@ -29,6 +43,9 @@ 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. """ def __cinit__( @@ -36,14 +53,22 @@ 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) def __dealloc__(self): ref: cython.pointer[lib.AVBufferRef] @@ -64,10 +89,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 +102,46 @@ 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 + + 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..0c11fb4fb 100644 --- a/tests/test_dlpack.py +++ b/tests/test_dlpack.py @@ -533,13 +533,23 @@ 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) def test_video_frame_from_dlpack_cuda_hw_frame_behavior_if_available() -> None: @@ -679,3 +689,52 @@ 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 + producer_stream = torch.cuda.Stream() + encoder_stream = torch.cuda.Stream() + encoder_stream_ptr = int(encoder_stream.cuda_stream) + + try: + with torch.cuda.stream(producer_stream): + y = torch.zeros((height, width), dtype=torch.uint8, device="cuda") + uv = torch.zeros( + (height // 2, width // 2, 2), + dtype=torch.uint8, + device="cuda", + ) + 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}")