From b3d2a632f3d44cf0fc1a266cf988f87e786f9c58 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 30 Jul 2021 18:47:56 +0800
Subject: [PATCH] fix(sublinear): fix lb_memory in sublinear

GitOrigin-RevId: fd51023c01c179cf88bb599e92fa58565781580d
---
 imperative/python/megengine/jit/sublinear_memory_config.py | 6 +++---
 imperative/python/megengine/jit/tracing.py                 | 2 +-
 imperative/python/src/graph_rt.cpp                         | 2 +-
 src/core/impl/graph/seq_sublinear_memory.cpp               | 4 ++--
 src/core/include/megbrain/graph/cg.h                       | 2 +-
 5 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/imperative/python/megengine/jit/sublinear_memory_config.py b/imperative/python/megengine/jit/sublinear_memory_config.py
index a84fac4ef..a2ce31008 100644
--- a/imperative/python/megengine/jit/sublinear_memory_config.py
+++ b/imperative/python/megengine/jit/sublinear_memory_config.py
@@ -22,7 +22,7 @@ class SublinearMemoryConfig:
     :param genetic_pool_size: number of samples for the crossover random selection
         during genetic optimization. Default: 20.
         It can also be set through the environmental variable 'MGB_SUBLINEAR_MEMORY_GENETIC_POOL_SIZE'.
-    :param lb_memory: memory lower bound of bottleneck size in MB for sublinear memory optimization.
+    :param lb_memory_mb: memory lower bound of bottleneck size in MB for sublinear memory optimization.
         It can be used to perform manual tradeoff between memory and speed. Default: 0.
         It can also be set through the environmental variable 'MGB_SUBLINEAR_MEMORY_LOWER_BOUND_MB'.
     :param num_worker: number of thread workers to search the optimum checkpoints
@@ -39,7 +39,7 @@ class SublinearMemoryConfig:
         thresh_nr_try: int = 10,
         genetic_nr_iter: int = 0,
         genetic_pool_size: int = 20,
-        lb_memory: int = 0,
+        lb_memory_mb: int = 0,
         num_worker: int = max(1, get_device_count("cpu") // 2),
     ):
         assert thresh_nr_try >= 0, "thresh_nr_try must be greater or equal to zero"
@@ -50,6 +50,6 @@ class SublinearMemoryConfig:
             genetic_pool_size >= 0
         ), "genetic_pool_size must be greater or equal to zero"
         self.genetic_pool_size = genetic_pool_size
-        self.lb_memory = lb_memory
+        self.lb_memory_mb = lb_memory_mb
         assert num_worker > 0, "num_worker must be greater or equal to one"
         self.num_worker = num_worker
diff --git a/imperative/python/megengine/jit/tracing.py b/imperative/python/megengine/jit/tracing.py
index 8946b88bd..53bd63387 100644
--- a/imperative/python/megengine/jit/tracing.py
+++ b/imperative/python/megengine/jit/tracing.py
@@ -532,7 +532,7 @@ class trace:
         if self._sublinear_memory_config is not None:
             graph.options.enable_sublinear_memory_opt = True
             sublinear_config = graph.options.sublinear_mem_config
-            sublinear_config.lb_memory = self._sublinear_memory_config.lb_memory
+            sublinear_config.lb_memory_mb = self._sublinear_memory_config.lb_memory_mb
             sublinear_config.genetic_nr_iter = (
                 self._sublinear_memory_config.genetic_nr_iter
             )
diff --git a/imperative/python/src/graph_rt.cpp b/imperative/python/src/graph_rt.cpp
index ab5ba647a..f17d50464 100644
--- a/imperative/python/src/graph_rt.cpp
+++ b/imperative/python/src/graph_rt.cpp
@@ -476,7 +476,7 @@ void init_graph_rt(py::module m) {
         DEF_READWRITE(thresh_nr_try)
         DEF_READWRITE(genetic_nr_iter)
         DEF_READWRITE(genetic_pool_size)
-        DEF_READWRITE(lb_memory)
+        DEF_READWRITE(lb_memory_mb)
         DEF_READWRITE(num_worker);
 
 #undef CURRENT_CLASS
diff --git a/src/core/impl/graph/seq_sublinear_memory.cpp b/src/core/impl/graph/seq_sublinear_memory.cpp
index eaac0aa81..a4075a013 100644
--- a/src/core/impl/graph/seq_sublinear_memory.cpp
+++ b/src/core/impl/graph/seq_sublinear_memory.cpp
@@ -562,7 +562,7 @@ public:
             m_config->genetic_pool_size = psize;
         }
         if (auto env = MGB_GETENV("MGB_SUBLINEAR_MEMORY_LOWER_BOUND_MB")) {
-            m_config->lb_memory = std::stoi(env) * 1024 * 1024;
+            m_config->lb_memory_mb = std::stoi(env);
         }
     }
 
@@ -781,7 +781,7 @@ void SeqModifierForSublinearMemory::ActionSearcherSingleCN::search_genetic() {
 }
 
 void SeqModifierForSublinearMemory::ActionSearcherSingleCN::search_refine() {
-    size_t lower_bound = m_par_modifier->m_config->lb_memory;
+    size_t lower_bound = static_cast<size_t>(m_par_modifier->m_config->lb_memory_mb) << 20;
     if (m_min_bottleneck >= lower_bound)
         return;
     OprFootprint footprint;
diff --git a/src/core/include/megbrain/graph/cg.h b/src/core/include/megbrain/graph/cg.h
index 8ca976e46..bbca52052 100644
--- a/src/core/include/megbrain/graph/cg.h
+++ b/src/core/include/megbrain/graph/cg.h
@@ -443,7 +443,7 @@ class ComputingGraph : public std::enable_shared_from_this<ComputingGraph>,
                 int thresh_nr_try = 10;
                 int genetic_nr_iter = 0;
                 int genetic_pool_size = 20;
-                int lb_memory = 0;
+                int lb_memory_mb = 0;
                 int num_worker = sys::get_cpu_count() / 2;
             } sublinear_mem_config;
 
-- 
GitLab