Expose the number of GPUs.

include/mxnet/base.h

@@ -226,6 +226,11 @@ struct Context {
    * \return GPU Context. -1 for current GPU.
    */
   inline static Context GPU(int32_t dev_id = -1);
+  /*!
+   * Get the number of GPUs available.
+   * \return The number of GPUs that are available.
+   */
+  inline static int32_t GetGPUCount();
   /*!
    * Create a pinned CPU context.
    * \param dev_id the device id for corresponding GPU.
@@ -316,6 +321,17 @@ inline Context Context::GPU(int32_t dev_id) {
   return Create(kGPU, dev_id);
 }
 
+inline int32_t Context::GetGPUCount() {
+#if MXNET_USE_CUDA
+  int32_t count;
+  cudaError_t e = cudaGetDeviceCount(&count);
+  CHECK_EQ(e, cudaSuccess) << " CUDA: " << cudaGetErrorString(e);
+  return count;
+#else
+  LOG(FATAL) << "Please compile with CUDA support to query the number of GPUs.";
+#endif
+}
+
 inline Context Context::FromString(const std::string& str) {
   Context ret;
   try {

include/mxnet/c_api.h

@@ -367,6 +367,13 @@ MXNET_DLL int MXSetNumOMPThreads(int thread_num);
  */
 MXNET_DLL int MXEngineSetBulkSize(int bulk_size, int* prev_bulk_size);
 
+/*!
+ * \brief Get the number of GPUs.
+ * \param pointer to int that will hold the number of GPUs available.
+ * \return 0 when success, -1 when failure happens.
+ */
+MXNET_DLL int MXGetGPUCount(int* out);
+
 /*!
  * \brief get the MXNet library version as an integer
  * \param pointer to the integer holding the version number

python/mxnet/context.py

@@ -19,6 +19,10 @@
 """Context management API of mxnet."""
 from __future__ import absolute_import
 
+import ctypes
+from .base import _LIB
+from .base import check_call
+
 class Context(object):
     """Constructs a context.
 
@@ -212,6 +216,25 @@ def gpu(device_id=0):
     return Context('gpu', device_id)
 
 
+def num_gpus():
+    """Query CUDA for the number of GPUs present.
+
+    Raises
+    ------
+    Will raise an exception on any CUDA error or in case MXNet was not
+    compiled with CUDA support.
+
+
+    Returns
+    -------
+    count : int
+        The number of GPUs.
+ 
+    """
+    count = ctypes.c_int()
+    check_call(_LIB.MXGetGPUCount(ctypes.byref(count)))
+    return count.value
+
 def current_context():
     """Returns the current context.
 

src/c_api/c_api.cc

@@ -108,6 +108,12 @@ int MXEngineSetBulkSize(int bulk_size, int* prev_bulk_size) {
   API_END();
 }
 
+int MXGetGPUCount(int* out) {
+  API_BEGIN();
+  *out = Context::GetGPUCount();
+  API_END();
+}
+
 int MXGetVersion(int *out) {
   API_BEGIN();
   *out = static_cast<int>(MXNET_VERSION);