--- /dev/null
+# 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
+from tvm.relay.testing import resnet
+
+# 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 = 'c -device=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
+ """
+ disable_vectorize = tvm.target.build_config(disable_vectorize=True)
+ disable_fusion = relay.build_config(disabled_pass={'FuseOps'})
+ with disable_vectorize, disable_fusion:
+ 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.target.build_config(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]')
projects / platform / upstream / tvm.git / commitdiff