File tree Expand file tree Collapse file tree 4 files changed +24
-3
lines changed Expand file tree Collapse file tree 4 files changed +24
-3
lines changed Original file line number Diff line number Diff line change 20
20
* Add an abstract type ` AbstractMXError ` as the parent type for all MXNet-related
21
21
API errors. (#16235 )
22
22
23
+ * Porting more ` context ` functions from Python.
24
+ * ` num_gpus() ` (#16236 )
25
+ * ` gpu_memory_info() ` (#16324 )
26
+
23
27
24
28
# v1.5.0
25
29
Original file line number Diff line number Diff line change @@ -79,7 +79,8 @@ export Executor,
79
79
export Context,
80
80
cpu,
81
81
gpu,
82
- num_gpus
82
+ num_gpus,
83
+ gpu_memory_info
83
84
84
85
# model.jl
85
86
export AbstractModel,
Original file line number Diff line number Diff line change @@ -68,3 +68,21 @@ function num_gpus()
68
68
@mxcall :MXGetGPUCount (Ref{Cint},) n
69
69
n[]
70
70
end
71
+
72
+ """
73
+ gpu_memory_info(dev_id = 0)::Tuple{UInt64,UInt64}
74
+
75
+ Query CUDA for the free and total bytes of GPU global memory.
76
+ It returns a tuple of `(free memory, total memory)`.
77
+
78
+ ```julia-repl
79
+ julia> mx.gpu_memory_info()
80
+ (0x00000003af240000, 0x00000003f9440000)
81
+ ```
82
+ """
83
+ function gpu_memory_info (dev_id = 0 )
84
+ free = Ref {UInt64} ()
85
+ n = Ref {UInt64} ()
86
+ @mxcall :MXGetGPUMemoryInformation64 (Cint, Ref{UInt64}, Ref{UInt64}) dev_id free n
87
+ free[], n[]
88
+ end
Original file line number Diff line number Diff line change @@ -292,8 +292,6 @@ def gpu_memory_info(device_id=0):
292
292
Returns
293
293
-------
294
294
(free, total) : (int, int)
295
- The number of GPUs.
296
-
297
295
"""
298
296
free = ctypes .c_uint64 ()
299
297
total = ctypes .c_uint64 ()
You can’t perform that action at this time.
0 commit comments