[MXNET-1430] julia: implement context.gpu_memory_info (apache#16324)

* julia: implement context.gpu_memory_info resolve MXNET-1430 * update export and NEWS
aaronmarkham · Oct 16, 2019 · 4547cc4 · 4547cc4
1 parent d948256
commit 4547cc4
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 3 deletions.
diff --git a/julia/NEWS.md b/julia/NEWS.md
@@ -20,6 +20,10 @@
 * Add an abstract type `AbstractMXError` as the parent type for all MXNet-related
   API errors. (#16235)
 
+* Porting more `context` functions from Python.
+    * `num_gpus()` (#16236)
+    * `gpu_memory_info()` (#16324)
+
 
 # v1.5.0
 

diff --git a/julia/src/MXNet.jl b/julia/src/MXNet.jl
@@ -79,7 +79,8 @@ export Executor,
 export Context,
        cpu,
        gpu,
-       num_gpus
+       num_gpus,
+       gpu_memory_info
 
 # model.jl
 export AbstractModel,

diff --git a/julia/src/context.jl b/julia/src/context.jl
@@ -68,3 +68,21 @@ function num_gpus()
   @mxcall :MXGetGPUCount (Ref{Cint},) n
   n[]
 end
+
+"""
+    gpu_memory_info(dev_id = 0)::Tuple{UInt64,UInt64}
+
+Query CUDA for the free and total bytes of GPU global memory.
+It returns a tuple of `(free memory, total memory)`.
+
+```julia-repl
+julia> mx.gpu_memory_info()
+(0x00000003af240000, 0x00000003f9440000)
+```
+"""
+function gpu_memory_info(dev_id = 0)
+  free = Ref{UInt64}()
+  n = Ref{UInt64}()
+  @mxcall :MXGetGPUMemoryInformation64 (Cint, Ref{UInt64}, Ref{UInt64}) dev_id free n
+  free[], n[]
+end
diff --git a/python/mxnet/context.py b/python/mxnet/context.py
@@ -292,8 +292,6 @@ def gpu_memory_info(device_id=0):
     Returns
     -------
     (free, total) : (int, int)
-        The number of GPUs.
-
     """
     free = ctypes.c_uint64()
     total = ctypes.c_uint64()