Add astra.astra.get_gpu_info utility function

2 files changed, 19 insertions, 0 deletions

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