| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| import os |
| |
| import numpy as np |
| import tvm |
| from tvm import te |
| from tvm.contrib import graph_runtime, util |
| from tvm import relay |
| import tvm.micro as micro |
| from tvm.micro import create_micro_mod |
| |
| # Use real micro device - an STM32F746 discovery board |
| # SETUP: |
| # Be sure to have openocd installed and running |
| # Ex : openocd -f board/stm32f7discovery.cfg |
| # Be sure to have the ST CMSIS library downloaded, installed and |
| # Ex : export CMSIS_ST_PATH="/home/yourid/st/STM32Cube_FW_F7_V1.16.0/Drivers/CMSIS" |
| DEV_CONFIG_A = micro.device.arm.stm32f746xx.generate_config("127.0.0.1", 6666) |
| DEV_CONFIG_B = micro.device.arm.stm32f746xx.generate_config("127.0.0.1", 6666) |
| TARGET = "micro_dev" |
| |
| |
| def relay_micro_build(func, dev_config, params=None): |
| """Create a graph runtime module with a micro device context from a Relay function. |
| |
| Parameters |
| ---------- |
| func : relay.Function |
| function to compile |
| |
| dev_config : Dict[str, Any] |
| MicroTVM config dict for the target device |
| |
| params : dict |
| input parameters that do not change during inference |
| |
| Return |
| ------ |
| mod : tvm.runtime.Module |
| graph runtime module for the target device |
| """ |
| with tvm.transform.PassContext( |
| disabled_pass={"FuseOps"}, config={"tir.disable_vectorize": True} |
| ): |
| graph, c_mod, params = relay.build(func, target=TARGET, params=params) |
| micro_mod = micro.create_micro_mod(c_mod, dev_config) |
| ctx = tvm.micro_dev(0) |
| mod = graph_runtime.create(graph, micro_mod, ctx) |
| mod.set_input(**params) |
| return mod |
| |
| |
| GDB_INIT_TEMPLATE = """ |
| layout asm |
| target remote localhost:{gdb_port} |
| set $pc = UTVMInit |
| break UTVMDone |
| """ |
| |
| |
| def reset_gdbinit(): |
| if "server_port" not in DEV_CONFIG_A: |
| return |
| try: |
| gdb_init_dir = os.environ["MICRO_GDB_INIT_DIR"] |
| except KeyError: |
| return |
| with open(f"{gdb_init_dir}/.gdbinit", "w") as f: |
| gdb_port = DEV_CONFIG_A["server_port"] - 3333 |
| f.write(GDB_INIT_TEMPLATE.format(gdb_port=gdb_port)) |
| |
| |
| def test_alloc(): |
| """Test tensor allocation on the device.""" |
| if not tvm.runtime.enabled("micro_dev"): |
| return |
| shape = (1024,) |
| dtype = "float32" |
| with micro.Session(DEV_CONFIG_A): |
| ctx = tvm.micro_dev(0) |
| np_tensor = np.random.uniform(size=shape).astype(dtype) |
| micro_tensor = tvm.nd.array(np_tensor, ctx) |
| tvm.testing.assert_allclose(np_tensor, micro_tensor.asnumpy()) |
| |
| |
| def test_add(): |
| """Test a module which performs addition.""" |
| if not tvm.runtime.enabled("micro_dev"): |
| return |
| shape = (1024,) |
| dtype = "float32" |
| |
| reset_gdbinit() |
| |
| # Construct TVM expression. |
| tvm_shape = tvm.runtime.convert(shape) |
| A = te.placeholder(tvm_shape, name="A", dtype=dtype) |
| B = te.placeholder(tvm_shape, name="B", dtype=dtype) |
| C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name="C") |
| s = te.create_schedule(C.op) |
| |
| func_name = "fadd" |
| c_mod = tvm.build(s, [A, B, C], target="c", name=func_name) |
| |
| with micro.Session(DEV_CONFIG_A) as sess: |
| micro_mod = micro.create_micro_mod(c_mod, DEV_CONFIG_A) |
| micro_func = micro_mod[func_name] |
| ctx = tvm.micro_dev(0) |
| |
| a_np = np.random.uniform(size=shape).astype(dtype) |
| a = tvm.nd.array(a_np, ctx) |
| b_np = np.random.uniform(size=shape).astype(dtype) |
| b = tvm.nd.array(b_np, ctx) |
| c = tvm.nd.array(np.zeros(shape, dtype=dtype), ctx) |
| micro_func(a, b, c) |
| |
| # ensure inputs weren't corrupted |
| tvm.testing.assert_allclose(a.asnumpy(), a_np) |
| tvm.testing.assert_allclose(b.asnumpy(), b_np) |
| # ensure output is correct |
| tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy()) |
| |
| |
| def test_workspace_add(): |
| """Test a module which uses a workspace to compute an intermediate value.""" |
| if not tvm.runtime.enabled("micro_dev"): |
| return |
| shape = (1024,) |
| dtype = "float32" |
| |
| reset_gdbinit() |
| |
| # Construct TVM expression. |
| tvm_shape = tvm.runtime.convert(shape) |
| A = te.placeholder(tvm_shape, name="A", dtype=dtype) |
| B = te.placeholder(tvm_shape, name="B", dtype=dtype) |
| B = te.compute(A.shape, lambda *i: A(*i) + 1, name="B") |
| C = te.compute(A.shape, lambda *i: B(*i) + 1, name="C") |
| s = te.create_schedule(C.op) |
| |
| func_name = "fadd_two_workspace" |
| c_mod = tvm.build(s, [A, C], target="c", name=func_name) |
| |
| with micro.Session(DEV_CONFIG_A) as sess: |
| micro_mod = micro.create_micro_mod(c_mod, DEV_CONFIG_A) |
| micro_func = micro_mod[func_name] |
| ctx = tvm.micro_dev(0) |
| a_np = np.random.uniform(size=shape).astype(dtype) |
| a = tvm.nd.array(a_np, ctx) |
| c = tvm.nd.array(np.zeros(shape, dtype=dtype), ctx) |
| micro_func(a, c) |
| |
| # ensure input wasn't corrupted |
| tvm.testing.assert_allclose(a.asnumpy(), a_np) |
| # ensure output is correct |
| tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + 2.0) |
| |
| |
| def test_graph_runtime(): |
| """Test a program which uses the graph runtime.""" |
| if not tvm.runtime.enabled("micro_dev"): |
| return |
| shape = (1024,) |
| dtype = "float32" |
| |
| # Construct Relay program. |
| x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype)) |
| xx = relay.multiply(x, x) |
| z = relay.add(xx, relay.const(1.0)) |
| func = relay.Function([x], z) |
| |
| with micro.Session(DEV_CONFIG_A): |
| mod = relay_micro_build(func, DEV_CONFIG_A) |
| |
| x_in = np.random.uniform(size=shape[0]).astype(dtype) |
| mod.run(x=x_in) |
| result = mod.get_output(0).asnumpy() |
| |
| tvm.testing.assert_allclose(mod.get_input(0).asnumpy(), x_in) |
| tvm.testing.assert_allclose(result, x_in * x_in + 1.0) |
| |
| |
| def test_conv2d(): |
| if not tvm.runtime.enabled("micro_dev"): |
| return |
| |
| from tvm.relay import create_executor |
| from tvm.relay import transform |
| |
| dshape = (1, 4, 16, 16) |
| dtype = "int8" |
| func_name = "fused_nn_conv2d" |
| |
| reset_gdbinit() |
| |
| # Construct Relay program. |
| x = relay.var("x", shape=dshape, dtype=dtype) |
| conv_expr = relay.nn.conv2d(x, relay.var("w"), kernel_size=(3, 3), padding=(1, 1), channels=4) |
| func = relay.Function(relay.analysis.free_vars(conv_expr), conv_expr) |
| mod = tvm.IRModule.from_expr(func) |
| mod = transform.InferType()(mod) |
| |
| x_shape = list(map(lambda x: x.value, mod["main"].params[0].checked_type.shape)) |
| w_shape = list(map(lambda x: x.value, mod["main"].params[1].checked_type.shape)) |
| out_shape = list(map(lambda x: x.value, mod["main"].ret_type.shape)) |
| |
| with tvm.transform.PassContext(config={"tir.disable_vectorize": True}): |
| graph, c_mod, params = relay.build(mod, target="c") |
| |
| with micro.Session(DEV_CONFIG_A): |
| micro_mod = micro.create_micro_mod(c_mod, DEV_CONFIG_A) |
| candidate_func_name = func_name |
| for i in range(100): |
| try: |
| micro_func = micro_mod[candidate_func_name] |
| break |
| except tvm.TVMError as e: |
| candidate_func_name = f"{func_name}_{i}" |
| else: |
| assert False |
| ctx = tvm.micro_dev(0) |
| |
| x_data = tvm.nd.array(np.random.uniform(size=x_shape).astype(dtype), ctx) |
| w_data = tvm.nd.array(np.random.uniform(size=w_shape).astype(dtype), ctx) |
| result = tvm.nd.array(np.zeros(shape=out_shape, dtype=dtype), ctx) |
| micro_func(x_data, w_data, result) |
| |
| out_data = np.zeros(out_shape, dtype=dtype) |
| params = {"x": x_data.asnumpy(), "w": w_data.asnumpy()} |
| intrp = create_executor("debug") |
| expected_result = intrp.evaluate(mod["main"])(x_data, w_data) |
| |
| tvm.testing.assert_allclose(result.asnumpy(), expected_result.asnumpy()) |
| |
| |
| def test_interleave_sessions(): |
| """Test closing and reopening sessions.""" |
| if not tvm.runtime.enabled("micro_dev"): |
| return |
| shape = (1024,) |
| dtype = "float32" |
| |
| # Construct Relay add program. |
| x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype)) |
| ret = relay.add(x, relay.const(1.0)) |
| add_const_func = relay.Function([x], ret) |
| |
| sess_a = micro.Session(DEV_CONFIG_A) |
| sess_b = micro.Session(DEV_CONFIG_B) |
| with sess_a: |
| np_tensor_a = np.random.uniform(size=shape).astype(dtype) |
| micro_tensor_a = tvm.nd.array(np_tensor_a, tvm.micro_dev(0)) |
| with sess_b: |
| np_tensor_b = np.random.uniform(size=shape).astype(dtype) |
| micro_tensor_b = tvm.nd.array(np_tensor_b, tvm.micro_dev(0)) |
| with sess_a: |
| add_const_mod = relay_micro_build(add_const_func, DEV_CONFIG_A) |
| add_const_mod.run(x=micro_tensor_a) |
| add_result = add_const_mod.get_output(0).asnumpy() |
| tvm.testing.assert_allclose(add_result, np_tensor_a + 1.0) |
| with sess_b: |
| add_const_mod = relay_micro_build(add_const_func, DEV_CONFIG_B) |
| add_const_mod.run(x=micro_tensor_b) |
| add_result = add_const_mod.get_output(0).asnumpy() |
| tvm.testing.assert_allclose(add_result, np_tensor_b + 1.0) |
| |
| |
| def test_nested_sessions(): |
| """Test entering and exiting nested session contexts.""" |
| if not tvm.runtime.enabled("micro_dev"): |
| return |
| shape = (1024,) |
| dtype = "float32" |
| |
| # Construct Relay add program. |
| x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype)) |
| ret = relay.add(x, relay.const(1.0)) |
| add_const_func = relay.Function([x], ret) |
| |
| sess_a = micro.Session(DEV_CONFIG_A) |
| sess_b = micro.Session(DEV_CONFIG_B) |
| with sess_a: |
| np_tensor_a = np.random.uniform(size=shape).astype(dtype) |
| micro_tensor_a = tvm.nd.array(np_tensor_a, tvm.micro_dev(0)) |
| with sess_b: |
| np_tensor_b = np.random.uniform(size=shape).astype(dtype) |
| micro_tensor_b = tvm.nd.array(np_tensor_b, tvm.micro_dev(0)) |
| add_const_mod = relay_micro_build(add_const_func, DEV_CONFIG_A) |
| add_const_mod.run(x=micro_tensor_a) |
| add_result = add_const_mod.get_output(0).asnumpy() |
| tvm.testing.assert_allclose(add_result, np_tensor_a + 1.0) |
| |
| |
| def test_inactive_session_use(): |
| """Test the use of objects allocated in a session that is no longer active.""" |
| if not tvm.runtime.enabled("micro_dev"): |
| return |
| shape = (1024,) |
| dtype = "float32" |
| |
| # Construct Relay add program. |
| x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype)) |
| ret = relay.add(x, relay.const(1.0)) |
| add_const_func = relay.Function([x], ret) |
| |
| sess_a = micro.Session(DEV_CONFIG_A) |
| sess_b = micro.Session(DEV_CONFIG_B) |
| with sess_a: |
| np_tensor_a = np.random.uniform(size=shape).astype(dtype) |
| micro_tensor_a = tvm.nd.array(np_tensor_a, tvm.micro_dev(0)) |
| add_const_mod = relay_micro_build(add_const_func, DEV_CONFIG_A) |
| |
| with sess_b: |
| # These objects belong to `sess_a`. |
| add_const_mod.run(x=micro_tensor_a) |
| add_result = add_const_mod.get_output(0).asnumpy() |
| tvm.testing.assert_allclose(add_result, np_tensor_a + 1.0) |
| |
| |
| # TODO add workspace alloc/free stress test |
| |
| if __name__ == "__main__": |
| test_alloc() |
| print() |
| print("finished alloc test") |
| input("[press enter to continue]") |
| test_add() |
| print() |
| print("finished add test") |
| input("[press enter to continue]") |
| test_workspace_add() |
| print() |
| print("finished workspace add test") |
| input("[press enter to continue]") |
| test_graph_runtime() |
| print() |
| print("finished graph runtime test") |
| input("[press enter to continue]") |
| test_conv2d() |
| print() |
| print("finished conv2d test") |
| input("[press enter to continue]") |
| # disable for now as these are currently broken |
| # test_interleave_sessions() |
| # print() |
| # print('finished interleaved sessions test') |
| # input('[press enter to continue]') |
| # test_nested_sessions() |
| # print() |
| # print('finished nested sessions test') |
| # input('[press enter to continue]') |
| test_inactive_session_use() |
| print() |
| print("finished use inactive session test") |
| input("[press enter to continue]") |
