From 3afc37ceec3e0589be61f0d316169ddc06df277c Mon Sep 17 00:00:00 2001
From: Alexey Spizhevoy <no@email>
Date: Thu, 27 Jan 2011 13:48:33 +0000
Subject: [PATCH] added HAVE_CUDA, HAVE_TBB handling into multi_gpu sample

---
 samples/gpu/multi_gpu.cpp | 94 +++++++++++++++++++++++++++++++++------
 1 file changed, 81 insertions(+), 13 deletions(-)

diff --git a/samples/gpu/multi_gpu.cpp b/samples/gpu/multi_gpu.cpp
index 262f6e03f..c569875c8 100644
--- a/samples/gpu/multi_gpu.cpp
+++ b/samples/gpu/multi_gpu.cpp
@@ -3,32 +3,100 @@
 
 #include <iostream>
 #include <cvconfig.h>
+#include "opencv2/core/core.hpp"
 #include <opencv2/gpu/gpu.hpp>
-#include <opencv2/highgui/highgui.hpp>
 
-#ifdef HAVE_CUDA
-#include <cuda_runtime.h>
-#endif
-
-using namespace std;
-using namespace cv;
+#if !defined(HAVE_CUDA) || !defined(HAVE_TBB)
 
 int main()
 {
-    bool can_run = true;
-
 #if !defined(HAVE_CUDA)
     cout << "CUDA support is required (CMake key 'WITH_CUDA' must be true).\n";
-    can_run = false;
 #endif
 
 #if !defined(HAVE_TBB)
     cout << "TBB support is required (CMake key 'WITH_TBB' must be true).\n";
-    can_run = false;
 #endif
 
-    if (!can_run)
+    return 0;
+}
+
+#else
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include "opencv2/core/internal.hpp" // For TBB wrappers
+
+using namespace std;
+using namespace cv;
+using namespace cv::gpu;
+
+
+void cuSafeCall(int code);
+struct Worker { void operator()(int device_id) const; };
+void destroy();
+
+
+// Each GPU is associated with its own context
+CUcontext contexts[2];
+
+// Auxiliary variable, stores previusly used context
+CUcontext prev_context;
+
+
+int main()
+{
+    if (getCudaEnabledDeviceCount() < 2)
+    {
+        cout << "Two or more GPUs are required\n";
         return -1;
+    }
+
+    // Save the default context
+    cuSafeCall(cuCtxAttach(&contexts[0], 0));
+    cuSafeCall(cuCtxDetach(contexts[0]));
+
+    // Create new context for the second GPU
+    CUdevice device;
+    cuSafeCall(cuDeviceGet(&device, 1));
+    cuSafeCall(cuCtxCreate(&contexts[1], 0, device));
+
+    // Restore the first GPU context
+    cuSafeCall(cuCtxPopCurrent(&prev_context));
+
+    // Run 
+    int devices[] = {0, 1};
+    parallel_do(devices, devices + 2, Worker());
+
+    // Destroy context of the second GPU
+    destroy();
 
     return 0;
-}
\ No newline at end of file
+}
+
+
+void Worker::operator()(int device_id) const
+{
+    cout << device_id << endl;
+}
+
+
+void cuSafeCall(int code)
+{
+    if (code != CUDA_SUCCESS) 
+    {
+        cout << "CUDA driver API error: code " << code 
+            << ", file " << __FILE__ 
+            << ", line " << __LINE__ << endl;
+        destroy();
+        exit(-1);
+    }
+}
+
+
+void destroy() 
+{
+    cuCtxDestroy(contexts[1]);
+}
+
+#endif
\ No newline at end of file