From 12ea34309b57ce15df9559ccc1aa055238df0a3e Mon Sep 17 00:00:00 2001
From: Roman Arzumanyan <r.arzumanyan@visionlabs.ai>
Date: Tue, 12 Sep 2023 17:31:30 +0300
Subject: [PATCH] libavutil/hwcontext_cuda: option added to use context which
 is current for calling thread

---
 doc/ffmpeg.texi            |  4 ++++
 libavutil/hwcontext_cuda.c | 26 +++++++++++++++++++++++++-
 libavutil/hwcontext_cuda.h |  5 +++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/doc/ffmpeg.texi b/doc/ffmpeg.texi
index 2273c39214..4b532eef89 100644
--- a/doc/ffmpeg.texi
+++ b/doc/ffmpeg.texi
@@ -1156,6 +1156,8 @@ The following options are recognized:
 @table @option
 @item primary_ctx
 If set to 1, uses the primary device context instead of creating a new one.
+@item current_ctx
+If set to 1, uses context current to calling thread instead of creating a new one.
 @end table
 
 Examples:
@@ -1165,6 +1167,8 @@ Choose the second device on the system.
 
 @item -init_hw_device cuda:0,primary_ctx=1
 Choose the first device and use the primary device context.
+@item -init_hw_device cuda:0,current_ctx=1
+Choose context current to calling thread.
 @end table
 
 @item dxva2
diff --git a/libavutil/hwcontext_cuda.c b/libavutil/hwcontext_cuda.c
index 4b298fa93e..ad9bf91717 100644
--- a/libavutil/hwcontext_cuda.c
+++ b/libavutil/hwcontext_cuda.c
@@ -361,6 +361,10 @@ static int cuda_context_init(AVHWDeviceContext *device_ctx, int flags) {
                                                     hwctx->internal->cuda_device));
         if (ret < 0)
             return ret;
+    } else if (flags & AV_CUDA_USE_CURRENT_CONTEXT) {
+        ret = CHECK_CU(cu->cuCtxGetCurrent(&hwctx->cuda_ctx));
+        if (ret < 0)
+            return ret;
     } else {
         ret = CHECK_CU(cu->cuCtxCreate(&hwctx->cuda_ctx, desired_flags,
                                        hwctx->internal->cuda_device));
@@ -381,19 +385,39 @@ static int cuda_context_init(AVHWDeviceContext *device_ctx, int flags) {
 static int cuda_flags_from_opts(AVHWDeviceContext *device_ctx,
                                 AVDictionary *opts, int *flags)
 {
+    int use_primary_ctx = 0, use_current_ctx = 0;
     AVDictionaryEntry *primary_ctx_opt = av_dict_get(opts, "primary_ctx", NULL, 0);
+    AVDictionaryEntry *current_ctx_opt = av_dict_get(opts, "current_ctx", NULL, 0);
+
+    if (primary_ctx_opt)
+        use_primary_ctx = strtol(primary_ctx_opt->value, NULL, 10);
+    
+    if (current_ctx_opt)
+        use_current_ctx = strtol(current_ctx_opt->value, NULL, 10);
 
-    if (primary_ctx_opt && strtol(primary_ctx_opt->value, NULL, 10)) {
+    if (use_primary_ctx && use_current_ctx)
+        av_log(device_ctx, AV_LOG_ERROR, "Can't use primary and current CUDA ctx simultaneously.\n");
+
+    if (use_primary_ctx) {
         av_log(device_ctx, AV_LOG_VERBOSE, "Using CUDA primary device context\n");
         *flags |= AV_CUDA_USE_PRIMARY_CONTEXT;
     } else if (primary_ctx_opt) {
         av_log(device_ctx, AV_LOG_VERBOSE, "Disabling use of CUDA primary device context\n");
         *flags &= ~AV_CUDA_USE_PRIMARY_CONTEXT;
+    } 
+    
+    if (use_current_ctx) {
+        av_log(device_ctx, AV_LOG_VERBOSE, "Using CUDA current device context\n");
+        *flags |= AV_CUDA_USE_CURRENT_CONTEXT;
+    } else if (current_ctx_opt) {
+        av_log(device_ctx, AV_LOG_VERBOSE, "Disabling use of CUDA current device context\n");
+        *flags &= ~AV_CUDA_USE_CURRENT_CONTEXT;
     }
 
     return 0;
 }
 
+
 static int cuda_device_create(AVHWDeviceContext *device_ctx,
                               const char *device,
                               AVDictionary *opts, int flags)
diff --git a/libavutil/hwcontext_cuda.h b/libavutil/hwcontext_cuda.h
index cefbe0ceab..cbad434fea 100644
--- a/libavutil/hwcontext_cuda.h
+++ b/libavutil/hwcontext_cuda.h
@@ -62,6 +62,11 @@ typedef struct AVCUDADeviceContext {
  */
 #define AV_CUDA_USE_PRIMARY_CONTEXT (1 << 0)
 
+/**
+ * Use current device context instead of creating a new one.
+ */
+#define AV_CUDA_USE_CURRENT_CONTEXT (1 << 1)
+
 /**
  * @}
  */
-- 
2.39.2 (Apple Git-143)