updates of SetDevice

source/tensor/XDevice.cpp

@@ -268,6 +268,28 @@ int XDevice::GetGPUDevice()
 #endif
 }
 
+/* 
+swith to a device (CPU or GPU) 
+>> devID - device id
+*/
+void XDevice::SetDevice(int devID)
+{
+    if(devID >= 0)
+        SetGPUDevice(devID);
+}
+
+/* 
+swith to a device (CPU or GPU) with a backup of the device id 
+>> devID - device id
+>> backupDevID - backup of the device id
+*/
+void XDevice::SetDevice(int devID, int &backupDevID)
+{
+    backupDevID = GetGPUDevice();
+    if (devID >= 0)
+        SetGPUDevice(devID);
+}
+
 /* reset cuda flag for more efficient cuda execution. It should be called after "SetGPUDevice" when
    no GPU context has been established. */
 void XDevice::SetFastFlags()

source/tensor/XDevice.h

@@ -138,7 +138,7 @@ public:
     cublasHandle_t * GetCublasHandle();
 #endif
 
-    /* switch to a device */
+    /* switch to a GPU device */
     static
     void SetGPUDevice(int devID);
 
@@ -146,10 +146,18 @@ public:
     static
     void SetGPUDeviceFast(int devID);
 
-    /* switch to a get current dev */
+    /* get current dev */
     static
     int GetGPUDevice();
 
+    /* swith to a device (CPU or GPU) */
+    static
+    void SetDevice(int devID);
+
+    /* swith to a device (CPU or GPU) with a backup of the device id */
+    static
+    void SetDevice(int devID, int &backupDevID);
+
     /* reset cuda flag for more efficient cuda execution */
     static
     void SetFastFlags();

source/tensor/XQueue.cpp

@@ -170,7 +170,7 @@ void XQueue::RunJobConsumer(int jobDevID)
     isJobQueue = true;
     jobDequeuerArgs->Clear();
 
-    // warning: this may cause unknown error
+    /* warning: this may cause unknown errors */
     jobDequeuerArgs->Add(this);
     jobDequeuerArgs->Add(jobDevID >= 0 ? (devids + jobDevID) : &cpuid);
 
@@ -214,10 +214,8 @@ void XQueue::DequeueJobs(XList * args)
     XQueue * q = (XQueue*)args->GetItem(0);
     int devID = *(int*)args->GetItem(1);
 
-    int devIDBackup = XDevice::GetGPUDevice();
-
-    if(devID >= 0)
-        XDevice::SetGPUDevice(devID);
+    int devIDBackup = -1;
+    XDevice::SetDevice(devID, devIDBackup);
 
     while(1){
         JobQueueNode * node = (JobQueueNode*)q->Dequeue();
@@ -238,8 +236,7 @@ void XQueue::DequeueJobs(XList * args)
 
     }
 
-    if(devID >= 0)
-        XDevice::SetGPUDevice(devIDBackup);
+    XDevice::SetDevice(devIDBackup);
 }
 
 /* get the break flag */

source/tensor/XQueue.h

@@ -129,7 +129,7 @@ public:
     void WaitForEmptyJobQueue();
 
     /* run the job consumer */
-    void RunJobConsumer(int jobDevID = 0);
+    void RunJobConsumer(int jobDevID = -1);
 
     /* stop the job consumer */
     void StopJobConsumer();