Added support for multi-gpu. see rgl-epfl#10

Steelwall2014 · Mar 2, 2024 · d8e8c9c · d8e8c9c
1 parent acb3342
commit d8e8c9c
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/src/cuda_driver.cpp b/src/cuda_driver.cpp
@@ -55,7 +55,7 @@ void cuda_check_impl(CUresult errval, const char *file, const int line) {
     }
 }
 
-bool init_cuda() {
+bool init_cuda(int device_id) {
 
     if (handle)
         return true;
@@ -121,7 +121,7 @@ bool init_cuda() {
     }
 
     cuda_check(cuInit(0));
-    cuda_check(cuDeviceGet(&cu_device, 0));
+    cuda_check(cuDeviceGet(&cu_device, device_id));
     cuda_check(cuDevicePrimaryCtxRetain(&cu_context, cu_device));
     cuda_check(cuCtxPushCurrent(cu_context));
     cuda_check(cuModuleLoadData(&cu_module, (void *) imageBytes));

diff --git a/src/cuda_driver.h b/src/cuda_driver.h
@@ -47,7 +47,7 @@ extern CUfunction solve_lower_double;
 extern CUfunction analysis_lower;
 extern CUfunction analysis_upper;
 
-extern bool init_cuda();
+extern bool init_cuda(int device_id);
 extern void shutdown_cuda();
 extern void cuda_check_impl(CUresult errval, const char *file, const int line);
 

diff --git a/src/main.cpp b/src/main.cpp
@@ -48,7 +48,7 @@ void declare_cholesky(nb::module_ &m, const std::string &typestr, const char *do
                 // GPU init
 
                 // Initialize CUDA and load the kernels if not already done
-                init_cuda();
+                init_cuda(ii.device_id());
 
                 scoped_set_context guard(cu_context);