Skip to content

Commit

Permalink
[µTVM] Enable AutoTVM for ARM STM32F746XX Boards (#4274)
Browse files Browse the repository at this point in the history
  • Loading branch information
@weberlo @tqchen
weberlo authored and tqchen committed Dec 2, 2019
1 parent 11af82c commit 47c870a
Show file tree
Hide file tree
Showing 37 changed files with 1,985 additions and 879 deletions.
11 changes: 7 additions & 4 deletions python/tvm/_ffi/ndarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down
223 changes: 128 additions & 95 deletions python/tvm/contrib/binutil.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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".
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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()
Expand Down
52 changes: 51 additions & 1 deletion python/tvm/exec/rpc_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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,
Expand All @@ -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",
Expand All @@ -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()
Expand Down
Loading

0 comments on commit 47c870a

Please sign in to comment.