[FFmpeg-cvslog] avcodec/nvenc: pass CUstream to nvenc when available

[Timo Rothenpieler]

ffmpeg | branch: master | Timo Rothenpieler <timo at rothenpieler.org> | Fri Sep 27 18:30:10 2019 +0200| [51a23343d9c736217d8845b7442fafc373726433] | committer: Timo Rothenpieler

avcodec/nvenc: pass CUstream to nvenc when available

Signed-off-by: Timo Rothenpieler <timo at rothenpieler.org>

> http://git.videolan.org/gitweb.cgi/ffmpeg.git/?a=commit;h=51a23343d9c736217d8845b7442fafc373726433
---

 libavcodec/nvenc.c | 20 ++++++++++++++++----
 libavcodec/nvenc.h |  2 ++
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/libavcodec/nvenc.c b/libavcodec/nvenc.c
index 8a7347b84e..2de2bfdacb 100644
--- a/libavcodec/nvenc.c
+++ b/libavcodec/nvenc.c
@@ -460,6 +460,7 @@ static av_cold int nvenc_check_device(AVCodecContext *avctx, int idx)
         goto fail;
 
     ctx->cu_context = ctx->cu_context_internal;
+    ctx->cu_stream = NULL;
 
     if ((ret = nvenc_pop_context(avctx)) < 0)
         goto fail2;
@@ -546,6 +547,7 @@ static av_cold int nvenc_setup_device(AVCodecContext *avctx)
 
         if (cuda_device_hwctx) {
             ctx->cu_context = cuda_device_hwctx->cuda_ctx;
+            ctx->cu_stream = cuda_device_hwctx->stream;
         }
 #if CONFIG_D3D11VA
         else if (d3d11_device_hwctx) {
@@ -1245,15 +1247,25 @@ static av_cold int nvenc_setup_encoder(AVCodecContext *avctx)
         return res;
 
     nv_status = p_nvenc->nvEncInitializeEncoder(ctx->nvencoder, &ctx->init_encode_params);
+    if (nv_status != NV_ENC_SUCCESS) {
+        nvenc_pop_context(avctx);
+        return nvenc_print_error(avctx, nv_status, "InitializeEncoder failed");
+    }
+
+#ifdef NVENC_HAVE_CUSTREAM_PTR
+    if (ctx->cu_context) {
+        nv_status = p_nvenc->nvEncSetIOCudaStreams(ctx->nvencoder, &ctx->cu_stream, &ctx->cu_stream);
+        if (nv_status != NV_ENC_SUCCESS) {
+            nvenc_pop_context(avctx);
+            return nvenc_print_error(avctx, nv_status, "SetIOCudaStreams failed");
+        }
+    }
+#endif
 
     res = nvenc_pop_context(avctx);
     if (res < 0)
         return res;
 
-    if (nv_status != NV_ENC_SUCCESS) {
-        return nvenc_print_error(avctx, nv_status, "InitializeEncoder failed");
-    }
-
     if (ctx->encode_config.frameIntervalP > 1)
         avctx->has_b_frames = 2;
 
diff --git a/libavcodec/nvenc.h b/libavcodec/nvenc.h
index 46ad7972bc..cdba998ffd 100644
--- a/libavcodec/nvenc.h
+++ b/libavcodec/nvenc.h
@@ -57,6 +57,7 @@ typedef void ID3D11Device;
 // SDK 9.1 compile time feature checks
 #if NVENCAPI_CHECK_VERSION(9, 1)
 #define NVENC_HAVE_MULTIPLE_REF_FRAMES
+#define NVENC_HAVE_CUSTREAM_PTR
 #endif
 
 typedef struct NvencSurface
@@ -132,6 +133,7 @@ typedef struct NvencContext
     NV_ENC_CONFIG encode_config;
     CUcontext cu_context;
     CUcontext cu_context_internal;
+    CUstream cu_stream;
     ID3D11Device *d3d11_device;
 
     int nb_surfaces;