mlonmcu 0.3.0.dev0__3-none-any.whl

mlonmcu/setup/tasks/spike.py

@@ -0,0 +1,190 @@
+#
+# 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.
+#
+"""Definition of tasks used to dynamically install MLonMCU dependencies"""
+
+import os
+import multiprocessing
+from pathlib import Path
+
+from mlonmcu.setup.task import TaskType
+from mlonmcu.context.context import MlonMcuContext
+from mlonmcu.setup import utils
+from mlonmcu.logging import get_logger
+
+from .common import get_task_factory
+
+logger = get_logger()
+
+Tasks = get_task_factory()
+
+
+def _validate_spike(context: MlonMcuContext, params=None):
+    if not context.environment.has_target("spike"):
+        return False
+    if params.get("vext", False):
+        if params.get("pext", False):
+            return False  # Can not use booth at a time
+        if not context.environment.supports_feature("vext"):
+            return False
+    if params.get("pext", False):
+        if params.get("vext", False):
+            return False  # Can not use booth at a time
+        if not context.environment.supports_feature("pext"):
+            return False
+    user_vars = context.environment.vars
+    if "spike.pk" not in user_vars:  # TODO: also check command line flags?
+        assert "spikepk" in context.environment.repos, "Undefined repository: 'spikepk'"
+    if "spike.exe" not in user_vars:  # TODO: also check command line flags?
+        assert "spike" in context.environment.repos, "Undefined repository: 'spike'"
+    return True
+
+
+@Tasks.provides(["spikepk.src_dir"])
+@Tasks.validate(_validate_spike)
+@Tasks.register(category=TaskType.TARGET)
+def clone_spike_pk(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Clone the spike proxt kernel."""
+    spikepkName = utils.makeDirName("spikepk")
+    spikepkSrcDir = context.environment.paths["deps"].path / "src" / spikepkName
+    user_vars = context.environment.vars
+    if "spike.pk" in user_vars:  # TODO: also check command line flags?
+        return False
+    if rebuild or not utils.is_populated(spikepkSrcDir):
+        spikepkRepo = context.environment.repos["spikepk"]
+        utils.clone(spikepkRepo.url, spikepkSrcDir, branch=spikepkRepo.ref)
+    context.cache["spikepk.src_dir"] = spikepkSrcDir
+
+
+@Tasks.needs(["spikepk.src_dir", "riscv_gcc.install_dir", "riscv_gcc.name"])
+@Tasks.provides(["spikepk.build_dir", "spike.pk"])
+@Tasks.validate(_validate_spike)
+@Tasks.register(category=TaskType.TARGET)
+def build_spike_pk(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Build Spike proxy kernel."""
+    if not params:
+        params = {}
+    user_vars = context.environment.vars
+    if "spike.pk" in user_vars:  # TODO: also check command line flags?
+        return False
+    spikepkName = utils.makeDirName("spikepk")
+    spikepkSrcDir = context.cache["spikepk.src_dir"]
+    spikepkBuildDir = context.environment.paths["deps"].path / "build" / spikepkName
+    spikepkInstallDir = context.environment.paths["deps"].path / "install" / spikepkName
+    spikepkBin = spikepkInstallDir / "pk"
+    if rebuild or not (utils.is_populated(spikepkBuildDir) and spikepkBin.is_file()):
+        # No need to build a vext and non-vext variant?
+        utils.mkdirs(spikepkBuildDir)
+        gccName = context.cache["riscv_gcc.name"]
+        # assert gccName == "riscv32-unknown-elf", "Spike PK requires a non-multilib toolchain!"
+        vext = params.get("vext", False)
+        pext = params.get("pext", False)
+        assert not (pext and vext), "Currently only p or vector extension can be enabled at a time."
+        if vext and "riscv_gcc.install_dir_vext" in user_vars:
+            riscv_gcc = user_vars["riscv_gcc.install_dir_vext"]
+        elif pext and "riscv_gcc.install_dir_pext" in user_vars:
+            riscv_gcc = user_vars["riscv_gcc.install_dir_pext"]
+        elif "riscv_gcc.install_dir" in user_vars:
+            riscv_gcc = user_vars["riscv_gcc.install_dir"]
+        else:
+            riscv_gcc = context.cache["riscv_gcc.install_dir"]
+        arch = "rv32imafdc"
+        spikepkArgs = []
+        spikepkArgs.append("--prefix=" + str(riscv_gcc))
+        spikepkArgs.append("--host=" + gccName)
+        spikepkArgs.append(f"--with-arch={arch}")
+        spikepkArgs.append("--with-abi=ilp32d")
+        env = os.environ.copy()
+        env["PATH"] = str(Path(riscv_gcc) / "bin") + ":" + env["PATH"]
+        utils.exec_getout(
+            str(spikepkSrcDir / "configure"),
+            *spikepkArgs,
+            cwd=spikepkBuildDir,
+            env=env,
+            live=False,
+            print_output=False,
+        )
+        utils.make(cwd=spikepkBuildDir, threads=threads, live=verbose, env=env)
+        # utils.make(target="install", cwd=spikepkBuildDir, live=verbose, env=env)
+        utils.mkdirs(spikepkInstallDir)
+        utils.move(spikepkBuildDir / "pk", spikepkBin)
+    context.cache["spikepk.build_dir"] = spikepkBuildDir
+    context.cache["spike.pk"] = spikepkBin
+
+
+@Tasks.provides(["spike.src_dir"])
+@Tasks.validate(_validate_spike)
+@Tasks.register(category=TaskType.TARGET)
+def clone_spike(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Clone the spike simulator."""
+    spikeName = utils.makeDirName("spike")
+    spikeSrcDir = context.environment.paths["deps"].path / "src" / spikeName
+    user_vars = context.environment.vars
+    if "spike.exe" in user_vars:  # TODO: also check command line flags?
+        return False
+    if rebuild or not utils.is_populated(spikeSrcDir):
+        spikeRepo = context.environment.repos["spike"]
+        utils.clone(spikeRepo.url, spikeSrcDir, branch=spikeRepo.ref)
+    context.cache["spike.src_dir"] = spikeSrcDir
+
+
+@Tasks.needs(["spike.src_dir", "riscv_gcc.install_dir", "riscv_gcc.name"])
+@Tasks.provides(["spike.build_dir", "spike.exe"])
+@Tasks.validate(_validate_spike)
+@Tasks.register(category=TaskType.TARGET)
+def build_spike(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Build Spike simulator."""
+    if not params:
+        params = {}
+    user_vars = context.environment.vars
+    if "spike.exe" in user_vars:  # TODO: also check command line flags?
+        return False
+    spikeName = utils.makeDirName("spike")
+    spikeSrcDir = context.cache["spike.src_dir"]
+    spikeBuildDir = context.environment.paths["deps"].path / "build" / spikeName
+    spikeInstallDir = context.environment.paths["deps"].path / "install" / spikeName
+    spikeExe = spikeInstallDir / "spike"
+    user_vars = context.environment.vars
+    if "spike.exe" in user_vars:  # TODO: also check command line flags?
+        return False
+    if rebuild or not (utils.is_populated(spikeBuildDir) and spikeExe.is_file()):
+        # No need to build a vext and non-vext variant?
+        utils.mkdirs(spikeBuildDir)
+        spikeArgs = []
+        spikeArgs.append("--prefix=" + str(context.cache["riscv_gcc.install_dir"]))
+        utils.exec_getout(
+            str(spikeSrcDir / "configure"),
+            *spikeArgs,
+            cwd=spikeBuildDir,
+            live=False,
+            print_output=False,
+        )
+        utils.make(cwd=spikeBuildDir, threads=threads, live=verbose)
+        # utils.make(target="install", cwd=spikeBuildDir, threads=threads, live=verbose)
+        utils.mkdirs(spikeInstallDir)
+        utils.move(spikeBuildDir / "spike", spikeExe)
+    context.cache["spike.build_dir"] = spikeBuildDir
+    context.cache["spike.exe"] = spikeExe

mlonmcu/setup/tasks/tf.py

@@ -0,0 +1,92 @@
+#
+# 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.
+#
+"""Definition of tasks used to dynamically install MLonMCU dependencies"""
+
+import multiprocessing
+from pathlib import Path
+
+from mlonmcu.setup.task import TaskType
+from mlonmcu.context.context import MlonMcuContext
+from mlonmcu.setup import utils
+from mlonmcu.logging import get_logger
+
+from .common import get_task_factory
+
+logger = get_logger()
+
+Tasks = get_task_factory()
+
+
+def _validate_tensorflow(context: MlonMcuContext, params=None):
+    return context.environment.has_framework("tflm")
+
+
+@Tasks.provides(["tf.src_dir"])
+@Tasks.validate(_validate_tensorflow)
+@Tasks.register(category=TaskType.FRAMEWORK)
+def clone_tensorflow(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Clone the TF/TFLM repository."""
+    tfName = utils.makeDirName("tf")
+    tfSrcDir = context.environment.paths["deps"].path / "src" / tfName
+    if rebuild or not utils.is_populated(tfSrcDir):
+        tfRepo = context.environment.repos["tensorflow"]
+        utils.clone(tfRepo.url, tfSrcDir, branch=tfRepo.ref, refresh=rebuild)
+    context.cache["tf.src_dir"] = tfSrcDir
+
+
+@Tasks.needs(["tf.src_dir"])
+@Tasks.provides(["tf.dl_dir", "tf.lib_path"])
+# @Tasks.param("dbg", False)
+@Tasks.param("dbg", True)
+@Tasks.validate(_validate_tensorflow)
+@Tasks.register(category=TaskType.FRAMEWORK)
+def build_tensorflow(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Download tensorflow dependencies and build lib."""
+    if not params:
+        params = {}
+    flags = utils.makeFlags((params["dbg"], "dbg"))
+    # tfName = utils.makeDirName("tf", flags=flags)
+    tfSrcDir = context.cache["tf.src_dir"]
+    tflmDir = Path(tfSrcDir) / "tensorflow" / "lite" / "micro"
+    tflmBuildDir = tflmDir / "tools" / "make"
+    tflmDownloadsDir = tflmBuildDir / "downloads"
+    if params["dbg"]:
+        tflmLib = (
+            tflmBuildDir / "gen" / "linux_x86_64" / "lib" / "libtensorflow-microlite.a"
+        )  # FIXME: add _dbg suffix is possible
+    else:
+        tflmLib = tflmBuildDir / "gen" / "linux_x86_64" / "lib" / "libtensorflow-microlite.a"
+    # if rebuild or not tflmLib.is_file() or not utils.is_populated(tflmDownloadsDir):
+    if rebuild or not utils.is_populated(tflmDownloadsDir):
+        tfDbgArg = ["BUILD_TYPE=debug"] if params["dbg"] else []
+        utils.make(
+            "-f",
+            str(tflmDir / "tools" / "make" / "Makefile"),
+            "third_party_downloads",
+            threads=threads,
+            *tfDbgArg,
+            cwd=tfSrcDir,
+            live=verbose,
+        )
+    context.cache["tf.dl_dir"] = tflmDownloadsDir
+    context.cache["tf.lib_path", flags] = tflmLib  # ignore!

mlonmcu/setup/tasks/tflite.py

@@ -0,0 +1,59 @@
+#
+# 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.
+#
+"""Definition of tasks used to dynamically install MLonMCU dependencies"""
+
+import multiprocessing
+
+from mlonmcu.setup.task import TaskType
+from mlonmcu.context.context import MlonMcuContext
+from mlonmcu.setup import utils
+from mlonmcu.logging import get_logger
+
+from .common import get_task_factory
+
+logger = get_logger()
+
+Tasks = get_task_factory()
+
+
+def _validate_tflite_visualize(context: MlonMcuContext, params=None):
+    return context.environment.has_frontend("tflite") and context.environment.has_feature("visualize")
+
+
+@Tasks.provides(["tflite_visualize.exe"])
+@Tasks.validate(_validate_tflite_visualize)
+@Tasks.register(category=TaskType.FEATURE)
+def download_tflite_vizualize(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Download the visualize.py script for TFLite."""
+    # This script is content of the tensorflow repo (not tflite-micro) and unfortunately not bundled
+    # into the tensorflow python package. Therefore just download this single file form GitHub
+
+    tfLiteVizualizeName = utils.makeDirName("tflite_visualize")
+    tfLiteVizualizeInstallDir = context.environment.paths["deps"].path / "install" / tfLiteVizualizeName
+    tfLiteVizualizeExe = tfLiteVizualizeInstallDir / "visualize.py"
+    user_vars = context.environment.vars
+    if "tflite_visualize.exe" in user_vars:  # TODO: also check command line flags?
+        return False
+    if rebuild or not utils.is_populated(tfLiteVizualizeInstallDir):
+        tfLiteVizualizeInstallDir.mkdir()
+        url = "https://raw.githubusercontent.com/tensorflow/tensorflow/master/tensorflow/lite/tools/visualize.py"
+        utils.download(url, tfLiteVizualizeExe, progress=verbose)
+    context.cache["tflite_visualize.exe"] = tfLiteVizualizeExe

mlonmcu/setup/tasks/tflmc.py

@@ -0,0 +1,141 @@
+#
+# 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.
+#
+"""Definition of tasks used to dynamically install MLonMCU dependencies"""
+
+import multiprocessing
+from pathlib import Path
+
+from mlonmcu.setup.task import TaskType
+from mlonmcu.context.context import MlonMcuContext
+from mlonmcu.setup import utils
+from mlonmcu.logging import get_logger
+
+from .tf import _validate_tensorflow
+from .common import get_task_factory
+
+logger = get_logger()
+
+Tasks = get_task_factory()
+
+
+def _validate_tflite_micro_compiler(context: MlonMcuContext, params=None):
+    if not _validate_tensorflow(context, params=params):
+        return False
+    if not context.environment.has_backend("tflmc"):
+        return False
+    return True
+
+
+@Tasks.provides(["tflmc.src_dir"])
+@Tasks.validate(_validate_tflite_micro_compiler)
+@Tasks.register(category=TaskType.BACKEND)
+def clone_tflite_micro_compiler(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Clone the preinterpreter repository."""
+    tflmcName = utils.makeDirName("tflmc")
+    tflmcSrcDir = context.environment.paths["deps"].path / "src" / tflmcName
+    if rebuild or not utils.is_populated(tflmcSrcDir):
+        tflmcRepo = context.environment.repos["tflite_micro_compiler"]
+        utils.clone(tflmcRepo.url, tflmcSrcDir, branch=tflmcRepo.ref)
+    context.cache["tflmc.src_dir"] = tflmcSrcDir
+
+
+def _validate_build_tflite_micro_compiler(context: MlonMcuContext, params=None):
+    if params:
+        muriscvnn = params.get("muriscvnn", False)
+        cmsisnn = params.get("cmsisnn", False)
+        if muriscvnn and cmsisnn:
+            # Not allowed
+            return False
+        elif muriscvnn:
+            if not context.environment.supports_feature("muriscvnn"):
+                return False
+        elif cmsisnn:
+            if not context.environment.supports_feature("cmsisnn"):
+                return False
+    return _validate_tflite_micro_compiler(context, params=params)
+
+
+@Tasks.needs(["tflmc.src_dir", "tf.src_dir"])
+@Tasks.optional(["muriscvnn.lib", "muriscvnn.inc_dir", "cmsisnn.dir"])
+@Tasks.provides(["tflmc.build_dir", "tflmc.exe"])
+@Tasks.param("muriscvnn", [False, True])
+@Tasks.param("cmsisnn", [False, True])
+@Tasks.param("dbg", [False, True])
+@Tasks.param("arch", ["x86"])  # TODO: compile for arm/riscv in the future
+@Tasks.validate(_validate_build_tflite_micro_compiler)
+@Tasks.register(category=TaskType.BACKEND)
+def build_tflite_micro_compiler(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Build the TFLM preinterpreter."""
+    muriscvnn = params.get("muriscvnn", False)
+    cmsisnn = params.get("cmsisnn", False)
+    dbg = params.get("dbg", False)
+    arch = params.get("arch", "x86")
+    flags = utils.makeFlags((True, arch), (muriscvnn, "muriscvnn"), (cmsisnn, "cmsisnn"), (dbg, "dbg"))
+    # flags_ = utils.makeFlags((dbg, "dbg"))
+    tflmcName = utils.makeDirName("tflmc", flags=flags)
+    tflmcBuildDir = context.environment.paths["deps"].path / "build" / tflmcName
+    tflmcInstallDir = context.environment.paths["deps"].path / "install" / tflmcName
+    tflmcExe = tflmcInstallDir / "compiler"
+    tfSrcDir = context.cache["tf.src_dir"]
+    tflmcSrcDir = context.cache["tflmc.src_dir"]
+    if rebuild or not utils.is_populated(tflmcBuildDir) or not tflmcExe.is_file():
+        cmakeArgs = [
+            "-DTF_SRC=" + str(tfSrcDir),
+            "-DGET_TF_SRC=ON",
+        ]
+        if muriscvnn:
+            flags__ = utils.makeFlags((True, arch), (True, "gcc"), (dbg, "dbg"))
+            muriscvnnLib = context.cache["muriscvnn.lib", flags__]
+            muriscvnnInc = context.cache["muriscvnn.inc_dir"]
+            cmakeArgs.append("-DTFLM_OPTIMIZED_KERNEL=cmsis_nn")
+            cmakeArgs.append(f"-DTFLM_OPTIMIZED_KERNEL_LIB={muriscvnnLib}")
+            cmakeArgs.append(f"-DTFLM_OPTIMIZED_KERNEL_INCLUDE_DIR={muriscvnnInc}")
+        elif cmsisnn:
+            flags__ = utils.makeFlags((True, arch), (dbg, "dbg"))
+            cmsisnnLib = context.cache["cmsisnn.lib", flags__]
+            cmsisDir = Path(context.cache["cmsisnn.dir"])
+            cmsisIncs = r"\;".join(
+                [
+                    str(cmsisDir),
+                    str(cmsisDir / "CMSIS" / "Core" / "Include"),
+                    str(cmsisDir / "CMSIS" / "NN" / "Include"),
+                    str(cmsisDir / "CMSIS" / "DSP" / "Include"),
+                ]
+            )
+            cmakeArgs.append("-DTFLM_OPTIMIZED_KERNEL=cmsis_nn")
+            cmakeArgs.append(f"-DTFLM_OPTIMIZED_KERNEL_LIB={cmsisnnLib}")
+            cmakeArgs.append(f"-DTFLM_OPTIMIZED_KERNEL_INCLUDE_DIR={cmsisIncs}")
+        utils.mkdirs(tflmcBuildDir)
+        # utils.cmake("-DTF_SRC=" + str(tfSrcDir), str(tflmcSrcDir), debug=params["dbg"], cwd=tflmcBuildDir)
+        utils.cmake(
+            *cmakeArgs,
+            str(tflmcSrcDir),
+            debug=dbg,
+            cwd=tflmcBuildDir,
+            live=verbose,
+        )
+        utils.make(cwd=tflmcBuildDir, threads=threads, live=verbose)
+        utils.mkdirs(tflmcInstallDir)
+        utils.move(tflmcBuildDir / "compiler", tflmcExe)
+    context.cache["tflmc.build_dir", flags] = tflmcBuildDir
+    context.cache["tflmc.exe", flags] = tflmcExe

mlonmcu/setup/tasks/tvm.py

@@ -0,0 +1,172 @@
+#
+# 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.
+#
+"""Definition of tasks used to dynamically install MLonMCU dependencies"""
+
+import multiprocessing
+from pathlib import Path
+
+from mlonmcu.setup.task import TaskType
+from mlonmcu.context.context import MlonMcuContext
+from mlonmcu.setup import utils
+from mlonmcu.logging import get_logger
+
+from .common import get_task_factory
+
+logger = get_logger()
+
+Tasks = get_task_factory()
+
+
+def _validate_tvm(context: MlonMcuContext, params=None):
+    user_vars = context.environment.vars
+    patch = bool(params.get("patch", False))
+    if patch:
+        if not context.environment.has_feature("disable_legalize"):
+            return False
+    return context.environment.has_framework("tvm")
+
+
+def _validate_tvm_build(context: MlonMcuContext, params=None):
+    user_vars = context.environment.vars
+    use_tlcpack = user_vars.get("tvm.use_tlcpack", False)
+    patch = bool(params.get("patch", False))
+    if patch:
+        if not context.environment.has_feature("disable_legalize"):
+            return False
+    if use_tlcpack:
+        assert not context.environment.has_feature("disable_legalize")
+        return False
+
+    return context.environment.has_framework("tvm")
+
+
+@Tasks.provides(["tvm.src_dir", "tvm.configs_dir"])
+@Tasks.optional(["tvm_extensions.src_dir"])
+@Tasks.validate(_validate_tvm)
+@Tasks.param("patch", [False, True])  # This is just a temporary workaround until the patch is hopefully upstreamed
+@Tasks.validate(_validate_tvm)
+@Tasks.register(category=TaskType.FRAMEWORK)
+def clone_tvm(context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()):
+    """Clone the TVM repository."""
+    if not params:
+        params = {}
+    patch = params["patch"]
+    flags = utils.makeFlags((patch, "patch"))
+    tvmName = utils.makeDirName("tvm", flags=flags)
+    tvmSrcDir = context.environment.paths["deps"].path / "src" / tvmName
+    tvmPythonPath = tvmSrcDir / "python"
+    if rebuild or not utils.is_populated(tvmSrcDir):
+        tvmRepo = context.environment.repos["tvm"]
+        utils.clone(tvmRepo.url, tvmSrcDir, branch=tvmRepo.ref, recursive=True)
+        if patch:
+            extSrcDir = context.cache["tvm_extensions.src_dir"]
+            patchFile = extSrcDir / "tvmc_diff.patch"
+            utils.apply(tvmSrcDir, patchFile)
+
+    context.cache["tvm.src_dir", flags] = tvmSrcDir
+    context.cache["tvm.configs_dir", flags] = tvmSrcDir / "configs"
+    context.cache["tvm.pythonpath", flags] = tvmPythonPath
+
+
+@Tasks.needs(["tvm.src_dir", "llvm.install_dir"])
+@Tasks.provides(["tvm.build_dir", "tvm.lib"])
+@Tasks.param("dbg", False)
+@Tasks.param("cmsisnn", [False, True])
+@Tasks.validate(_validate_tvm_build)
+@Tasks.register(category=TaskType.FRAMEWORK)
+def build_tvm(context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()):
+    """Build the TVM framework."""
+    if not params:
+        params = {}
+    flags = utils.makeFlags((params["dbg"], "dbg"), (params["cmsisnn"], "cmsisnn"))
+    dbg = bool(params["dbg"])
+    cmsisnn = bool(params["cmsisnn"])
+    # FIXME: Try to use TVM dir outside of src dir to allow multiple versions/dbg etc!
+    # This should help: TVM_LIBRARY_PATH -> tvm.build_dir
+    tvmName = utils.makeDirName("tvm", flags=flags)
+    tvmSrcDir = context.cache["tvm.src_dir", ()]  # params["patch"] does not affect the build
+    tvmBuildDir = context.environment.paths["deps"].path / "build" / tvmName
+    tvmLib = tvmBuildDir / "libtvm.so"
+    user_vars = context.environment.vars
+    if rebuild or not utils.is_populated(tvmBuildDir) or not tvmLib.is_file():
+        ninja = False
+        if "tvm.make_tool" in user_vars:
+            if user_vars["tvm.make_tool"] == "ninja":
+                ninja = True
+        utils.mkdirs(tvmBuildDir)
+        cfgFileSrc = Path(tvmSrcDir) / "cmake" / "config.cmake"
+        cfgFile = tvmBuildDir / "config.cmake"
+        llvmConfig = str(context.cache["llvm.install_dir"] / "bin" / "llvm-config")
+        llvmConfigEscaped = str(llvmConfig).replace("/", "\\/")
+        utils.copy(cfgFileSrc, cfgFile)
+        utils.exec(
+            "sed",
+            "-i",
+            "--",
+            r"s/USE_MICRO OFF/USE_MICRO ON/g",
+            str(cfgFile),
+        )
+        utils.exec(
+            "sed",
+            "-i",
+            "--",
+            r"s/USE_MICRO_STANDALONE_RUNTIME OFF/USE_MICRO_STANDALONE_RUNTIME ON/g",
+            str(cfgFile),
+        )
+        utils.exec(
+            "sed",
+            "-i",
+            "--",
+            r"s/USE_LLVM \(OFF\|ON\)/USE_LLVM " + llvmConfigEscaped + "/g",
+            str(cfgFile),
+        )
+        if cmsisnn:
+            utils.exec(
+                "sed",
+                "-i",
+                "--",
+                r"s/USE_CMSISNN \(OFF\|ON\)/USE_CMSISNN ON/g",
+                str(cfgFile),
+            )
+
+        utils.cmake(tvmSrcDir, cwd=tvmBuildDir, debug=dbg, use_ninja=ninja, live=verbose)
+        utils.make(cwd=tvmBuildDir, threads=threads, use_ninja=ninja, live=verbose)
+    context.cache["tvm.build_dir", flags] = tvmBuildDir
+    context.cache["tvm.lib", flags] = tvmLib
+
+
+def _validate_tvm_extensions(context: MlonMcuContext, params=None):
+    return _validate_tvm_build(context, params=params) and context.environment.has_feature("disable_legalize")
+
+
+@Tasks.provides(["tvm_extensions.src_dir", "tvm_extensions.wrapper"])
+@Tasks.validate(_validate_tvm_extensions)
+@Tasks.register(category=TaskType.FEATURE)
+def clone_tvm_extensions(
+    context: MlonMcuContext, params=None, rebuild=False, verbose=False, threads=multiprocessing.cpu_count()
+):
+    """Clone the TVM extensions repository."""
+    extName = utils.makeDirName("tvm_extensions")
+    extSrcDir = context.environment.paths["deps"].path / "src" / extName
+    extWrapper = extSrcDir / "tvmc_wrapper.py"
+    if rebuild or not utils.is_populated(extSrcDir):
+        extRepo = context.environment.repos["tvm_extensions"]
+        utils.clone(extRepo.url, extSrcDir, branch=extRepo.ref, refresh=rebuild)
+    context.cache["tvm_extensions.src_dir"] = extSrcDir
+    context.cache["tvm_extensions.wrapper"] = extWrapper