From 182ca25d45982d72f121baef6d023baf4c57ddca Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 15 Oct 2021 14:20:32 +0800
Subject: [PATCH] feat(mge/imperative): add coalesce_all_free_memory mge api

GitOrigin-RevId: ea05b1e9a17df17d853b8d7ff788ed0321a668ae
---
 imperative/python/megengine/device.py | 17 +++++++++++++++++
 imperative/python/src/utils.cpp       |  4 ++++
 2 files changed, 21 insertions(+)

diff --git a/imperative/python/megengine/device.py b/imperative/python/megengine/device.py
index dbd54fd21..374744a31 100644
--- a/imperative/python/megengine/device.py
+++ b/imperative/python/megengine/device.py
@@ -16,6 +16,7 @@ from .core._imperative_rt.common import (
 )
 from .core._imperative_rt.common import set_prealloc_config as _set_prealloc_config
 from .core._imperative_rt.common import what_is_xpu as _what_is_xpu
+from .core._imperative_rt.utils import _try_coalesce_all_free_memory
 
 __all__ = [
     "is_cuda_available",
@@ -25,6 +26,7 @@ __all__ = [
     "get_mem_status_bytes",
     "get_cuda_compute_capability",
     "set_prealloc_config",
+    "coalesce_free_memory",
     "DeviceType",
 ]
 
@@ -186,3 +188,18 @@ def set_prealloc_config(
 
 def what_is_xpu():
     return _what_is_xpu().name.lower()
+
+
+def coalesce_free_memory():
+    r"""This function will try it best to free all consecutive free chunks back to operating system,
+    small pieces may not be returned.
+    
+    because of the async processing of megengine, the effect of this func may not be reflected 
+    immediately. if you want to see the effect immediately, you can call megengine._full_sync after
+    this func was called
+
+    .. note::
+    Please notice that this function will not move any memory in-use.
+    Please notice that this function may do nothing if there are no chunks that can be freed
+    """
+    return _try_coalesce_all_free_memory()
diff --git a/imperative/python/src/utils.cpp b/imperative/python/src/utils.cpp
index fcc05bf0e..5ea8cf51a 100644
--- a/imperative/python/src/utils.cpp
+++ b/imperative/python/src/utils.cpp
@@ -194,6 +194,10 @@ void init_utils(py::module m) {
     m.def("_get_device_count", &mgb::CompNode::get_device_count,
           "Get total number of specific devices on this system");
 
+    m.def("_try_coalesce_all_free_memory", &mgb::CompNode::try_coalesce_all_free_memory,
+          "This function will try it best to free all consecutive free chunks back to "
+          "operating system");
+
     using mgb::imperative::TensorSanityCheck;
     py::class_<TensorSanityCheck>(m, "TensorSanityCheckImpl")
             .def(py::init<>())
-- 
GitLab