mlonmcu 0.3.0.dev0__3-none-any.whl

mlonmcu/flow/tvm/backend/backend.py

@@ -0,0 +1,282 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed 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 multiprocessing
+
+from mlonmcu.flow.backend import Backend
+from mlonmcu.setup import utils
+from mlonmcu.config import str2bool
+from .model_info import get_model_info, get_supported_formats
+from .python_utils import prepare_python_environment
+from .tvmc_utils import (
+    get_target_tvmc_args,
+    get_pass_config_tvmc_args,
+    get_disabled_pass_tvmc_args,
+    get_runtime_executor_tvmc_args,
+    get_input_shapes_tvmc_args,
+    get_tuning_records_tvmc_args,
+)
+
+
+class TVMBackend(Backend):
+    registry = {}
+
+    name = None
+
+    FEATURES = ["autotuned", "cmsisnnbyoc", "muriscvnnbyoc", "disable_legalize", "moiopt"]
+
+    DEFAULTS = {
+        "print_outputs": False,
+        "opt_level": 3,
+        "target_device": None,
+        "target_mcpu": None,
+        "target_march": None,
+        "target_model": None,
+        "target_mtriple": None,
+        "target_mabi": None,
+        "target_mattr": None,
+        "extra_target": None,
+        "extra_target_mcpu": None,
+        "desired_layout": None,  # optional: NCHW or NHWC
+        "disabled_passes": [],  # i.e. AlterOpLayout
+        "extra_pass_config": {},  # TODO: some example (fuse_max_depth etc.)
+        "use_tuning_results": False,
+        "tvmc_extra_args": [],  # Currently compile subcommand only!
+        "tvmc_custom_script": None,
+        # See https://github.com/apache/tvm/blob/1115fd9bc261619ffa0539746ae0aebc46232dc6/python/tvm/autotvm/tophub.py
+        "tophub_url": None,
+        "num_threads": multiprocessing.cpu_count(),
+    }
+
+    REQUIRED = []
+
+    OPTIONAL = ["tvm.build_dir", "tvm.pythonpath", "tvm.configs_dir", "tvm.use_tlcpack"]
+
+    def __init__(self, target="c", executor=None, runtime="crt", fmt="mlf", features=None, config=None):
+        super().__init__(framework="tvm", features=features, config=config)
+
+        self.model = None  # Actual filename!
+        self.model_info = None
+        self.input_shapes = None
+        self.supported_formats = get_supported_formats()
+        self.target = target
+        self.runtime = runtime
+        self.executor = executor
+        self.fmt = fmt
+
+        self.prefix = "default"
+        self.artifacts = (
+            []
+        )  # TODO: either make sure that ony one model is processed at a time or move the artifacts to the methods
+        self._tuning_records = None
+
+    @property
+    def tuning_records(self):
+        if self._tuning_records:
+            return self._tuning_records
+        elif "autotuning_results_file" in self.config and self.config["autotuning_results_file"]:
+            return self.config["autotuning_results_file"]
+        else:
+            return None
+
+    @tuning_records.setter
+    def tuning_records(self, filepath):
+        self._tuning_records = filepath
+
+    @property
+    def pass_config(self):
+        base = {"tir.disable_vectorize": True}
+        extra = self.config["extra_pass_config"]
+        if isinstance(extra, str):
+            import ast
+
+            extra = ast.literal_eval(extra)
+        assert isinstance(extra, dict)
+        base.update(extra)
+        return base
+
+    @property
+    def target_device(self):
+        return self.config["target_device"]
+
+    @property
+    def target_mcpu(self):
+        return self.config["target_mcpu"]
+
+    @property
+    def target_march(self):
+        return self.config["target_march"]
+
+    @property
+    def target_mtriple(self):
+        return self.config["target_mtriple"]
+
+    @property
+    def target_mabi(self):
+        return self.config["target_mabi"]
+
+    @property
+    def target_mattr(self):
+        return self.config["target_mattr"]
+
+    @property
+    def target_model(self):
+        return self.config["target_model"]
+
+    @property
+    def extra_target(self):
+        return self.config["extra_target"]
+
+    @property
+    def extra_target_mcpu(self):
+        return self.config["extra_target_mcpu"]
+
+    @property
+    def desired_layout(self):
+        return self.config["desired_layout"]
+
+    @property
+    def opt_level(self):
+        return self.config["opt_level"]
+
+    @property
+    def use_tuning_results(self):
+        value = self.config["use_tuning_results"]
+        return str2bool(value) if not isinstance(value, (bool, int)) else value
+
+    @property
+    def tvmc_extra_args(self):
+        return self.config["tvmc_extra_args"]
+
+    @property
+    def tvmc_custom_script(self):
+        return self.config["tvmc_custom_script"]
+
+    @property
+    def disabled_passes(self):
+        return self.config["disabled_passes"]
+
+    @property
+    def tvm_pythonpath(self):
+        return self.config["tvm.pythonpath"]
+
+    @property
+    def tvm_build_dir(self):
+        return self.config["tvm.build_dir"]
+
+    @property
+    def tvm_configs_dir(self):
+        return self.config["tvm.configs_dir"]
+
+    @property
+    def tophub_url(self):
+        return self.config["tophub_url"]
+
+    @property
+    def print_outputs(self):
+        value = self.config["print_outputs"]
+        return str2bool(value) if not isinstance(value, (bool, int)) else value
+
+    @property
+    def use_tlcpack(self):
+        value = self.config["tvm.use_tlcpack"]
+        return str2bool(value, allow_none=True) if not isinstance(value, (bool, int)) else value
+
+    def num_threads(self):
+        return self.config["num_threads"]
+
+    def get_target_details(self):
+        ret = {}
+        if self.target_device:
+            ret["device"] = self.target_device
+        if self.target_mcpu:
+            ret["mcpu"] = self.target_mcpu
+        if self.target_march:
+            ret["march"] = self.target_march
+        if self.target_mtriple:
+            ret["mtriple"] = self.target_mtriple
+        if self.target_mabi:
+            ret["mabi"] = self.target_mabi
+        if self.target_mattr:
+            ret["mattr"] = self.target_mattr
+        if self.target_model:
+            ret["model"] = self.target_model
+        return ret
+
+    def get_extra_target_details(self):
+        ret = {}
+        if self.extra_target_mcpu:
+            ret["mcpu"] = self.extra_target_mcpu
+        return ret
+
+    def get_tvmc_compile_args(self, out, dump=None):
+        assert self.executor is not None
+        assert self.executor in ["aot", "graph"], "Unsupported TVM executor"
+        args = [
+            self.model,
+            *get_target_tvmc_args(
+                self.target,
+                extra_target=self.extra_target,
+                target_details=self.get_target_details(),
+                extra_target_details=self.get_extra_target_details(),
+            ),
+            *get_runtime_executor_tvmc_args(self.runtime, self.executor),
+            *get_pass_config_tvmc_args(self.pass_config),
+            *get_disabled_pass_tvmc_args(self.disabled_passes),
+            *get_input_shapes_tvmc_args(self.input_shapes),
+            *get_tuning_records_tvmc_args(self.use_tuning_results, self.tuning_records),
+            *(["--desired-layout", self.desired_layout] if self.desired_layout is not None else []),
+            *(["--dump-code", ",".join(dump)] if dump is not None else []),
+            *self.tvmc_extra_args,
+            *["--opt-level", str(self.opt_level)],
+            *["--output", str(out)],
+            *["-f", self.fmt],
+            *["--model-format", self.model_format],
+        ]
+        return args
+
+    def invoke_tvmc(self, command, *args, cwd=None):
+        env = prepare_python_environment(
+            None if self.use_tlcpack else self.tvm_pythonpath,
+            None if self.use_tlcpack else self.tvm_build_dir,
+            None if self.use_tlcpack else self.tvm_configs_dir,
+            tophub_url=self.tophub_url,
+            num_threads=self.num_threads,
+        )
+        if self.use_tlcpack:
+            pre = ["tvmc"]
+            return utils.exec_getout(
+                *pre, command, *args, live=self.print_outputs, print_output=False, env=env, cwd=cwd
+            )
+        else:
+            if self.tvmc_custom_script is None:
+                pre = ["-m", "tvm.driver.tvmc"]
+            else:
+                pre = [self.tvmc_custom_script]
+            return utils.python(*pre, command, *args, live=self.print_outputs, print_output=False, env=env, cwd=cwd)
+
+    def invoke_tvmc_compile(self, out, dump=None, cwd=None):
+        args = self.get_tvmc_compile_args(out)
+        return self.invoke_tvmc("compile", *args, cwd=cwd)
+
+    def load_model(self, model):
+        self.model = model
+        # TODO: path model class instead of path!
+        # fmt = self.model.formats[0]
+        self.model_format, self.model_info = get_model_info(model, backend_name=self.name)
+        self.input_shapes = {tensor.name: tensor.shape for tensor in self.model_info.in_tensors}

mlonmcu/flow/tvm/backend/model_info.py

@@ -0,0 +1,307 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed 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 re
+import os
+
+from mlonmcu.models.model import ModelFormats
+
+
+class TensorInfo:
+    def __init__(self, name, shape, dtype, fix_names=False):
+        self.name = name
+        if fix_names:
+            self.name = self.name.replace("/", "_").replace(";", "_")
+        assert isinstance(shape, (tuple, list))
+        self.shape = shape
+        size_lookup = {
+            "float32": 4,
+            "uint8": 1,
+            "int8": 1,
+            "int32": 4,
+            "int64": 8,
+        }
+        assert dtype in size_lookup
+        self.dtype = dtype
+        self.type_size = size_lookup[self.dtype]
+
+    @property
+    def size(self):
+        ret = self.type_size
+        for dim in self.shape:
+            ret *= dim
+        return ret
+
+
+class TfLiteTensorInfo(TensorInfo):
+    def __init__(self, t, fix_names=False):
+        # Local imports to get rid of tflite dependency for non-tflite models
+        from tflite.TensorType import TensorType as TType
+
+        name = t.Name().decode()
+        shape = tuple([t.Shape(si) for si in range(0, t.ShapeLength())])
+
+        type_lookup = {
+            TType.FLOAT32: "float32",
+            TType.UINT8: "uint8",
+            TType.INT8: "int8",
+            TType.INT32: "int32",
+            TType.BOOL: "int8",
+            TType.INT64: "int64",
+        }
+        dtype = type_lookup[t.Type()]
+        super().__init__(name, shape, dtype, fix_names=fix_names)
+
+
+class RelayTensorInfo(TensorInfo):
+    def __init__(self, name, shape, dtype, fix_names=False):
+        super().__init__(name, shape, dtype, fix_names=fix_names)
+
+
+class ModelInfo:
+    def __init__(self, in_tensors, out_tensors, fix_names=False):
+        self.in_tensors = in_tensors
+        self.out_tensors = out_tensors
+
+
+class TfLiteModelInfo(ModelInfo):
+    def __init__(self, model, fix_names=False):
+        assert model.SubgraphsLength() == 1
+        g = model.Subgraphs(0)
+
+        in_tensors = []
+        for i in range(0, g.InputsLength()):
+            t = g.Tensors(g.Inputs(i))
+            in_tensors.append(TfLiteTensorInfo(t, fix_names=fix_names))
+
+        out_tensors = []
+        for i in range(0, g.OutputsLength()):
+            t = g.Tensors(g.Outputs(i))
+            out_tensors.append(TfLiteTensorInfo(t, fix_names=fix_names))
+        super().__init__(in_tensors, out_tensors)
+
+
+def shape_from_str(shape_str):
+    return tuple(map(int, shape_str.replace(" ", "").split(",")))
+
+
+def parse_relay_main(line):
+    input_tensors = []
+    output_tensors = []
+
+    input_tensors_strs = re.compile(r"%[a-zA-Z0-9_]+: Tensor\[\((?:\d+)(?:,\s*\d+)*\), (?:[a-zA-Z0-9_]+)\]").findall(
+        line
+    )
+    for input_tensors_str in input_tensors_strs:
+        res = re.compile(r"%([a-zA-Z0-9]+): Tensor\[\((\d+(?:, \d+)+)\), ([a-zA-Z0-9_]+)\]").match(input_tensors_str)
+        assert res is not None
+        groups = res.groups()
+        assert len(groups) == 3
+        input_name, input_shape_str, input_type = groups
+        input_shape = shape_from_str(input_shape_str)
+        input_tensor = TensorInfo(input_name, input_shape, input_type)
+        input_tensors.append(input_tensor)
+
+    output_tensor_names_str = re.compile(r"output_tensor_names=\[(\".*\")\]").findall(line)
+    output_tensor_names = re.compile(r"\"([a-zA-Z0-9_]+)\"").findall(output_tensor_names_str[0])
+
+    output_tensors_str = re.compile(r"-> (.+) {").findall(line)
+    output_tensor_strs = re.compile(r"Tensor\[\(\d+(?:, \d+)+\), [a-zA-Z0-9_]+\]").findall(output_tensors_str[0])
+
+    assert len(output_tensor_names) == len(output_tensor_strs)
+
+    for i, output_name in enumerate(output_tensor_names):
+        res = re.compile(r"Tensor\[\((\d+(?:, \d+)+)\), ([a-zA-Z0-9_]+)\]").match(output_tensor_strs[i])
+        assert res is not None
+        groups = res.groups()
+        assert len(groups) == 2
+        output_shape_str, output_type = groups
+        output_shape = shape_from_str(output_shape_str)
+        output_tensor = TensorInfo(output_name, output_shape, output_type)
+        output_tensors.append(output_tensor)
+    return input_tensors, output_tensors
+
+
+class RelayModelInfo(ModelInfo):
+    def __init__(self, mod_text, fix_names=False):
+        in_tensors = None
+        out_tensors = None
+        for line in mod_text.split("\n"):
+            if "def @main(" in line:
+                in_tensors, out_tensors = parse_relay_main(line)
+                break
+        assert in_tensors is not None and out_tensors is not None
+        super().__init__(in_tensors, out_tensors)
+
+
+def get_tfgraph_inout(graph):
+    ops = graph.get_operations()
+    outputs_set = set(ops)
+    inputs = []
+    for op in ops:
+        if op.type == "Placeholder":
+            inputs.append(op)
+        else:
+            for input_tensor in op.inputs:
+                if input_tensor.op in outputs_set:
+                    outputs_set.remove(input_tensor.op)
+    outputs = list(outputs_set)
+    return (inputs, outputs)
+
+
+class PBModelInfo(ModelInfo):
+    def __init__(self, model_file):
+        import tensorflow as tf
+
+        with tf.io.gfile.GFile(model_file, "rb") as f:
+            graph_def = tf.compat.v1.GraphDef()
+            graph_def.ParseFromString(f.read())
+        with tf.Graph().as_default() as graph:
+            tf.import_graph_def(graph_def)
+            inputs, outputs = get_tfgraph_inout(graph)
+        in_tensors = [TensorInfo(t.name, t.shape.as_list(), t.dtype.name) for op in inputs for t in op.outputs]
+        out_tensors = [TensorInfo(t.name, t.shape.as_list(), t.dtype.name) for op in outputs for t in op.outputs]
+        super().__init__(in_tensors, out_tensors)
+
+
+class ONNXModelInfo(ModelInfo):
+    def __init__(self, model_file):
+        from google.protobuf.json_format import MessageToDict
+        import onnx
+        from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
+
+        model = onnx.load(model_file)
+
+        def _helper(tensors):
+            ret = []
+            for tensor in tensors:
+                d = MessageToDict(tensor)
+                name = d["name"]
+                tensor_type = d["type"]["tensorType"]
+                elem_type = tensor_type["elemType"]
+                dims = tensor_type["shape"]["dim"]
+                shape = [int(x["dimValue"]) if "dimValue" in x else 40 for x in dims]  # TODO: dyn shape
+                dtype = str(TENSOR_TYPE_TO_NP_TYPE[elem_type])
+                ret.append(TensorInfo(name, shape, dtype))
+            return ret
+
+        in_tensors = _helper(model.graph.input)
+        out_tensors = _helper(model.graph.output)
+
+        # TVM seems to ignore the original output names for ONNX models
+        for i, t in enumerate(out_tensors):
+            t.name = f"output{i}"
+        super().__init__(in_tensors, out_tensors)
+
+
+class PaddleModelInfo(ModelInfo):
+    def __init__(self, model_file):
+        import paddle
+
+        paddle.enable_static()
+        paddle.disable_signal_handler()
+
+        exe = paddle.static.Executor(paddle.CPUPlace())
+        prefix = "".join(model_file.strip().split(".")[:-1])
+        program, _, _ = paddle.static.load_inference_model(prefix, exe)
+
+        output_names = list()
+        for block in program.blocks:
+            for op in block.ops:
+                if op.type == "fetch":
+                    output_names.append(op.input("X")[0])
+
+        def _convert_dtype_value(val):
+            """Converts a Paddle type id to a string."""
+            # See: https://github.com/apache/tvm/blob/cc769fdc951707b8be991949864817b955a4dbc7/python/tvm/
+            # relay/frontend/paddlepaddle.py#L70
+
+            convert_dtype_map = {
+                21: "int8",
+                20: "uint8",
+                6: "float64",
+                5: "float32",
+                4: "float16",
+                3: "int64",
+                2: "int32",
+                1: "int16",
+                0: "bool",
+            }
+            if val not in convert_dtype_map:
+                msg = "Paddle data type value %d is not handled yet." % (val)
+                raise NotImplementedError(msg)
+            return convert_dtype_map[val]
+
+        raise NotImplementedError
+        # super().__init__(in_tensors, out_tensors)
+
+
+def get_tflite_model_info(model_buf):
+    # Local imports to get rid of tflite dependency for non-tflite models
+    import tflite
+
+    tflite_model = tflite.Model.GetRootAsModel(model_buf, 0)
+    model_info = TfLiteModelInfo(tflite_model)
+    return model_info
+
+
+def get_relay_model_info(mod_text):
+    model_info = RelayModelInfo(mod_text)
+    return model_info
+
+
+def get_pb_model_info(model_file):
+    model_info = PBModelInfo(model_file)
+    return model_info
+
+
+def get_paddle_model_info(model_file):
+    model_info = PaddleModelInfo(model_file)
+    return model_info
+
+
+def get_onnx_model_info(model_file):
+    model_info = ONNXModelInfo(model_file)
+    return model_info
+
+
+def get_model_info(model, backend_name="unknown"):
+    ext = os.path.splitext(model)[1][1:]
+    fmt = ModelFormats.from_extension(ext)
+    if fmt == ModelFormats.TFLITE:
+        with open(model, "rb") as handle:
+            model_buf = handle.read()
+            return "tflite", get_tflite_model_info(model_buf)
+    elif fmt == ModelFormats.RELAY:
+        # Warning: the wrapper generateion does currently not work because of the
+        # missing possibility to get the relay models input names and shapes
+        with open(model, "r") as handle:
+            mod_text = handle.read()
+        return "relay", get_relay_model_info(mod_text)
+    elif fmt == ModelFormats.PB:
+        return "pb", get_pb_model_info(model)
+    elif fmt == ModelFormats.ONNX:
+        return "onnx", get_onnx_model_info(model)
+    elif fmt == ModelFormats.PADDLE:
+        return "pdmodel", get_pb_model_info(model)
+    else:
+        raise RuntimeError(f"Unsupported model format '{fmt.name}' for backend '{backend_name}'")
+
+
+def get_supported_formats():
+    return [ModelFormats.TFLITE, ModelFormats.RELAY, ModelFormats.PB, ModelFormats.ONNX, ModelFormats.PADDLE]

mlonmcu/flow/tvm/backend/python_utils.py

@@ -0,0 +1,36 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed 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
+
+
+def prepare_python_environment(pythonpath, tvm_build_dir, tvm_configs_dir, tophub_url=None, num_threads=None):
+    env = os.environ.copy()
+    if pythonpath:
+        env["PYTHONPATH"] = str(pythonpath)
+    if tvm_build_dir:
+        env["TVM_LIBRARY_PATH"] = str(tvm_build_dir)
+    if tvm_configs_dir:
+        env["TVM_CONFIGS_JSON_DIR"] = str(tvm_configs_dir)
+    if tophub_url:
+        env["TOPHUB_LOCATION"] = tophub_url
+    if num_threads:
+        env["TVM_NUM_THREADS"] = str(num_threads)
+    # Use all cores/threads for building models and let OS scheduler decide on the mapping
+    env["TVM_BIND_THREADS"] = str(0)
+    return env

mlonmcu/flow/tvm/backend/tuner.py

@@ -0,0 +1,41 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed 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.
+#
+
+
+class TVMTuner:
+    DEFAULTS = {
+        "enable": False,
+        "print_outputs": False,
+        "results_file": None,
+        "append": None,
+        "tuner": "ga",  # Options: ga,gridsearch,random,xgb,xgb_knob,xgb-rank
+        "trials": 10,  # TODO: increase to 100?
+        "early_stopping": None,  # calculate default dynamically
+        "num_workers": 1,
+        "max_parallel": 1,
+        "use_rpc": False,
+        "timeout": 100,
+        "mode": "autotvm",  # Options: autotvm, auto_scheduler
+        "visualize": False,
+        "tasks": None,
+    }
+
+    def __init__(self, backend, config=None):
+        self.backend = backend
+        self.config = config if config is not None else {}