Andrew Reusch
251 lines
8.4 KiB
Python
251 lines
8.4 KiB
Python
# 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.
|
|
|
|
"""Defines a top-level glue class that operates the Transport and Flasher classes."""
|
|
|
|
import logging
|
|
import sys
|
|
|
|
from ..error import register_error
|
|
from .._ffi import get_global_func
|
|
from ..contrib import graph_executor
|
|
from ..contrib.debugger import debug_executor
|
|
from ..rpc import RPCSession
|
|
from .transport import IoTimeoutError
|
|
from .transport import TransportLogger
|
|
|
|
try:
|
|
from .base import _rpc_connect
|
|
except ImportError:
|
|
raise ImportError("micro tvm is not enabled. Set USE_MICRO to ON in config.cmake")
|
|
|
|
|
|
@register_error
|
|
class SessionTerminatedError(Exception):
|
|
"""Raised when a transport read operationd discovers that the remote session is terminated."""
|
|
|
|
|
|
class Session:
|
|
"""MicroTVM Device Session
|
|
|
|
Parameters
|
|
----------
|
|
config : dict
|
|
configuration for this session (as generated by
|
|
`tvm.micro.device.host.default_config()`, for example)
|
|
|
|
Example
|
|
--------
|
|
.. code-block:: python
|
|
|
|
c_mod = ... # some module generated with "c" as the target
|
|
dev_config = micro.device.arm.stm32f746xx.default_config('127.0.0.1', 6666)
|
|
with tvm.micro.Session(dev_config) as sess:
|
|
micro_mod = sess.create_micro_mod(c_mod)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
binary=None,
|
|
flasher=None,
|
|
transport_context_manager=None,
|
|
session_name="micro-rpc",
|
|
timeout_override=None,
|
|
):
|
|
"""Configure a new session.
|
|
|
|
Parameters
|
|
----------
|
|
binary : MicroBinary
|
|
If given, `flasher` must also be given. During session initialization, this binary will
|
|
be flashed to the device before the transport is created.
|
|
flasher : Flasher
|
|
If given, `binary` must also be given. Used to flash `binary` during session
|
|
initialization.
|
|
transport_context_manager : ContextManager[transport.Transport]
|
|
If given, `flasher` and `binary` should not be given. On entry, this context manager
|
|
should establish a tarnsport between this TVM instance and the device.
|
|
session_name : str
|
|
Name of the session, used for debugging.
|
|
timeout_override : TransportTimeouts
|
|
If given, TransportTimeouts that govern the way Receive() behaves. If not given, this is
|
|
determined by calling has_flow_control() on the transport.
|
|
"""
|
|
self.binary = binary
|
|
self.flasher = flasher
|
|
self.transport_context_manager = transport_context_manager
|
|
self.session_name = session_name
|
|
self.timeout_override = timeout_override
|
|
|
|
self._rpc = None
|
|
self._graph_executor = None
|
|
|
|
def get_system_lib(self):
|
|
return self._rpc.get_function("runtime.SystemLib")()
|
|
|
|
def _wrap_transport_read(self, n, timeout_microsec):
|
|
try:
|
|
return self.transport.read(
|
|
n, float(timeout_microsec) / 1e6 if timeout_microsec is not None else None
|
|
)
|
|
except IoTimeoutError:
|
|
return bytes([])
|
|
|
|
def _wrap_transport_write(self, data, timeout_microsec):
|
|
try:
|
|
return self.transport.write(
|
|
data, float(timeout_microsec) / 1e6 if timeout_microsec is not None else None
|
|
)
|
|
except IoTimeoutError:
|
|
return 0
|
|
|
|
def __enter__(self):
|
|
"""Initialize this session and establish an RPC session with the on-device RPC server.
|
|
|
|
Returns
|
|
-------
|
|
Session :
|
|
Returns self.
|
|
"""
|
|
if self.flasher is not None:
|
|
self.transport_context_manager = self.flasher.flash(self.binary)
|
|
|
|
self.transport = TransportLogger(
|
|
self.session_name, self.transport_context_manager, level=logging.DEBUG
|
|
).__enter__()
|
|
|
|
try:
|
|
timeouts = self.timeout_override
|
|
if timeouts is None:
|
|
timeouts = self.transport.timeouts()
|
|
|
|
self._rpc = RPCSession(
|
|
_rpc_connect(
|
|
self.session_name,
|
|
self._wrap_transport_write,
|
|
self._wrap_transport_read,
|
|
int(timeouts.session_start_retry_timeout_sec * 1e6),
|
|
int(timeouts.session_start_timeout_sec * 1e6),
|
|
int(timeouts.session_established_timeout_sec * 1e6),
|
|
)
|
|
)
|
|
self.device = self._rpc.cpu(0)
|
|
return self
|
|
|
|
except:
|
|
self.transport.__exit__(*sys.exc_info())
|
|
raise
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_traceback):
|
|
"""Tear down this session and associated RPC session resources."""
|
|
self.transport.__exit__(exc_type, exc_value, exc_traceback)
|
|
|
|
|
|
def lookup_remote_linked_param(mod, storage_id, template_tensor, device):
|
|
"""Lookup a parameter that has been pre-linked into a remote (i.e. over RPC) Module.
|
|
|
|
This function signature matches the signature built by
|
|
|
|
Parameters
|
|
----------
|
|
mod : tvm.runtime.Module
|
|
The remote Module containing the pre-linked parameters.
|
|
storage_id : int
|
|
An integer identifying the pre-linked paramter to find
|
|
template_tensor : DLTensor
|
|
A DLTensor containing metadata that should be filled-in to the returned NDArray. This
|
|
function should mostly not inspect this, and just pass it along to
|
|
NDArrayFromRemoteOpaqueHandle.
|
|
device : Device
|
|
The remote CPU device to be used with the returned NDArray.
|
|
|
|
Returns
|
|
-------
|
|
tvm.nd.NDArray :
|
|
NDArray containing the pre-linked parameter.
|
|
"""
|
|
try:
|
|
lookup_linked_param = mod.get_function("_lookup_linked_param")
|
|
except AttributeError:
|
|
return None
|
|
|
|
remote_data = lookup_linked_param(storage_id)
|
|
if remote_data is None:
|
|
return None
|
|
|
|
return get_global_func("tvm.rpc.NDArrayFromRemoteOpaqueHandle")(
|
|
mod, remote_data, template_tensor, device, None
|
|
)
|
|
|
|
|
|
def create_local_graph_executor(graph_json_str, mod, device):
|
|
"""Create a local graph executor driving execution on the remote CPU device given.
|
|
|
|
Parameters
|
|
----------
|
|
graph_json_str : str
|
|
A string containing the graph representation.
|
|
|
|
mod : tvm.runtime.Module
|
|
The remote module containing functions in graph_json_str.
|
|
|
|
device : tvm.runtime.Device
|
|
The remote CPU execution device.
|
|
|
|
Returns
|
|
-------
|
|
tvm.contrib.GraphExecutor :
|
|
A local graph executor instance that executes on the remote device.
|
|
"""
|
|
device_type_id = [device.device_type, device.device_id]
|
|
fcreate = get_global_func("tvm.graph_executor.create")
|
|
return graph_executor.GraphModule(
|
|
fcreate(graph_json_str, mod, lookup_remote_linked_param, *device_type_id)
|
|
)
|
|
|
|
|
|
def create_local_debug_executor(graph_json_str, mod, device, dump_root=None):
|
|
"""Create a local debug runtime driving execution on the remote CPU device given.
|
|
|
|
Parameters
|
|
----------
|
|
graph_json_str : str
|
|
A string containing the graph representation.
|
|
|
|
mod : tvm.runtime.Module
|
|
The remote module containing functions in graph_json_str.
|
|
|
|
device : tvm.runtime.Device
|
|
The remote CPU execution device.
|
|
|
|
dump_root : Optional[str]
|
|
If given, passed as dump_root= to GraphModuleDebug.
|
|
|
|
Returns
|
|
-------
|
|
tvm.contrib.GraphExecutor :
|
|
A local graph executor instance that executes on the remote device.
|
|
"""
|
|
device_type_id = [device.device_type, device.device_id]
|
|
fcreate = get_global_func("tvm.graph_executor_debug.create")
|
|
return debug_executor.GraphModuleDebug(
|
|
fcreate(graph_json_str, mod, lookup_remote_linked_param, *device_type_id),
|
|
[device],
|
|
graph_json_str,
|
|
dump_root=dump_root,
|
|
)
|