From 7b2ad11903781f05b97a10bb9b3b7fd53d859289 Mon Sep 17 00:00:00 2001 From: Yozer Date: Sun, 26 Jul 2026 17:29:22 +0200 Subject: [PATCH 1/4] feat(nvenc): add support for explicit cuda streams --- av/video/frame.pxd | 3 ++- av/video/frame.py | 22 +++++++++++++++- av/video/frame.pyi | 3 +++ include/avutil.pxd | 2 +- tests/test_dlpack.py | 60 +++++++++++++++++++++++++++++++++++++++++++- 5 files changed, 86 insertions(+), 4 deletions(-) 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..d79d15eac 100644 --- a/av/video/frame.py +++ b/av/video/frame.py @@ -16,7 +16,8 @@ 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.final @@ -29,6 +30,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 +40,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] @@ -83,6 +95,14 @@ def _get_device_ref(self) -> cython.pointer[lib.AVBufferRef]: 0, ) ) + if self.cuda_stream: + 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.stream = cython.cast(CUstream, self.cuda_stream) 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..540a62d59 100644 --- a/include/avutil.pxd +++ b/include/avutil.pxd @@ -192,7 +192,7 @@ cdef extern from "libavutil/frame.h" nogil: cdef extern from "libavutil/hwcontext.h" nogil: cdef struct AVHWDeviceContext: - pass + void *hwctx enum AVHWDeviceType: AV_HWDEVICE_TYPE_NONE diff --git a/tests/test_dlpack.py b/tests/test_dlpack.py index 15a513e6d..643e2ef09 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") + + 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,51 @@ 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 + stream = torch.cuda.Stream() + stream_ptr = int(stream.cuda_stream) + + try: + with torch.cuda.stream(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=stream_ptr, + ) + frame = VideoFrame.from_dlpack( + (y, uv), + format="nv12", + stream=stream_ptr, + cuda_context=cuda_context, + ) + + assert cuda_context.cuda_stream == stream_ptr + cc = cast(VideoCodecContext, 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}") From 81cc489e2484ba6bfc6041bb0cb35eec335fa3eb Mon Sep 17 00:00:00 2001 From: Yozer Date: Sun, 26 Jul 2026 18:09:36 +0200 Subject: [PATCH 2/4] fix: add missing files --- include/hwcontext_cuda.pxd | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 include/hwcontext_cuda.pxd diff --git a/include/hwcontext_cuda.pxd b/include/hwcontext_cuda.pxd new file mode 100644 index 000000000..ccc21cb77 --- /dev/null +++ b/include/hwcontext_cuda.pxd @@ -0,0 +1,17 @@ +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 *CUstream + ctypedef struct AVCUDADeviceContext: + CUstream stream From 21f7e82fa1c72f52f27b97ba085c180628e3e54d Mon Sep 17 00:00:00 2001 From: Yozer Date: Sun, 26 Jul 2026 18:10:59 +0200 Subject: [PATCH 3/4] fix: external cuda stream test --- tests/test_dlpack.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/test_dlpack.py b/tests/test_dlpack.py index 643e2ef09..f02fd9738 100644 --- a/tests/test_dlpack.py +++ b/tests/test_dlpack.py @@ -700,11 +700,12 @@ def test_encode_cuda_frame_with_nvenc_external_stream_if_available() -> None: pytest.skip("CUDA is not available.") width, height = 256, 256 - stream = torch.cuda.Stream() - stream_ptr = int(stream.cuda_stream) + producer_stream = torch.cuda.Stream() + encoder_stream = torch.cuda.Stream() + encoder_stream_ptr = int(encoder_stream.cuda_stream) try: - with torch.cuda.stream(stream): + with torch.cuda.stream(producer_stream): y = torch.zeros((height, width), dtype=torch.uint8, device="cuda") uv = torch.zeros( (height // 2, width // 2, 2), @@ -715,17 +716,17 @@ def test_encode_cuda_frame_with_nvenc_external_stream_if_available() -> None: device_id=int(y.__dlpack_device__()[1]), primary_ctx=False, current_ctx=True, - cuda_stream=stream_ptr, + cuda_stream=encoder_stream_ptr, ) frame = VideoFrame.from_dlpack( (y, uv), format="nv12", - stream=stream_ptr, + stream=encoder_stream_ptr, cuda_context=cuda_context, ) - assert cuda_context.cuda_stream == stream_ptr - cc = cast(VideoCodecContext, av.CodecContext.create("h264_nvenc", "w")) + 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) From dbef919b14e5cd2ef3606b04faf5cc16fc8085f5 Mon Sep 17 00:00:00 2001 From: Yozer Date: Sun, 26 Jul 2026 18:24:45 +0200 Subject: [PATCH 4/4] fix: cuda stream type test --- tests/test_dlpack.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_dlpack.py b/tests/test_dlpack.py index f02fd9738..0c11fb4fb 100644 --- a/tests/test_dlpack.py +++ b/tests/test_dlpack.py @@ -546,7 +546,7 @@ def test_cuda_context_flags() -> None: assert ctx.cuda_stream == 1234 with pytest.raises(TypeError, match="integer or None"): - av.video.frame.CudaContext(cuda_stream="1234") + 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)