From ab82c8da7943669c2b4c2ba8d7464f3b5aa4a1d3 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Wed, 18 Nov 2020 19:31:49 +0800
Subject: [PATCH] feat(mge/device): add python binding for get_mem_status_bytes

GitOrigin-RevId: aa0d3264d6042b1f3c7dfe243ef9a97794868c42
---
 imperative/python/megengine/device.py | 12 ++++++++++++
 imperative/python/src/common.cpp      |  3 +++
 2 files changed, 15 insertions(+)

diff --git a/imperative/python/megengine/device.py b/imperative/python/megengine/device.py
index e7a826d4..36845d69 100644
--- a/imperative/python/megengine/device.py
+++ b/imperative/python/megengine/device.py
@@ -8,6 +8,7 @@
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 import os
 import re
+from typing import Optional
 
 from .core._imperative_rt.common import CompNode, DeviceType
 from .core._imperative_rt.common import set_prealloc_config as _set_prealloc_config
@@ -17,6 +18,7 @@ __all__ = [
     "get_device_count",
     "get_default_device",
     "set_default_device",
+    "get_mem_status_bytes",
     "set_prealloc_config",
     "DeviceType",
 ]
@@ -92,6 +94,16 @@ def get_default_device() -> str:
     return CompNode._get_default_device()
 
 
+def get_mem_status_bytes(device: Optional[str] = None):
+    r"""
+    Get total and free memory on the computing device in bytes.
+    """
+    if device is None:
+        device = get_default_device()
+    tot, free = CompNode(device).get_mem_status_bytes
+    return tot, free
+
+
 set_default_device(os.getenv("MGE_DEFAULT_DEVICE", "xpux"))
 
 
diff --git a/imperative/python/src/common.cpp b/imperative/python/src/common.cpp
index 136d6dde..fd5c6602 100644
--- a/imperative/python/src/common.cpp
+++ b/imperative/python/src/common.cpp
@@ -58,6 +58,9 @@ void init_common(py::module m) {
         .def_property_readonly("logical_name", [](const CompNode& cn) {
             return cn.to_string_logical();
         })
+        .def_property_readonly("get_mem_status_bytes", [](const CompNode& cn) {
+            return cn.get_mem_status_bytes();
+        })
         .def("create_event", &CompNode::create_event, py::arg("flags") = 0ul)
         .def("_set_default_device", &set_default_device)
         .def("_get_default_device", &get_default_device)