avutil/hwcontext_cuda: allow using primary CUDA device context

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

4 files changed, 59 insertions, 11 deletions

diff --git a/libavutil/hwcontext_cuda.c b/libavutil/hwcontext_cuda.c
index cca39e9fc7..30611b1912 100644
--- a/libavutil/hwcontext_cuda.c
+++ b/libavutil/hwcontext_cuda.c
@@ -280,10 +280,16 @@ static void cuda_device_uninit(AVHWDeviceContext *device_ctx)
 
     if (hwctx->internal) {
         CudaFunctions *cu = hwctx->internal->cuda_dl;
+
         if (hwctx->internal->is_allocated && hwctx->cuda_ctx) {
-            CHECK_CU(cu->cuCtxDestroy(hwctx->cuda_ctx));
+            if (hwctx->internal->flags & AV_CUDA_USE_PRIMARY_CONTEXT)
+                CHECK_CU(cu->cuDevicePrimaryCtxRelease(hwctx->internal->cuda_device));
+            else
+                CHECK_CU(cu->cuCtxDestroy(hwctx->cuda_ctx));
+
             hwctx->cuda_ctx = NULL;
         }
+
         cuda_free_functions(&hwctx->internal->cuda_dl);
     }
 
@@ -322,9 +328,11 @@ static int cuda_device_create(AVHWDeviceContext *device_ctx,
 {
     AVCUDADeviceContext *hwctx = device_ctx->hwctx;
     CudaFunctions *cu;
-    CUdevice cu_device;
     CUcontext dummy;
-    int ret, device_idx = 0;
+    int ret, dev_active = 0, device_idx = 0;
+    unsigned int dev_flags = 0;
+
+    const unsigned int desired_flags = CU_CTX_SCHED_BLOCKING_SYNC;
 
     if (device)
         device_idx = strtol(device, NULL, 0);
@@ -338,21 +346,42 @@ static int cuda_device_create(AVHWDeviceContext *device_ctx,
     if (ret < 0)
         goto error;
 
-    ret = CHECK_CU(cu->cuDeviceGet(&cu_device, device_idx));
+    ret = CHECK_CU(cu->cuDeviceGet(&hwctx->internal->cuda_device, device_idx));
     if (ret < 0)
         goto error;
 
-    ret = CHECK_CU(cu->cuCtxCreate(&hwctx->cuda_ctx, CU_CTX_SCHED_BLOCKING_SYNC, cu_device));
-    if (ret < 0)
-        goto error;
+    hwctx->internal->flags = flags;
 
-    // Setting stream to NULL will make functions automatically use the default CUstream
-    hwctx->stream = NULL;
+    if (flags & AV_CUDA_USE_PRIMARY_CONTEXT) {
+        ret = CHECK_CU(cu->cuDevicePrimaryCtxGetState(hwctx->internal->cuda_device, &dev_flags, &dev_active));
+        if (ret < 0)
+            goto error;
 
-    CHECK_CU(cu->cuCtxPopCurrent(&dummy));
+        if (dev_active && dev_flags != desired_flags) {
+            av_log(device_ctx, AV_LOG_ERROR, "Primary context already active with incompatible flags.\n");
+            goto error;
+        } else if (dev_flags != desired_flags) {
+            ret = CHECK_CU(cu->cuDevicePrimaryCtxSetFlags(hwctx->internal->cuda_device, desired_flags));
+            if (ret < 0)
+                goto error;
+        }
+
+        ret = CHECK_CU(cu->cuDevicePrimaryCtxRetain(&hwctx->cuda_ctx, hwctx->internal->cuda_device));
+        if (ret < 0)
+            goto error;
+    } else {
+        ret = CHECK_CU(cu->cuCtxCreate(&hwctx->cuda_ctx, desired_flags, hwctx->internal->cuda_device));
+        if (ret < 0)
+            goto error;
+
+        CHECK_CU(cu->cuCtxPopCurrent(&dummy));
+    }
 
     hwctx->internal->is_allocated = 1;
 
+    // Setting stream to NULL will make functions automatically use the default CUstream
+    hwctx->stream = NULL;
+
     return 0;
 
 error:
diff --git a/libavutil/hwcontext_cuda.h b/libavutil/hwcontext_cuda.h
index 81a0552cab..cefbe0ceab 100644
--- a/libavutil/hwcontext_cuda.h
+++ b/libavutil/hwcontext_cuda.h
@@ -49,4 +49,21 @@ typedef struct AVCUDADeviceContext {
  * AVHWFramesContext.hwctx is currently not used
  */
 
+/**
+ * @defgroup hwcontext_cuda Device context creation flags
+ *
+ * Flags for av_hwdevice_ctx_create.
+ *
+ * @{
+ */
+
+/**
+ * Use primary device context instead of creating a new one.
+ */
+#define AV_CUDA_USE_PRIMARY_CONTEXT (1 << 0)
+
+/**
+ * @}
+ */
+
 #endif /* AVUTIL_HWCONTEXT_CUDA_H */
diff --git a/libavutil/hwcontext_cuda_internal.h b/libavutil/hwcontext_cuda_internal.h
index e1bc6ff350..d5633c58d5 100644
--- a/libavutil/hwcontext_cuda_internal.h
+++ b/libavutil/hwcontext_cuda_internal.h
@@ -31,6 +31,8 @@
 struct AVCUDADeviceContextInternal {
     CudaFunctions *cuda_dl;
     int is_allocated;
+    CUdevice cuda_device;
+    int flags;
 };
 
 #endif /* AVUTIL_HWCONTEXT_CUDA_INTERNAL_H */
diff --git a/libavutil/version.h b/libavutil/version.h
index af3abf7265..e18163388d 100644
--- a/libavutil/version.h
+++ b/libavutil/version.h
@@ -80,7 +80,7 @@
 
 #define LIBAVUTIL_VERSION_MAJOR  56
 #define LIBAVUTIL_VERSION_MINOR  36
-#define LIBAVUTIL_VERSION_MICRO 100
+#define LIBAVUTIL_VERSION_MICRO 101
 
 #define LIBAVUTIL_VERSION_INT   AV_VERSION_INT(LIBAVUTIL_VERSION_MAJOR, \
                                                LIBAVUTIL_VERSION_MINOR, \