Add astra.astra.get_gpu_info utility function

4 files changed, 47 insertions, 0 deletions

diff --git a/cuda/2d/util.cu b/cuda/2d/util.cu
index 09d1a2b..9c1bb28 100644
--- a/cuda/2d/util.cu
+++ b/cuda/2d/util.cu
@@ -274,5 +274,29 @@ void reportCudaError(cudaError_t err)
 }
 
 
+_AstraExport std::string getCudaDeviceString(int device)
+{
+	char buf[1024];
+	cudaError_t err;
+	if (device == -1) {
+		err = cudaGetDevice(&device);
+		if (err != cudaSuccess) {
+			return "Error getting current GPU index";
+		}
+	}
+
+	cudaDeviceProp prop;
+	err = cudaGetDeviceProperties(&prop, device);
+	if (err != cudaSuccess) {
+		snprintf(buf, 1024, "GPU #%d: Invalid device (%d): %s", device, err, cudaGetErrorString(err));
+		return buf;
+	}
+
+	long mem = prop.totalGlobalMem / (1024*1024);
+	snprintf(buf, 1024, "GPU #%d: %s, with %ldMB", device, prop.name, mem);
+	return buf;
+}
+
+
 
 }
diff --git a/cuda/2d/util.h b/cuda/2d/util.h
index 875aae3..31fcfbd 100644
--- a/cuda/2d/util.h
+++ b/cuda/2d/util.h
@@ -30,6 +30,7 @@ along with the ASTRA Toolbox. If not, see <http://www.gnu.org/licenses/>.
 
 #include <cuda.h>
 #include <driver_types.h>
+#include <string>
 
 #ifdef _MSC_VER
 
@@ -92,6 +93,9 @@ void reportCudaError(cudaError_t err);
 float dotProduct2D(float* D_data, unsigned int pitch,
                    unsigned int width, unsigned int height);
 
+// Return string with CUDA device number, name and memory size.
+// Use device == -1 to get info for the current device.
+_AstraExport std::string getCudaDeviceString(int device);
 
 }
 
diff --git a/python/astra/astra.py b/python/astra/astra.py
index 3804d51..434ccb9 100644
--- a/python/astra/astra.py
+++ b/python/astra/astra.py
@@ -45,6 +45,16 @@ def set_gpu_index(idx, memory=0):
     """
     a.set_gpu_index(idx, memory)
 
+def get_gpu_info(idx=-1):
+    """Get GPU info.
+    
+    :param idx: GPU index, or -1 for current device
+    :type idx: :class:`int`
+    :returns: :class:`str` -- GPU info
+    """
+    return a.get_gpu_info(idx)
+
+
 def delete(ids):
     """Delete an astra object.
     
diff --git a/python/astra/astra_c.pyx b/python/astra/astra_c.pyx
index 6de10da..f39b0a1 100644
--- a/python/astra/astra_c.pyx
+++ b/python/astra/astra_c.pyx
@@ -42,9 +42,14 @@ cdef extern from "astra/Globals.h" namespace "astra":
 IF HAVE_CUDA==True:
   cdef extern from "../cuda/2d/darthelper.h" namespace "astraCUDA":
       bool setGPUIndex(int)
+  cdef extern from "../cuda/2d/util.h" namespace "astraCUDA":
+      string getCudaDeviceString(int)
 ELSE:
   def setGPUIndex():
     pass
+  def getCudaDeviceString(idx):
+    pass
+
 cdef extern from "astra/CompositeGeometryManager.h" namespace "astra":
     cdef cppclass SGPUParams:
         vector[int] GPUIndices
@@ -85,9 +90,13 @@ IF HAVE_CUDA==True:
         ret = setGPUIndex(params.GPUIndices[0])
         if not ret:
             six.print_("Failed to set GPU " + str(params.GPUIndices[0]))
+  def get_gpu_info(idx=-1):
+    return wrap_from_bytes(getCudaDeviceString(idx))
 ELSE:
   def set_gpu_index(idx, memory=0):
     raise NotImplementedError("CUDA support is not enabled in ASTRA")
+  def get_gpu_info(idx=-1):
+    raise NotImplementedError("CUDA support is not enabled in ASTRA")
 
 def delete(ids):
     import collections