211 lines
7.5 KiB
Python
211 lines
7.5 KiB
Python
import sys
|
|
import os
|
|
import logging
|
|
import shutil
|
|
import subprocess
|
|
from utils import util
|
|
|
|
|
|
class OptModules(object):
|
|
|
|
def __init__(self, platform, model, base_config, opt_dir, log):
|
|
self._platform = platform
|
|
self._model = os.path.abspath(model)
|
|
self._base_config = os.path.abspath(base_config)
|
|
self._modules = self._get_modules()
|
|
self._opt_dir = opt_dir
|
|
self._log = log
|
|
self._logger = log.logger
|
|
self._debug_mode = '-d' if (log.logger.level == logging.DEBUG) else ''
|
|
|
|
# check if file exist
|
|
arr = [self._model, self._base_config]
|
|
for path in arr:
|
|
if not util.check_file_exist(path):
|
|
self._logger.error('"{}" not exist!'.format(path))
|
|
sys.exit(-1)
|
|
|
|
@property
|
|
def modules(self):
|
|
return self._modules
|
|
|
|
@property
|
|
def logger(self):
|
|
return self._logger
|
|
|
|
def _get_modules(self):
|
|
# mandatory config for sub-module
|
|
m = {
|
|
"520": [],
|
|
"720": ["image_cut_search"],
|
|
"530": ["image_cut_search"],
|
|
"730": ["image_cut_search"],
|
|
"630": ["image_cut_search"],
|
|
"540": ["image_cut_search"],
|
|
}
|
|
return m[self._platform]
|
|
|
|
def _get_module_command(self, module):
|
|
# mandatory config for sub-module
|
|
m = {
|
|
"example1": "main.py",
|
|
"example2": "main.py",
|
|
"image_cut_search": "all_rule_auto_release.py -r -u -t"
|
|
}
|
|
return '{}/modules/{}/{}'.format(self._opt_dir, module, m[module])
|
|
|
|
def _get_module_environ_var(self, module):
|
|
# optional config for sub-module
|
|
m = {
|
|
"image_cut_search": {
|
|
# speed up compilation time (mainly reduce weight gen time) for evaluation only mode
|
|
# this is useful when running IP evaluator only (e.g. running opt compiler wrapper)
|
|
"KNERON_EVAL_ONLY": "1"
|
|
}
|
|
}
|
|
return m.get(module, {})
|
|
|
|
def _get_module_overwrite_key(self, module):
|
|
# mandatory config for sub-module
|
|
m = {
|
|
"example1": {
|
|
"fmap_cut": {
|
|
"optimize": True
|
|
}
|
|
},
|
|
"example2": {
|
|
"weight_compress": True
|
|
},
|
|
"image_cut_search": {
|
|
"fmap_cut": {
|
|
"mode": "default",
|
|
"cut_evenly": True,
|
|
"optimize": True,
|
|
"cut_piece_search": True
|
|
}
|
|
}
|
|
}
|
|
return m[module]
|
|
|
|
def _get_module_ret_code(self, ret_code):
|
|
if ret_code == 0:
|
|
return 'Success'
|
|
elif ret_code in range(1, 30):
|
|
return 'Compile failed'
|
|
elif ret_code in range(31, 50):
|
|
# 2-50: image-cut-search failed
|
|
msg = '[image_cut_search failed] '
|
|
if ret_code == 32:
|
|
return msg + 'Image cut search only support version 720/530/630/730/540 right now'
|
|
elif ret_code == 33:
|
|
return msg + 'Cannot find valid cut group since no info_cutting.log'
|
|
else:
|
|
return msg + 'Unknown error code'
|
|
else:
|
|
return 'Unknown error code'
|
|
|
|
def run(self):
|
|
for i, module in enumerate(self._modules):
|
|
if i == 0:
|
|
config = self._base_config
|
|
|
|
# create sub-module output directory
|
|
self._log.create_module_dir(module)
|
|
|
|
# set sub-module environment variable
|
|
self.set_module_environ_var(module)
|
|
|
|
# merge sub-module patch config
|
|
config = self.merge_module_compiler_config(module, config)
|
|
|
|
# run sub-module
|
|
config = self.run_one_module(module, config)
|
|
|
|
# unset sub-module environment variable
|
|
self.unset_module_environ_var(module)
|
|
|
|
# copy results in last module to top directory
|
|
if (len(self._modules) > 0):
|
|
self.copy_module_results_to_top_dir(module)
|
|
|
|
def run_one_module(self, module, config):
|
|
# step 1: run optimized module
|
|
c = self._get_module_command(module)
|
|
command = '{command} {platform} {model} {config} {output_dir} {debug_mode}'.format(
|
|
command=c,
|
|
platform=self._platform,
|
|
model=self._model,
|
|
config=config,
|
|
output_dir=self._log.get_module_dir(module),
|
|
debug_mode=self._debug_mode)
|
|
|
|
stdout = subprocess.PIPE
|
|
stderr = subprocess.PIPE
|
|
|
|
self._logger.info('sub-module : "{}"...'.format(module))
|
|
self._logger.info('environment var : {}'.format(self._get_module_environ_var(module)))
|
|
self._logger.info('script : {}'.format(command))
|
|
cmd_args = command.split()
|
|
o = subprocess.run(cmd_args, cwd=self._opt_dir, stdout=stdout, stderr=stderr)
|
|
|
|
if o.stdout:
|
|
self._logger.info('\n\n' + o.stdout.decode())
|
|
|
|
if o.stderr:
|
|
self._logger.error('\n\n' + o.stderr.decode())
|
|
|
|
if o.returncode != 0:
|
|
self._logger.error('run sub-module "{}" failed !!! [ret_code={}. msg="{}"]'.format(
|
|
module, o.returncode, self._get_module_ret_code(o.returncode)))
|
|
sys.exit(o.returncode)
|
|
|
|
# step 2: check if best_config.json and best_profile_result.txt in specific path
|
|
best_config = self._log.get_module_best_config(module)
|
|
best_profile_result = self._log.get_module_best_profile_result(module)
|
|
arr = [best_config, best_profile_result]
|
|
for path in arr:
|
|
if not util.check_file_exist(path):
|
|
self._logger.error('module[{}] should output specific file to "{}"'.format(module, path))
|
|
sys.exit(-1)
|
|
|
|
self._logger.info('... successful.')
|
|
return best_config
|
|
|
|
def set_module_environ_var(self, module):
|
|
envs = self._get_module_environ_var(module)
|
|
for key, value in envs.items():
|
|
os.environ[key] = value
|
|
|
|
def unset_module_environ_var(self, module):
|
|
envs = self._get_module_environ_var(module)
|
|
for key in envs:
|
|
del os.environ[key]
|
|
|
|
def merge_module_compiler_config(self, module, config):
|
|
# step 1: copy config to base_config.json
|
|
base_config = self._log.get_module_base_config(module)
|
|
shutil.copyfile(config, base_config)
|
|
|
|
# step 2: merge module patched config
|
|
cfg = util.load_json_file(config)
|
|
patch = self._get_module_overwrite_key(module)
|
|
cfg = util.merge_dict(cfg, patch)
|
|
|
|
# step 3: save merged config to merge_config.json
|
|
merge_config = self._log.get_module_merge_config(module)
|
|
util.save_json_file(merge_config, cfg)
|
|
|
|
return merge_config
|
|
|
|
def copy_module_results_to_top_dir(self, module):
|
|
# copy best_profile_result.txt/best_config.json/best_compiler_result/ to top directory
|
|
best_profile_result = self._log.get_module_best_profile_result(module)
|
|
shutil.copy(best_profile_result, self._log.log_dir)
|
|
best_config = self._log.get_module_best_config(module)
|
|
shutil.copy(best_config, self._log.log_dir)
|
|
model_best_compiler_result = self._log.get_best_compiler_result(self._log.log_dir)
|
|
if util.check_file_exist(model_best_compiler_result):
|
|
shutil.rmtree(model_best_compiler_result)
|
|
best_compiler_result = self._log.get_module_best_compiler_result(module)
|
|
shutil.copytree(best_compiler_result, model_best_compiler_result)
|