[µTVM] Enable AutoTVM for ARM STM32F746XX Boards

python/tvm/_ffi/ndarray.py

@@ -75,10 +75,13 @@ def context(dev_type, dev_id=0):
       assert tvm.context("cuda", 0) == tvm.gpu(0)
     """
     if isinstance(dev_type, string_types):
-        dev_type = dev_type.split()[0]
-        if dev_type not in TVMContext.STR2MASK:
-            raise ValueError("Unknown device type %s" % dev_type)
-        dev_type = TVMContext.STR2MASK[dev_type]
+        if '-device=micro_dev' in dev_type:
+            dev_type = 'micro_dev'
+        else:
+            dev_type = dev_type.split()[0]
+            if dev_type not in TVMContext.STR2MASK:
+                raise ValueError("Unknown device type %s" % dev_type)
+            dev_type = TVMContext.STR2MASK[dev_type]
     return TVMContext(dev_type, dev_id)
 
 

python/tvm/contrib/binutil.py

@@ -19,9 +19,81 @@
 import os
 import subprocess
 from . import util
-from .._ffi.base import py_str
 from ..api import register_func
 
+RELOCATION_LD_SCRIPT_TEMPLATE = """
+/* linker symbol for use in UTVMInit */
+_utvm_stack_pointer_init = 0x{stack_pointer_init:x};
+
+SECTIONS
+{{
+  . = 0x{text_start:x};
+  . = ALIGN({word_size});
+  .text :
+  {{
+    . = ALIGN({word_size});
+    KEEP(*(.text))
+    KEEP(*(.text*))
+    . = ALIGN({word_size});
+  }}
+
+  . = 0x{rodata_start:x};
+  . = ALIGN({word_size});
+  .rodata :
+  {{
+    . = ALIGN({word_size});
+    KEEP(*(.rodata))
+    KEEP(*(.rodata*))
+    . = ALIGN({word_size});
+  }}
+
+  . = 0x{data_start:x};
+  . = ALIGN({word_size});
+  .data :
+  {{
+    . = ALIGN({word_size});
+    KEEP(*(.data))
+    KEEP(*(.data*))
+    . = ALIGN({word_size});
+  }}
+
+  . = 0x{bss_start:x};
+  . = ALIGN({word_size});
+  .bss :
+  {{
+    . = ALIGN({word_size});
+    KEEP(*(.bss))
+    KEEP(*(.bss*))
+    . = ALIGN({word_size});
+  }}
+}}
+"""
+
+def run_cmd(cmd):
+    """Runs `cmd` in a subprocess and awaits its completion.
+
+    Parameters
+    ----------
+    cmd : List[str]
+        list of command-line arguments
+
+    Returns
+    -------
+    output : str
+        resulting stdout capture from the subprocess
+    """
+    proc = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT)
+    (output, _) = proc.communicate()
+    output = output.decode("utf-8")
+    if proc.returncode != 0:
+        cmd_str = " ".join(cmd)
+        msg = f"error while running command \"{cmd_str}\":\n{output}"
+        raise RuntimeError(msg)
+    return output
+
 
 @register_func("tvm_callback_get_section_size")
 def tvm_callback_get_section_size(binary_path, section_name, toolchain_prefix):
@@ -48,14 +120,7 @@ def tvm_callback_get_section_size(binary_path, section_name, toolchain_prefix):
         raise RuntimeError("no such file \"{}\"".format(binary_path))
     # We use the "-A" flag here to get the ".rodata" section's size, which is
     # not included by default.
-    size_proc = subprocess.Popen(
-        ["{}size".format(toolchain_prefix), "-A", binary_path], stdout=subprocess.PIPE)
-    (size_output, _) = size_proc.communicate()
-    size_output = size_output.decode("utf-8")
-    if size_proc.returncode != 0:
-        msg = "error in finding section size:\n"
-        msg += py_str(size_output)
-        raise RuntimeError(msg)
+    size_output = run_cmd(["{}size".format(toolchain_prefix), "-A", binary_path])
 
     # TODO(weberlo): Refactor this method and `*relocate_binary` so they are
     # both aware of [".bss", ".sbss", ".sdata"] being relocated to ".bss".
@@ -74,13 +139,15 @@ def tvm_callback_get_section_size(binary_path, section_name, toolchain_prefix):
             continue
         entry_name = tokens[0]
         entry_size = int(tokens[1])
-        if entry_name in sections_to_sum:
-            section_size += entry_size
+        for section in sections_to_sum:
+            if entry_name.startswith(section):
+                section_size += entry_size
+                break
 
     # NOTE: For some reason, the size of the BSS section on the RISC-V
     # GCC is sometimes reported to be smaller than it is, so we need to adjust
     # for this.
-    if "riscv" in toolchain_prefix and section_name == 'bss':
+    if "riscv" in toolchain_prefix and section_name == "bss":
         # TODO(weberlo): Figure out why 32 is the minimum constant that works.
         #
         # The current hypothesis is that the last symbols in the ".bss" and
@@ -97,25 +164,38 @@ def tvm_callback_get_section_size(binary_path, section_name, toolchain_prefix):
 
 @register_func("tvm_callback_relocate_binary")
 def tvm_callback_relocate_binary(
-        binary_path, text_addr, rodata_addr, data_addr, bss_addr, toolchain_prefix):
+        binary_path,
+        word_size,
+        text_start,
+        rodata_start,
+        data_start,
+        bss_start,
+        stack_end,
+        toolchain_prefix):
     """Relocates sections in the binary to new addresses
 
     Parameters
     ----------
     binary_path : str
         path of the binary file
 
-    text_addr : str
-        text section absolute address
+    word_size : int
+        word size on the target machine
+
+    text_start : int
+        text section address
+
+    rodata_start : int
+        rodata section address
 
-    rodata_addr : str
-        rodata section absolute address
+    data_start : int
+        data section address
 
-    data_addr : str
-        data section absolute address
+    bss_start : int
+        bss section address
 
-    bss_addr : str
-        bss section absolute address
+    stack_end : int
+        stack section end address
 
     toolchain_prefix : str
         prefix for binary names in target compiler toolchain
@@ -125,68 +205,29 @@ def tvm_callback_relocate_binary(
     rel_bin : bytearray
         the relocated binary
     """
-    tmp_dir = util.tempdir()
-    rel_obj_path = tmp_dir.relpath("relocated.o")
+    stack_pointer_init = stack_end - word_size
     ld_script_contents = ""
     # TODO(weberlo): There should be a better way to configure this for different archs.
     if "riscv" in toolchain_prefix:
         ld_script_contents += "OUTPUT_ARCH( \"riscv\" )\n\n"
-    # TODO(weberlo): Generate the script in a more procedural manner.
-    ld_script_contents += """
-SECTIONS
-{
-  . = %s;
-  . = ALIGN(8);
-  .text :
-  {
-    *(.text)
-    . = ALIGN(8);
-    *(.text*)
-  }
-  . = %s;
-  . = ALIGN(8);
-  .rodata :
-  {
-    *(.rodata)
-    . = ALIGN(8);
-    *(.rodata*)
-  }
-  . = %s;
-  . = ALIGN(8);
-  .data :
-  {
-    *(.data)
-    . = ALIGN(8);
-    *(.data*)
-    . = ALIGN(8);
-    *(.sdata)
-  }
-  . = %s;
-  . = ALIGN(8);
-  .bss :
-  {
-    *(.bss)
-    . = ALIGN(8);
-    *(.bss*)
-    . = ALIGN(8);
-    *(.sbss)
-  }
-}
-    """ % (text_addr, rodata_addr, data_addr, bss_addr)
+    ld_script_contents += RELOCATION_LD_SCRIPT_TEMPLATE.format(
+        word_size=word_size,
+        text_start=text_start,
+        rodata_start=rodata_start,
+        data_start=data_start,
+        bss_start=bss_start,
+        stack_pointer_init=stack_pointer_init)
+
+    tmp_dir = util.tempdir()
+    rel_obj_path = tmp_dir.relpath("relocated.obj")
     rel_ld_script_path = tmp_dir.relpath("relocated.lds")
     with open(rel_ld_script_path, "w") as f:
         f.write(ld_script_contents)
-    ld_proc = subprocess.Popen(["{}ld".format(toolchain_prefix), binary_path,
-                                "-T", rel_ld_script_path,
-                                "-o", rel_obj_path],
-                               stdout=subprocess.PIPE,
-                               stderr=subprocess.STDOUT)
-    (out, _) = ld_proc.communicate()
-    if ld_proc.returncode != 0:
-        msg = "linking error using ld:\n"
-        msg += py_str(out)
-        raise RuntimeError(msg)
-
+    run_cmd([
+        "{}ld".format(toolchain_prefix),
+        binary_path,
+        "-T", rel_ld_script_path,
+        "-o", rel_obj_path])
     with open(rel_obj_path, "rb") as f:
         rel_bin = bytearray(f.read())
     return rel_bin
@@ -217,16 +258,11 @@ def tvm_callback_read_binary_section(binary, section, toolchain_prefix):
     tmp_section = tmp_dir.relpath("tmp_section.bin")
     with open(tmp_bin, "wb") as out_file:
         out_file.write(bytes(binary))
-    objcopy_proc = subprocess.Popen(["{}objcopy".format(toolchain_prefix), "--dump-section",
-                                     ".{}={}".format(section, tmp_section),
-                                     tmp_bin],
-                                    stdout=subprocess.PIPE,
-                                    stderr=subprocess.STDOUT)
-    (out, _) = objcopy_proc.communicate()
-    if objcopy_proc.returncode != 0:
-        msg = "error in using objcopy:\n"
-        msg += py_str(out)
-        raise RuntimeError(msg)
+    run_cmd([
+        "{}objcopy".format(toolchain_prefix),
+        "--dump-section",
+        ".{}={}".format(section, tmp_section),
+        tmp_bin])
     if os.path.isfile(tmp_section):
         # Get section content if it exists.
         with open(tmp_section, "rb") as f:
@@ -259,15 +295,12 @@ def tvm_callback_get_symbol_map(binary, toolchain_prefix):
     tmp_obj = tmp_dir.relpath("tmp_obj.bin")
     with open(tmp_obj, "wb") as out_file:
         out_file.write(bytes(binary))
-    nm_proc = subprocess.Popen(["{}nm".format(toolchain_prefix), "-C", "--defined-only", tmp_obj],
-                               stdout=subprocess.PIPE,
-                               stderr=subprocess.STDOUT)
-    (nm_output, _) = nm_proc.communicate()
-    if nm_proc.returncode != 0:
-        msg = "error in using nm:\n"
-        msg += py_str(nm_output)
-        raise RuntimeError(msg)
-    nm_output = nm_output.decode("utf8").splitlines()
+    nm_output = run_cmd([
+        "{}nm".format(toolchain_prefix),
+        "-C",
+        "--defined-only",
+        tmp_obj])
+    nm_output = nm_output.splitlines()
     map_str = ""
     for line in nm_output:
         line = line.split()

python/tvm/exec/rpc_server.py

@@ -19,14 +19,22 @@
 from __future__ import absolute_import
 
 import argparse
+import ast
 import multiprocessing
 import sys
 import logging
+import tvm
+from tvm import micro
 from .. import rpc
 
 def main(args):
-    """Main function"""
+    """Main function
 
+    Parameters
+    ----------
+    args : argparse.Namespace
+        parsed args from command-line invocation
+    """
     if args.tracker:
         url, port = args.tracker.rsplit(":", 1)
         port = int(port)
@@ -37,6 +45,9 @@ def main(args):
     else:
         tracker_addr = None
 
+    if args.utvm_dev_config or args.utvm_dev_id:
+        init_utvm(args)
+
     server = rpc.Server(args.host,
                         args.port,
                         args.port_end,
@@ -48,6 +59,38 @@ def main(args):
     server.proc.join()
 
 
+def init_utvm(args):
+    """MicroTVM-specific RPC initialization
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        parsed args from command-line invocation
+    """
+    if args.utvm_dev_config and args.utvm_dev_id:
+        raise RuntimeError('only one of --utvm-dev-config and --utvm-dev-id allowed')
+
+    if args.utvm_dev_config:
+        with open(args.utvm_dev_config, 'r') as dev_conf_file:
+            dev_config = json.load(dev_conf_file)
+    else:
+        dev_config_args = ast.literal_eval(args.utvm_dev_config_args)
+        default_config_func = micro.device.get_device_funcs(args.utvm_dev_id)['default_config']
+        dev_config = default_config_func(*dev_config_args)
+
+    if args.utvm_dev_config or args.utvm_dev_id:
+        # add MicroTVM overrides
+        @tvm.register_func('tvm.rpc.server.start', override=True)
+        def server_start():
+            # pylint: disable=unused-variable
+            session = micro.Session(dev_config)
+            session._enter()
+
+            @tvm.register_func('tvm.rpc.server.shutdown', override=True)
+            def server_shutdown():
+                session._exit()
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('--host', type=str, default="0.0.0.0",
@@ -71,6 +114,13 @@ def main(args):
                          and ROCM compilers.")
     parser.add_argument('--custom-addr', type=str,
                         help="Custom IP Address to Report to RPC Tracker")
+    parser.add_argument('--utvm-dev-config', type=str,
+                        help='JSON config file for the target device (if using MicroTVM)')
+    parser.add_argument('--utvm-dev-id', type=str,
+                        help='Unique ID for the target device (if using MicroTVM)')
+    parser.add_argument('--utvm-dev-config-args', type=str,
+                        help=('Python list of literals required to generate a default'
+                              ' MicroTVM config (if --utvm-dev-id is specified)'))
 
     parser.set_defaults(fork=True)
     args = parser.parse_args()