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)