def test_autotune():
"""Verify that autotune works with micro."""
import tvm.relay as relay
data = relay.var("data", relay.TensorType((1, 3, 64, 64), "float32"))
weight = relay.var("weight", relay.TensorType((8, 3, 5, 5), "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)
main_func = mod["main"]
shape_dict = {
p.name_hint: p.checked_type.concrete_shape
for p in main_func.params
}
type_dict = {p.name_hint: p.checked_type.dtype for p in main_func.params}
weight_data = np.ones(shape_dict["weight"]).astype(type_dict["weight"])
input_data = np.ones(shape_dict["data"]).astype(type_dict["data"])
params = {"weight": weight_data}
inputs = {"data": input_data}
target = tvm.target.target.micro("host")
template_project_dir = pathlib.Path(
tvm.micro.get_standalone_crt_dir()) / "template" / "host"
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
module_loader = tvm.micro.AutoTvmModuleLoader(
template_project_dir=template_project_dir,
project_options={},
)
builder = tvm.autotvm.LocalBuilder(
n_parallel=1,
build_kwargs={"build_option": {
"tir.disable_vectorize": True
}},
do_fork=True,
build_func=tvm.micro.autotvm_build_func,
)
runner = tvm.autotvm.LocalRunner(number=1,
repeat=1,
module_loader=module_loader)
measure_option = tvm.autotvm.measure_option(builder=builder, runner=runner)
tune_log_file = pathlib.Path("crt_autotune.log")
if tune_log_file.exists():
tune_log_file.unlink()
num_trials = 10
for task in tasks:
tuner = tvm.autotvm.tuner.GATuner(task)
tuner.tune(
n_trial=num_trials,
measure_option=measure_option,
callbacks=[
tvm.autotvm.callback.log_to_file(str(tune_log_file)),
tvm.autotvm.callback.progress_bar(num_trials, si_prefix="M"),
],
si_prefix="M",
)
assert tuner.best_flops > 0
check_tune_log(tune_log_file)
# Build without tuning
with pass_context:
lowered = tvm.relay.build(mod, target=TARGET, params=params)
temp_dir = tvm.contrib.utils.tempdir()
project = tvm.micro.generate_project(template_project_dir, lowered,
temp_dir / "project")
project.build()
with tvm.micro.Session(project.transport()) 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(**inputs)
expected_output = graph_mod.get_output(0).numpy()
del graph_mod
# Build using autotune logs
with tvm.autotvm.apply_history_best(str(tune_log_file)):
with pass_context:
lowered_tuned = tvm.relay.build(mod, target=target, params=params)
temp_dir = tvm.contrib.utils.tempdir()
project = tvm.micro.generate_project(template_project_dir, lowered_tuned,
temp_dir / "project")
project.build()
with tvm.micro.Session(project.transport()) 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(**inputs)
output = graph_mod.get_output(0).numpy()
del graph_mod
tvm.testing.assert_allclose(output, expected_output, rtol=1e-4, atol=1e-5)
def test_autotune_conv2d(temp_dir, board, west_cmd, tvm_debug):
"""Test AutoTune for microTVM Zephyr"""
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,
)
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, 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, 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)