diff options
| -rw-r--r-- | cuda/2d/util.cu | 24 | ||||
| -rw-r--r-- | cuda/2d/util.h | 4 | ||||
| -rw-r--r-- | python/astra/astra.py | 10 | ||||
| -rw-r--r-- | python/astra/astra_c.pyx | 9 |
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 |