Support self containing sub dir & loop

Change-Id: I5c48890021d46694c3b0e3fb52b30194049ed006

python/tpu_perf/build.py

@@ -5,6 +5,7 @@
 import re
 from .buildtree import check_buildtree, BuildTree
 from .subp import CommandExecutor, sys_memory_size
+from .util import *
 
 def replace_shape_batch(cmd, batch_size):
  match = re.search('-shapes *(\[.*\])', cmd)
@@ -79,31 +80,6 @@ def build_nntc(tree, path, config):
  if 'shape_key' in config:
  name = f'{name}-{config["shape_key"]}'
 
- if not option_time_only and 'fp32_compile_options' in config:
- fp32_pool = CommandExecutor(workdir, env)
- cmd = tree.expand_variables(config, config["fp32_compile_options"])
- logging.info(f'Building FP32 bmodel {name}...')
- if 'fp32_batch_sizes' in config:
- for batch_size in config['fp32_batch_sizes']:
- batch_cmd = replace_shape_batch(cmd, batch_size)
- fp32_pool.put(
- f'{batch_size}b.fp32.compile',
- f'{batch_cmd} --outdir {batch_size}b.fp32.compilation')
- else:
- outname = 'fp32.compilation'
- match = re.search('-shapes *(\[.*\])', cmd)
- if match is not None:
- shapes = eval(match.group(1))
- if type(shapes[0]) != list:
- shapes = [shapes]
- batch = shapes[0][0]
- outname = f'{batch}b.' + outname
- fp32_pool.put(
- f'fp32.compile',
- f'{cmd} --outdir {outname}')
- fp32_pool.wait()
- logging.info(f'FP32 bmodel {name} done.')
-
  int8_pool = CommandExecutor(
  workdir, env,
  memory_hint=config.get('memory_hint'))
@@ -125,16 +101,47 @@ def build_nntc(tree, path, config):
  logging.info(f'Calibrating {name}...')
  int8_pool.run(cali_key, cmd)
 
+ if 'fp32_compile_options' in config:
+ fp32_pool = CommandExecutor(workdir, env)
+ logging.info(f'Building FP32 bmodel {name}...')
+ batch_sizes = config.get('fp32_batch_sizes', [1]) \
+ if not option_time_only else [1]
+ fp32_loops = config.get('fp32_loops') or \
+ tree.global_config.get('fp32_loops') or [dict()]
+ for loop in fp32_loops:
+ loop_config = dict_override(config, loop)
+ cmd = tree.expand_variables(loop_config, loop_config["fp32_compile_options"])
+ for batch_size in batch_sizes:
+ if 'fp32_batch_sizes' in config:
+ batch_cmd = replace_shape_batch(cmd, batch_size)
+ else:
+ batch_cmd = cmd
+ outdir = loop_config.get(
+ 'fp32_outdir_template',
+ '{}b.fp32.compilation').format(batch_size)
+ fp32_pool.put(
+ outdir,
+ f'{batch_cmd} --outdir {outdir}')
+ fp32_pool.wait()
+ logging.info(f'FP32 bmodel {name} done.')
+
  if 'bmnetu_options' in config:
  # Build int8 bmodel
- cmd = f'python3 -m bmnetu {config["bmnetu_options"]}'
- cmd = tree.expand_variables(config, cmd)
  logging.info(f'Compiling {name}...')
- for b in config['bmnetu_batch_sizes']:
- int8_pool.put(
- f'compile.{b}',
- f'{cmd} --max_n {b} --outdir {b}b.compilation')
+ int8_loops = config.get('int8_loops') or \
+ tree.global_config.get('int8_loops') or [dict()]
+ for loop in int8_loops:
+ loop_config = dict_override(config, loop)
+ cmd = f'python3 -m bmnetu {loop_config["bmnetu_options"]}'
+ cmd = tree.expand_variables(loop_config, cmd)
+ for b in loop_config['bmnetu_batch_sizes']:
+ outdir = loop_config.get(
+ 'int8_outdir_template', '{}b.compilation').format(b)
+ int8_pool.put(
+ outdir,
+ f'{cmd} --max_n {b} --outdir {outdir}')
  int8_pool.wait()
+ logging.info(f'INT8 bmodel {name} done.')
 
 def main():
  logging.basicConfig(

python/tpu_perf/buildtree.py

@@ -5,6 +5,8 @@
 import copy
 import logging
 
+from .util import *
+
 def read_config(path):
  fn = os.path.join(path, 'config.yaml')
  if not os.path.exists(fn):
@@ -13,15 +15,6 @@ def read_config(path):
  with open(fn) as f:
  return yaml.load(f, yaml.Loader)
 
-def hash_name(config):
- from hashlib import md5
- m = md5()
- keys = [k for k in config]
- keys.sort()
- for k in keys:
- m.update(str(config[k]).encode())
- return m.hexdigest()
-
 def shape_key_and_param(shape):
  if type(shape) != list:
  logging.error(f'shape should be list')
@@ -47,8 +40,8 @@ def __init__(self, root, args = None):
  self.root = root
  self.global_config = global_config = read_config(root) or dict()
  global_config['root'] = root
- if 'workdir' not in global_config:
- global_config['workdir'] = os.path.join(root, 'output')
+ if 'outdir' not in global_config:
+ global_config['outdir'] = os.path.join(root, 'output')
 
  self.cases = []
  if not args.full:
@@ -61,8 +54,18 @@ def __init__(self, root, args = None):
  self.cases = lines
  if args.models:
  self.cases = args.models
+ global_config['target'] = self.target = args.target
+ global_config['devices'] = args.devices
 
- self.output_names = set()
+ # Target specific configs
+ if self.target in global_config:
+ specific = global_config.pop(self.target)
+ self.global_config = global_config = dict_override(global_config, specific)
+
+ # Expand all $(home) in global config
+ # This allow sub directory rejection
+ self.global_config = self.expand_all_variables(
+ dict(home=self.root), global_config, shallow=True, no_except=True)
 
  @staticmethod
  def add_arguments(parser):
@@ -71,6 +74,13 @@ def add_arguments(parser):
  help='model directories to run')
  parser.add_argument('--full', action='store_true', help='Run all cases')
  parser.add_argument('--list', '-l', type=str, help='Case list')
+ parser.add_argument('--devices', '-d',
+ type=int, nargs='*', help='Devices',
+ default=[0])
+ parser.add_argument(
+ '--target', '-t', type=str, default='BM1684X',
+ choices=['BM1684', 'BM1684X'],
+ help='Target chip')
 
  def read_global_variable(self, name, config = dict(), default=None):
  if default is None and name not in self.global_config:
@@ -81,7 +91,10 @@ def read_global_variable(self, name, config = dict(), default=None):
 
  whole_var_pattern = '^\$\(([a-z0-9_]+)\)$'
 
- def expand_variables(self, config, string, stack=None, shallow=False):
+ def expand_variables(
+ self, config, string,
+ stack=None, shallow=False, no_except=False):
+
  if type(string) != str:
  return string
  if stack is None:
@@ -92,6 +105,8 @@ def expand_variables(self, config, string, stack=None, shallow=False):
  var = var_match.group(1)
  value = config.get(var) or global_config.get(var)
  if not value:
+ if no_except:
+ return string
  logging.error(f'Invalid variable "{var}" in {string}')
  raise Exception('invalid variable')
  if shallow:
@@ -104,6 +119,8 @@ def expand_variables(self, config, string, stack=None, shallow=False):
  var = m.group(1)
  value = config.get(var) or global_config.get(var)
  if not value:
+ if no_except:
+ continue
  logging.error(f'Invalid variable "{var}" in {string}')
  raise Exception('invalid variable')
  if var in stack:
@@ -116,17 +133,78 @@ def expand_variables(self, config, string, stack=None, shallow=False):
  out = out.replace(raw, str(value))
  return out
 
+ def expand_all_variables(self, config, data, **kw_args):
+ if type(data) == str:
+ return self.expand_variables(config, data, **kw_args)
+ elif type(data) == list:
+ data = data.copy()
+ for i in range(len(data)):
+ data[i] = self.expand_all_variables(config, data[i], **kw_args)
+ elif type(data) == dict:
+ data = data.copy()
+ for k, v in data.items():
+ data[k] = self.expand_all_variables(config, v, **kw_args)
+ return data
+
+ def expand_all_whole_variables(self, config, data, **kw_args):
+ if type(data) == str:
+ if not re.match(self.whole_var_pattern, data):
+ return data
+ return self.expand_variables(config, data, **kw_args)
+ elif type(data) == list:
+ data = data.copy()
+ for i in range(len(data)):
+ data[i] = self.expand_all_whole_variables(config, data[i], **kw_args)
+ elif type(data) == dict:
+ data = data.copy()
+ for k, v in data.items():
+ data[k] = self.expand_all_whole_variables(config, v, **kw_args)
+ return data
+
  def read_dir(self, path):
+ path = os.path.abspath(path)
+ context = dict()
+ p = path
+ while True:
+ p = os.path.dirname(p)
+ if p == self.root:
+ break
+ if len(p) < 2:
+ break
+ fn = os.path.join(p, 'config.yaml')
+ if not os.path.isfile(fn):
+ continue
+ context = read_config(p) or dict()
+ if self.target in context:
+ specific = context.pop(self.target)
+ context = dict_override(context, specific)
+ context = self.expand_all_variables(
+ dict(home=p), context, shallow=True, no_except=True)
+ break
  for fn in os.listdir(path):
  if not fn.endswith('config.yaml'):
  continue
  fn = os.path.join(path, fn)
  if not os.path.isfile(fn):
  continue
- for ret in self._read_dir(fn):
+ for ret in self._read_dir(fn, context):
  yield ret
 
- def _read_dir(self, config_fn):
+ def hash_name(self, config):
+ from hashlib import md5
+ m = md5()
+ keys = [k for k in config]
+ keys.sort()
+ for k in keys:
+ v = config[k]
+ if type(v) == str and '/' in v and os.path.exists(v):
+ # v is a path
+ # Convert to relative path
+ v = os.path.relpath(v, self.root)
+ m.update(f'{k}: {v};'.encode())
+ return m.hexdigest()
+
+ def _read_dir(self, config_fn, context = dict()):
  path = os.path.dirname(config_fn)
  global_config = self.global_config
 
@@ -138,22 +216,20 @@ def _read_dir(self, config_fn):
  if config.get('ignore'):
  return
 
+ if self.target in config:
+ specific = config.pop(self.target)
+ config = dict_override(config, specific)
+
+ config = dict_override(config, context)
+
  if 'name' not in config:
  logging.error(f'Invalid config {config_fn}')
  raise RuntimeError('Invalid config')
 
  # Pre expand non-string variables. Because variable
  # processing logic might rely on this.
- for k in config:
- if type(config[k]) != str:
- continue
- if not re.match(self.whole_var_pattern, config[k]):
- continue
- try:
- config[k] = self.expand_variables(config, config[k], shallow=True)
- except:
- print(k, config[k])
- raise
+ config = self.expand_all_whole_variables(config, config, shallow=True)
+
  if 'harness' in config and type(config['harness']['args']) != list:
  config['harness']['args'] = [config['harness']['args']]
 
@@ -180,7 +256,7 @@ def _read_dir(self, config_fn):
  logging.error(f'Duplicate output name {name}')
  raise RuntimeError('invalid output name')
  self.output_names.add(name)
- workdir = os.path.join(global_config['workdir'], name)
+ workdir = os.path.join(global_config['outdir'], name)
  os.makedirs(workdir, exist_ok=True)
 
  # Default configuration
@@ -190,14 +266,15 @@ def _read_dir(self, config_fn):
  config['bmnetu_batch_sizes'] = [1]
 
  if 'input' in config:
- key = hash_name(config['input'])
+ key = self.hash_name(config['input'])
  data_dir = self.read_global_variable(
  'data_dir', default='$(root)/data')
  config['lmdb_out'] = os.path.join(data_dir, key)
 
  yield path, copy.deepcopy(config)
 
  def walk(self, path=None):
+ self.output_names = set()
  if path is None:
  if self.cases:
  for path in self.cases:

python/tpu_perf/eject.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+
+import os
+
+def main():
+ import argparse
+ parser = argparse.ArgumentParser(description='Eject dir in model zoo')
+ parser.add_argument(
+ '--extra', '-e', type=str,
+ help='Extra file list', required=True)
+ parser.add_argument(
+ '--out', '-O', type=str,
+ help='Output tar file', required=True)
+ parser.add_argument(
+ 'dir', metavar='DIR', type=str,
+ help='Directory to eject')
+ args = parser.parse_args()
+
+ with open(args.extra) as f:
+ extra_files = [line.strip(' \n') for line in f.readlines()]
+ extra_files = [line for line in extra_files if line]
+
+ import tarfile
+ with tarfile.open(args.out, 'w:bz2') as tar:
+ for fn in extra_files:
+ tar.add(
+ fn, recursive=False,
+ arcname=os.path.join(f'model-zoo-{args.dir}', fn))
+ for fn in os.listdir(args.dir):
+ tar.add(
+ os.path.join(args.dir, fn),
+ arcname=os.path.join(f'model-zoo-{args.dir}', fn))
+
+if __name__ == '__main__':
+ main()

python/tpu_perf/harness/__init__.py

@@ -1,12 +1,8 @@
 _harness_functions = dict()
 
 def load_plugins():
- import logging
- import importlib
- try:
- importlib.import_module('harness')
- except ModuleNotFoundError:
- pass
+ from .. import util
+ util.load_plugins('harness')
 
 def harness(key):
  def register(fn):