[incubator-mxnet] 01/01: julia: implement context.gpu_memory_info

[ib...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

iblis pushed a commit to branch ib/jl-gpu-mem
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git

commit a20dcd5c47f8a179fb84b7b36943efe2b675593f
Author: Iblis Lin <ib...@hs.ntnu.edu.tw>
AuthorDate: Sat Sep 28 05:49:23 2019 +0000

    julia: implement context.gpu_memory_info
    
    resolve MXNET-1430
---
 julia/src/context.jl    | 18 ++++++++++++++++++
 python/mxnet/context.py |  2 --
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/julia/src/context.jl b/julia/src/context.jl
index bce67a5..68e6913 100644
--- 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
index decea71..99998ea 100644
--- 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()