Adding set device function

common/ScaLBL.h

@@ -11,6 +11,8 @@
 #define ScalLBL_H
 #include "common/Domain.h"
 
+extern "C" int ScaLBL_SetDevice();
+
 extern "C" void ScaLBL_AllocateDeviceMemory(void** address, size_t size);
 
 extern "C" void ScaLBL_FreeDeviceMemory(void* pointer);

gpu/Extras.cu

@@ -2,6 +2,15 @@
 #include <cuda.h>
 #include <stdio.h>
 
+extern "C" int ScaLBL_SetDevice(){
+	int n_devices; 
+	int local_rank = atoi(getenv("MV2_COMM_WORLD_LOCAL_RANK"));
+	cudaGetDeviceCount(&n_devices); 
+	int device = local_rank % n_devices; 
+	cudaSetDevice(device); 
+	return device;
+}
+
 extern "C" void ScaLBL_AllocateDeviceMemory(void** address, size_t size){
        	cudaMalloc(address,size);
 	cudaError_t err = cudaGetLastError();

tests/lbpm_color_macro_simulator.cpp

@@ -27,7 +27,10 @@ using namespace std;
 
 int main(int argc, char **argv)
 {
-  // Initialize MPI
+  // Initialize compute device
+	int device=ScaLBL_SetDevice();
+	
+	// Initialize MPI
 	int provided_thread_support = -1;
 	MPI_Init_thread(&argc,&argv,MPI_THREAD_MULTIPLE,&provided_thread_support);
 	MPI_Comm comm;