def _make_session(temp_dir, zephyr_board, west_cmd, mod, build_config, use_fvp):
config_main_stack_size = None
if test_utils.qemu_boards(zephyr_board):
config_main_stack_size = 1536
project_options = {
"project_type": "host_driven",
"west_cmd": west_cmd,
"verbose": bool(build_config.get("debug")),
"zephyr_board": zephyr_board,
"arm_fvp_path": "/opt/arm/FVP_Corstone_SSE-300/models/Linux64_GCC-6.4/FVP_Corstone_SSE-300_Ethos-U55",
"use_fvp": bool(use_fvp),
}
if config_main_stack_size is not None:
project_options["config_main_stack_size"] = config_main_stack_size
project = tvm.micro.generate_project(
str(test_utils.TEMPLATE_PROJECT_DIR),
mod,
temp_dir / "project",
project_options,
)
project.build()
project.flash()
return tvm.micro.Session(project.transport())
def _make_session(workspace_dir, zephyr_board, west_cmd, mod, build_config,
use_fvp):
config_main_stack_size = None
if test_utils.qemu_boards(zephyr_board):
# fyi: qemu_riscv64 seems to be the greediest stack user
config_main_stack_size = 4096
else:
# increase stack size for HW platforms
config_main_stack_size = 2048
project_options = {
"project_type": "host_driven",
"west_cmd": west_cmd,
"verbose": bool(build_config.get("debug")),
"zephyr_board": zephyr_board,
"arm_fvp_path":
"/opt/arm/FVP_Corstone_SSE-300/models/Linux64_GCC-6.4/FVP_Corstone_SSE-300_Ethos-U55",
"use_fvp": bool(use_fvp),
}
if config_main_stack_size is not None:
project_options["config_main_stack_size"] = config_main_stack_size
project = tvm.micro.generate_project(
str(test_utils.TEMPLATE_PROJECT_DIR),
mod,
workspace_dir / "project",
project_options,
)
project.build()
project.flash()
return tvm.micro.Session(project.transport())
def _make_session(temp_dir, zephyr_board, west_cmd, mod, build_config):
config_main_stack_size = None
if test_utils.qemu_boards(zephyr_board):
config_main_stack_size = 1536
project_options = {
"project_type": "host_driven",
"west_cmd": west_cmd,
"verbose": bool(build_config.get("debug")),
"zephyr_board": zephyr_board,
}
if config_main_stack_size is not None:
project_options["config_main_stack_size"] = config_main_stack_size
project = tvm.micro.generate_project(
str(test_utils.TEMPLATE_PROJECT_DIR),
mod,
temp_dir / "project",
project_options,
)
project.build()
project.flash()
return tvm.micro.Session(project.transport())
def test_autotune_conv2d(temp_dir, board, west_cmd, tvm_debug):
"""Test AutoTune for microTVM Zephyr"""
if board != "qemu_x86":
pytest.xfail(f"Autotune fails on {board}.")
runtime = Runtime("crt", {"system-lib": True})
model = test_utils.ZEPHYR_BOARDS[board]
build_config = {"debug": tvm_debug}
# Create a Relay model
data_shape = (1, 3, 16, 16)
weight_shape = (8, 3, 5, 5)
data = relay.var("data", relay.TensorType(data_shape, "float32"))
weight = relay.var("weight", relay.TensorType(weight_shape, "float32"))
y = relay.nn.conv2d(
data,
weight,
padding=(2, 2),
kernel_size=(5, 5),
kernel_layout="OIHW",
out_dtype="float32",
)
f = relay.Function([data, weight], y)
mod = tvm.IRModule.from_expr(f)
mod = relay.transform.InferType()(mod)
data_sample = np.random.rand(data_shape[0], data_shape[1], data_shape[2],
data_shape[3]).astype("float32")
weight_sample = np.random.rand(weight_shape[0], weight_shape[1],
weight_shape[2],
weight_shape[3]).astype("float32")
params = {mod["main"].params[1].name_hint: weight_sample}
target = tvm.target.target.micro(model)
pass_context = tvm.transform.PassContext(
opt_level=3, config={"tir.disable_vectorize": True})
with pass_context:
tasks = tvm.autotvm.task.extract_from_program(mod["main"], {}, target)
assert len(tasks) > 0
config_main_stack_size = None
if test_utils.qemu_boards(board):
config_main_stack_size = 1536
project_options = {
"zephyr_board": board,
"west_cmd": west_cmd,
"verbose": 1,
"project_type": "host_driven",
}
if config_main_stack_size is not None:
project_options["config_main_stack_size"] = config_main_stack_size
module_loader = tvm.micro.AutoTvmModuleLoader(
template_project_dir=test_utils.TEMPLATE_PROJECT_DIR,
project_options=project_options,
)
timeout = 200
builder = tvm.autotvm.LocalBuilder(
timeout=timeout,
n_parallel=1,
build_kwargs={"build_option": {
"tir.disable_vectorize": True
}},
do_fork=True,
build_func=tvm.micro.autotvm_build_func,
runtime=runtime,
)
runner = tvm.autotvm.LocalRunner(number=1,
repeat=1,
timeout=timeout,
module_loader=module_loader)
measure_option = tvm.autotvm.measure_option(builder=builder, runner=runner)
log_path = pathlib.Path("zephyr_autotune.log")
if log_path.exists():
log_path.unlink()
n_trial = 10
for task in tasks:
tuner = tvm.autotvm.tuner.GATuner(task)
tuner.tune(
n_trial=n_trial,
measure_option=measure_option,
callbacks=[
tvm.autotvm.callback.log_to_file(str(log_path)),
tvm.autotvm.callback.progress_bar(n_trial, si_prefix="M"),
],
si_prefix="M",
)
assert tuner.best_flops > 0
check_tune_log(log_path)
# Build without tuning
with pass_context:
lowered = tvm.relay.build(mod,
target=target,
runtime=runtime,
params=params)
temp_dir = utils.tempdir()
with _make_session(temp_dir, board, west_cmd, lowered,
build_config) as session:
graph_mod = tvm.micro.create_local_graph_executor(
lowered.get_graph_json(), session.get_system_lib(), session.device)
graph_mod.set_input(**lowered.get_params())
graph_mod.run(data=data_sample)
expected_output = graph_mod.get_output(0).numpy()
del graph_mod
# Build using autotune logs
with tvm.autotvm.apply_history_best(str(log_path)):
with pass_context:
lowered_tuned = tvm.relay.build(mod,
target=target,
runtime=runtime,
params=params)
temp_dir = utils.tempdir()
with _make_session(temp_dir, board, west_cmd, lowered_tuned,
build_config) as session:
graph_mod = tvm.micro.create_local_graph_executor(
lowered_tuned.get_graph_json(), session.get_system_lib(),
session.device)
graph_mod.set_input(**lowered_tuned.get_params())
graph_mod.run(data=data_sample)
output = graph_mod.get_output(0).numpy()
del graph_mod
tvm.testing.assert_allclose(output, expected_output, rtol=1e-4, atol=1e-5)
