Merge branch 'issue_203'

Author: rarzumanyan

PyNvCodec/src/PyNvCodec.cpp

@@ -57,25 +57,22 @@ static auto ThrowOnCudaError = [](CUresult res, int lineNum = -1) {
 
 class CudaResMgr {
   CudaResMgr() {
+    lock_guard<mutex> lock(gContextsMutex);
+    lock_guard<mutex> lock(gStreamsMutex);
+
     ThrowOnCudaError(cuInit(0), __LINE__);
 
     int nGpu;
     ThrowOnCudaError(cuDeviceGetCount(&nGpu), __LINE__);
-    {
-      lock_guard<mutex> lock(gContextsMutex);
-      for (int i = 0; i < nGpu; i++) {
-        CUcontext cuContext = nullptr;
 
-        g_Contexts.push_back(cuContext);
-      }
-    }
-    {
-      lock_guard<mutex> lock(gStreamsMutex);
-      for (int i = 0; i < nGpu; i++) {
-        CUstream cuStream = nullptr;
 
-        g_Streams.push_back(cuStream);
-      }
+    for (int i = 0; i < nGpu; i++) {
+      CUdevice cuDevice = 0;
+      CUcontext cuContext = nullptr;
+      g_Contexts.push_back(make_pair(cuDevice,cuContext));
+
+      CUstream cuStream = nullptr;
+      g_Streams.push_back(cuStream);
     }
     return;
   }
@@ -90,21 +87,23 @@ class CudaResMgr {
     if (idx >= GetNumGpus()) {
       return nullptr;
     }
+
     lock_guard<mutex> lock(gContextsMutex);
     auto &ctx = g_Contexts[idx];
-    if (!ctx) {
+    if (!ctx.second) {
       CUdevice cuDevice = 0;
       ThrowOnCudaError(cuDeviceGet(&cuDevice, idx), __LINE__);
-      ThrowOnCudaError(cuCtxCreate(&ctx, 0, cuDevice), __LINE__);
+      ThrowOnCudaError(cuDevicePrimaryCtxRetain(&ctx.second, cuDevice), __LINE__);
     }
 
-    return g_Contexts[idx];
+    return g_Contexts[idx].second;
   }
 
   CUstream GetStream(size_t idx) {
     if (idx >= GetNumGpus()) {
       return nullptr;
     }
+
     lock_guard<mutex> lock(gStreamsMutex);
     auto &str = g_Streams[idx];
     if (!str) {
@@ -131,15 +130,15 @@ class CudaResMgr {
         }
         g_Streams.clear();
       }
+
       {
         lock_guard<mutex> lock(gContextsMutex);
-        for (auto &cuContext : g_Contexts) {
-          if (cuContext) {
-            ThrowOnCudaError(cuCtxDestroy(cuContext), __LINE__);
+        for (int i=0;i<g_Contexts.size();i++) {
+          if (g_Contexts[i].second) {
+            ThrowOnCudaError(cuDevicePrimaryCtxRelease(g_Contexts[i].first), __LINE__);
           }
         }
         g_Contexts.clear();
-
       }
     } catch (runtime_error &e) {
       cerr << e.what() << endl;
@@ -154,7 +153,7 @@ class CudaResMgr {
 
   static size_t GetNumGpus() { return Instance().g_Contexts.size(); }
 
-  vector<CUcontext> g_Contexts;
+  vector<pair<CUdevice,CUcontext>> g_Contexts;
   vector<CUstream> g_Streams;
   mutex gContextsMutex;
   mutex gStreamsMutex;