[AUTOTVM] End2End autotvm support for vta (apache#18)

* support tuning a whole network

* pass unit test

* update tune resnet

* update all

python/tvm/autotvm/measure/measure_methods.py

@@ -207,7 +207,18 @@ def set_task(self, task):
  for x in arg_bufs]
  func = build(s, arg_bufs, "llvm")
  tvm_buf = [nd.array(x) for x in self.ref_input]
- func(*tvm_buf)
+
+ def _run_func():
+ """Run tvm function in a thread.
+ Because there is some issues with python multiprocessing and the thread pool in tvm
+ """
+ func(*tvm_buf)
+
+ thread = threading.Thread(target=_run_func)
+ thread.start()
+ thread.join()
+ del thread
+
  self.ref_output = [x.asnumpy() for x in tvm_buf]
 
  def get_build_kwargs(self):

python/tvm/autotvm/task/nnvm_integration.py

@@ -5,6 +5,7 @@
 """
 import warnings
 import logging
+import sys
 
 
 from ... import target as _target
@@ -18,8 +19,7 @@
 def extract_from_graph(graph, shape, dtype, target, symbols, target_host=None):
  """ Extract tuning tasks from a nnvm graph.
 
- This function collects tuning tasks by building the graph
- with a "tracing" target and tracing all the calls to topi.
+ This function collects tuning tasks by building the graph and trace all the calls to topi.
 
  Parameters
  ----------
@@ -45,7 +45,7 @@ def extract_from_graph(graph, shape, dtype, target, symbols, target_host=None):
  import nnvm
  import topi
 
- env = TaskExtractEnv.get()
+ env = TaskExtractEnv(symbols)
 
  #NOTE: To add more symbols, you only need to change the following lists
  #nnvm symbol -> topi compute
@@ -63,26 +63,23 @@ def extract_from_graph(graph, shape, dtype, target, symbols, target_host=None):
  else:
  warnings.warn("Symbol %s is not tunable, ignored" % sym_name)
 
- # run compiler to collect all TOPI calls during compilation
- env.reset(topi_funcs)
+ # run compiler to collect all TOPI calls during compilation
+ nnvm.compiler.engine.clear_cache()
+ nnvm.compiler.build(graph, target=target, shape=shape, dtype=dtype)
+ nnvm.compiler.engine.clear_cache()
 
- # disable logger temporarily
- old_state = logger.disabled
- logger.disabled = True
-
- # use a "tracing" target to do a fake compile for collecting topi calls
- tracing_target = _target.create("llvm -device=tracing")
- nnvm.compiler.engine.clear_cache()
- nnvm.compiler.build(graph, target=tracing_target, shape=shape, dtype=dtype)
-
- logger.disabled = old_state
+ logger.disabled = old_state
 
  # create tasks for target
  tasks = []
  for task_name, args in env.get_tasks():
- tasks.append(create(task_name, args,
- target=target, target_host=target_host,
- template_key='direct'))
+ try:
+ tsk = create(task_name, args,
+ target=target, target_host=target_host,
+ template_key='direct')
+ tasks.append(tsk)
+ except topi.InvalidShapeError:
+ print("shape error")
 
  return tasks
 

python/tvm/target.py

@@ -473,6 +473,13 @@ def rasp(options=None):
  return arm_cpu('rasp3b', options)
 
 
+def vta(model='unknown', options=None):
+ opts = ["-device=vta", '-keys=cpu', '-model=%s' % model]
+ opts = _merge_opts(opts, options)
+ ret = _api_internal._TargetCreate("ext_dev", *opts)
+ return ret
+
+
 def create(target_str):
  """Get a target given target string.
 

src/codegen/build_module.cc

@@ -39,6 +39,7 @@ Target CreateTarget(const std::string& target_name,
 
  std::string libs_flag = "-libs=";
  std::string device_flag = "-device=";
+ std::string keys_flag = "-keys=";
  for (auto& item : options) {
  t->options_array.push_back(ir::StringImm::make(item));
 
@@ -50,12 +51,16 @@ Target CreateTarget(const std::string& target_name,
  }
  } else if (item.find(device_flag) == 0) {
  t->device_name = item.substr(device_flag.length());
+ t->keys_array.push_back(ir::StringImm::make(t->device_name));
+ } else if (item.find(keys_flag) == 0) {
+ std::stringstream ss(item.substr(keys_flag.length()));
+ std::string key_item;
+ while (std::getline(ss, key_item, ',')) {
+ t->keys_array.push_back(ir::StringImm::make(key_item));
+ }
  }
  }
 
- if (t->device_name.length() > 0) {
- t->keys_array.push_back(ir::StringImm::make(t->device_name));
- }
  t->device_type = kDLCPU;
  t->thread_warp_size = 1;
  if (target_name == "c" || target_name == "llvm") {

topi/python/topi/__init__.py

@@ -34,6 +34,10 @@
 from . import image
 from . import sparse
 from . import hls
+
+# some short cut
+from .util import InvalidShapeError
+
 # not import testing by default
 # because testing can have extra deps that are not necessary
 # we can import them from test cases explicitly

topi/python/topi/util.py

@@ -6,6 +6,10 @@
 import tvm
 from . import tag
 
+class InvalidShapeError(ValueError):
+ """Invalid shape for a topi function. i.e. call winograd template for non-3x3 kernel)"""
+ pass
+
 def traverse_inline(s, final_op, callback):
  """Traverse computation graph and do auto inline
 

vta/python/vta/__init__.py

@@ -18,6 +18,7 @@
 # to maintain minimum dependency on the board
 if sys.argv[0] not in ("-c", "-m"):
  from . import top
- from .build_module import build_config, lower, build
  from . import graph
+
+ from .build_module import build_config, lower, build, vta_autotvm_build_func
  from .ptr_alias import reinterpret

vta/python/vta/build_module.py

@@ -102,3 +102,39 @@ def build(*args, **kwargs):
  with build_config():
  return tvm.build(*args, **kwargs)
  return tvm.build(*args, **kwargs)
+
+
+def vta_autotvm_build_func(measure_input, tmp_dir, **kwargs):
+ """Custom build func for VTA. Used for autotvm"""
+
+ import time
+ import os
+ from random import getrandbits
+ from tvm.autotvm.util import get_const_tuple
+ from tvm.autotvm.measure.measure_methods import BuildResult, InstantiationError
+
+ tic = time.time()
+ try:
+ filename = os.path.join(tmp_dir, "tmp_func_%0x.tar" % getrandbits(64))
+ target, task, config = measure_input
+
+ with target:
+ s, args = task.instantiate(config)
+ if not config.valid():
+ raise InstantiationError(config.errors)
+
+ func = build(s, args, target_host=task.target_host)
+ func2 = build(s, args)
+
+ arg_info = tuple((get_const_tuple(x.shape), x.dtype) for x in args)
+ func.export_library(filename)
+
+ # check by local simulator
+ ctx = tvm.context(str(target))
+ args = [tvm.nd.empty(x[0], dtype=x[1], ctx=ctx) for x in arg_info]
+ func2(*args)
+
+ except Exception as e: # pylint: disable=broad-except
+ return BuildResult(None, None, e, time.time() - tic)
+ return BuildResult(filename, arg_info, None, time.time() - tic)
+

vta/python/vta/environment.py

@@ -227,34 +227,24 @@ def gemm(self):
  """GEMM intrinsic"""
  return self.dev.gemm
 
- # TODO get rid of it
  @property
- def target_host(self):
- """The target host"""
- return "llvm " + self.llvm_triple
+ def target(self):
+ return tvm.target.vta(model=self.TARGET)
 
  @property
- def target_vta_cpu(self):
+ def target_host(self):
  """The target host"""
  if self.TARGET == "pynq":
- return "llvm -device=arm_cpu -model=pynq {}".format(self.llvm_triple)
+ return "llvm -target=armv7-none-linux-gnueabihf"
  elif self.TARGET == "ultra96":
- return "llvm -device=arm_cpu -model=ultra96 {}".format(self.llvm_triple)
+ return "llvm -target=aarch64-linux-gnu"
  elif self.TARGET == "sim":
  return "llvm"
  raise ValueError("Unknown target %s" % self.TARGET)
 
  @property
- def llvm_triple(self):
- """The llvm flags for the target platform"""
- if self.TARGET == "pynq":
- return "-target=armv7-none-linux-gnueabihf"
- elif self.TARGET == "ultra96":
- return "-target=aarch64-linux-gnu"
- elif self.TARGET == "sim":
- return ""
- else:
- raise ValueError("Unknown target %s" % self.TARGET)
+ def target_vta_cpu(self):
+ return tvm.target.arm_cpu(model=self.TARGET)
 
 
 def get_env():

vta/python/vta/top/__init__.py

@@ -1,11 +1,4 @@
 """TVM TOPI connector, eventually most of these should go to TVM repo"""
 
 from . import vta_conv2d
-from . import arm_conv2d
-from . import testing
-
-from .bitpack import bitpack
-from .vta_dense import packed_dense, schedule_packed_dense
-from .vta_conv2d import packed_conv2d, schedule_packed_conv2d
-from .vta_group_conv2d import packed_group_conv2d, schedule_packed_group_conv2d
-from .vta_conv2d_transpose import packed_conv2d_transpose, schedule_packed_conv2d_transpose
+from . import op