From ce5c6a3ad19fd94ed88e4635824bbfcd0f9262ab Mon Sep 17 00:00:00 2001
From: Wen Jufa <jufa.wen@sophgo.com>
Date: Mon, 1 Aug 2022 14:38:38 +0800
Subject: [PATCH] Add memory hint & calculate times

Change-Id: I0dbcc6ea3110e2a9bab2716409aa01cd56ddbf12
---
 python/tpu_perf/build.py |  4 +++-
 python/tpu_perf/run.py   | 12 ++++++++++--
 python/tpu_perf/subp.py  |  6 ++++--
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/python/tpu_perf/build.py b/python/tpu_perf/build.py
index 8e87f8a..629b3a9 100644
--- a/python/tpu_perf/build.py
+++ b/python/tpu_perf/build.py
@@ -62,7 +62,9 @@ def build(tree, path, config):
         fp32_pool.wait()
         logging.info(f'FP32 bmodel {name} done.')
 
-    int8_pool = CommandExecutor(workdir, env)
+    int8_pool = CommandExecutor(
+        workdir, env,
+        memory_hint=config.get('memory_hint'))
 
     cali_key = 'time_only_cali' if option_time_only else 'cali'
 
diff --git a/python/tpu_perf/run.py b/python/tpu_perf/run.py
index b968b3c..aff294f 100644
--- a/python/tpu_perf/run.py
+++ b/python/tpu_perf/run.py
@@ -68,6 +68,9 @@ def run(tree, path, config, stat_f):
     pool = CommandExecutor(workdir, env)
     rounds = config.get('time_rounds', 2000)
     rt_cmp = config.get('runtime_cmp')
+    iter_opt = tree.global_config.get('iter_opt', '--loopnum')
+    if 'iter_opt' in config:
+        iter_opt = config['iter_opt']
     for b in config['bmnetu_batch_sizes']:
         name = f'{b}b.compilation'
         bmodel_dir = os.path.join(workdir, name)
@@ -84,6 +87,9 @@ def run(tree, path, config, stat_f):
             info = parse_profile(profile_path)
             rounds = int(1200 / info['runtime'])
             logging.info(f'Run {rounds} times for {name}.{b}b')
+        max_rounds = 500000
+        if rounds > max_rounds:
+            rounds = max_rounds
 
         title = f'run.{b}'
         ref_fn = os.path.join(bmodel_dir, 'output_ref_data.dat')
@@ -92,13 +98,13 @@ def run(tree, path, config, stat_f):
             logging.info(f'Runtime test {name}')
             pool.put(
                 title,
-                ['bmrt_test', '--loopnum', str(rounds), '--context', bmodel_dir],
+                ['bmrt_test', iter_opt, str(rounds), '--context', bmodel_dir],
                 env=env, shell=False)
         else:
             logging.info(f'Runtime test {name} without reference')
             pool.put(
                 title,
-                ['bmrt_test', '--loopnum', str(rounds), '--bmodel', bmodel],
+                ['bmrt_test', iter_opt, str(rounds), '--bmodel', bmodel],
                 env=env, shell=False)
         try:
             pool.fire()
@@ -115,6 +121,8 @@ def run(tree, path, config, stat_f):
             stats = parse_stats(f.read())
         from math import nan
         real_time = stats['calculate'] * 1000 if 'calculate' in stats else nan
+        if 'calculate_times' in iter_opt:
+            real_time /= rounds
         row = [
             config['name'],
             stats['shape'],
diff --git a/python/tpu_perf/subp.py b/python/tpu_perf/subp.py
index f69d19b..8842f2f 100644
--- a/python/tpu_perf/subp.py
+++ b/python/tpu_perf/subp.py
@@ -19,14 +19,16 @@ def sys_memory_size():
        return int(m.group(1))
 
 class CommandExecutor:
-    def __init__(self, cwd, env):
+    def __init__(self, cwd, env, memory_hint = None):
+        if memory_hint is None:
+            memory_hint = 1024 * 1024 * 7
         import os
         self.env = os.environ.copy()
         for v in env:
             pair = v.split('=')
             self.env[pair[0].strip()] = pair[1].strip() if len(pair) > 1 else ""
         mem_size = sys_memory_size()
-        max_threads = max(1, int(mem_size / 1024 / 1024 / 4))
+        max_threads = max(1, int(mem_size / memory_hint))
         self.threads = 4
         if self.threads > max_threads:
             self.threads = max_threads