From 7ea3c555495dd63f3981c74069de9fae5b5947ae Mon Sep 17 00:00:00 2001
From: Seungha Yang <seungha@centricular.com>
Date: Tue, 13 Sep 2022 01:02:55 +0900
Subject: [PATCH] cudaupload,cudadownload: Use shared GstD3D11Device context if
 possible

Handle d3d11 device context in set_context() method with
additional device compatibility check so that only NVIDIA GPU
associated d3d11 device can be configured in the element.
And clear old d3d11 device per set_info() for d3d11 device to be
updated as well.

Part-of: <https://gitlab.freedesktop.org/gstreamer/gstreamer/-/merge_requests/3018>
---
 .../sys/nvcodec/gstcudamemorycopy.c           | 52 ++++++++++++++++++-
 1 file changed, 50 insertions(+), 2 deletions(-)

diff --git a/subprojects/gst-plugins-bad/sys/nvcodec/gstcudamemorycopy.c b/subprojects/gst-plugins-bad/sys/nvcodec/gstcudamemorycopy.c
index b766ec3b88..1d16590002 100644
--- a/subprojects/gst-plugins-bad/sys/nvcodec/gstcudamemorycopy.c
+++ b/subprojects/gst-plugins-bad/sys/nvcodec/gstcudamemorycopy.c
@@ -114,6 +114,11 @@ static gboolean gst_cuda_memory_copy_set_info (GstCudaBaseTransform * btrans,
     GstVideoInfo * out_info);
 static GstFlowReturn gst_cuda_memory_copy_transform (GstBaseTransform * trans,
     GstBuffer * inbuf, GstBuffer * outbuf);
+#ifdef GST_CUDA_HAS_D3D
+static gboolean
+gst_cuda_memory_copy_ensure_d3d11_interop (GstCudaContext * context,
+    GstD3D11Device * device);
+#endif
 
 static void
 gst_cuda_memory_copy_class_init (GstCudaMemoryCopyClass * klass)
@@ -147,12 +152,51 @@ static void
 gst_cuda_memory_copy_set_context (GstElement * element, GstContext * context)
 {
   /* CUDA context is handled by parent class, handle only non-CUDA context */
-#ifdef HAVE_NVCODEC_GST_GL
+#if defined (HAVE_NVCODEC_GST_GL) || defined (GST_CUDA_HAS_D3D)
   GstCudaMemoryCopy *self = GST_CUDA_MEMORY_COPY (element);
 
+#ifdef HAVE_NVCODEC_GST_GL
   gst_gl_handle_set_context (element, context, &self->gl_display,
       &self->other_gl_context);
-#endif
+#endif /* HAVE_NVCODEC_GST_GL */
+
+#ifdef GST_CUDA_HAS_D3D
+  GstCudaBaseTransform *base = GST_CUDA_BASE_TRANSFORM (element);
+  if (gst_d3d11_handle_set_context (element, context, -1, &self->d3d11_device)) {
+    gboolean compatible = TRUE;
+
+    if (base->context) {
+      if (!gst_cuda_memory_copy_ensure_d3d11_interop (base->context,
+              self->d3d11_device)) {
+        GST_INFO_OBJECT (self, "%" GST_PTR_FORMAT
+            " is not CUDA compatible with %" GST_PTR_FORMAT,
+            self->d3d11_device, base->context);
+        compatible = FALSE;
+      }
+    } else {
+      guint device_count = 0;
+      CUdevice device_list[1] = { 0, };
+      CUresult cuda_ret;
+
+      cuda_ret = CuD3D11GetDevices (&device_count, device_list, 1,
+          gst_d3d11_device_get_device_handle (self->d3d11_device),
+          CU_D3D11_DEVICE_LIST_ALL);
+      if (cuda_ret != CUDA_SUCCESS || device_count == 0) {
+        GST_INFO_OBJECT (self, "%" GST_PTR_FORMAT " is not CUDA compatible",
+            self->d3d11_device);
+        compatible = FALSE;
+      }
+    }
+
+    if (!compatible) {
+      gst_clear_object (&self->d3d11_device);
+    } else {
+      GST_INFO_OBJECT (self, "%" GST_PTR_FORMAT " is CUDA compatible",
+          self->d3d11_device);
+    }
+  }
+#endif /* GST_CUDA_HAS_D3D */
+#endif /* defined (HAVE_NVCODEC_GST_GL) || defined (GST_CUDA_HAS_D3D) */
 
   GST_ELEMENT_CLASS (parent_class)->set_context (element, context);
 }
@@ -790,6 +834,10 @@ gst_cuda_memory_copy_set_info (GstCudaBaseTransform * btrans,
           GST_CAPS_FEATURE_MEMORY_D3D11_MEMORY)) {
     self->out_type = GST_CUDA_BUFFER_COPY_D3D11;
   }
+
+  /* Clear d3d11 device, this set_info() might be called to update
+   * cuda context and therefore d3d11 device object should be updated as well */
+  gst_clear_object (&self->d3d11_device);
 #endif
 
 #ifdef HAVE_NVCODEC_NVMM