mehrdadh commented on code in PR #11464:
URL: https://github.com/apache/tvm/pull/11464#discussion_r897143785
##########
python/tvm/micro/model_library_format.py:
##########
@@ -24,6 +24,7 @@
import re
import tarfile
import typing
+from typing import Union
Review Comment:
done.
##########
python/tvm/micro/contrib/stm32/emitter.py:
##########
@@ -482,8 +482,18 @@ def parse_library_format(self, model_library_format_path,
quantization=None):
with tarfile.TarFile(model_library_format_path) as f:
f.extractall(extract_path)
+ with open(os.path.join(extract_path, "metadata.json")) as metadata_f:
+ metadata = json.load(metadata_f)
+
+ all_module_names = []
+ for name in metadata["modules"].keys():
+ all_module_names.append(name)
+ assert len(all_module_names) == 1, "Multiple modules is not supported."
Review Comment:
done
##########
tests/python/unittest/test_micro_model_library_format.py:
##########
@@ -95,8 +101,35 @@ def test_export_operator_model_library_format():
assert tir_f.read() == str(ir_mod)
[email protected]_micro
+def test_export_multiple_operator_model_library_format():
+ target = tvm.target.target.micro("host")
+ with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
+ A = tvm.te.placeholder((2,), dtype="int8")
+ B = tvm.te.placeholder((1,), dtype="int8")
+ C = tvm.te.compute(A.shape, lambda i: A[i] + B[0], name="C")
+ sched = tvm.te.create_schedule(C.op)
+ mod = tvm.build(
+ sched,
+ [A, B, C],
+ tvm.target.Target(target, target),
+ runtime=Runtime("crt", {"system-lib": True}),
+ name="add",
+ )
+
+ temp_dir = utils.tempdir()
+ mlf_tar_path = temp_dir.relpath("lib.tar")
+
+ with pytest.raises(RuntimeError) as exc:
+ micro.export_model_library_format([mod, mod], mlf_tar_path)
+
+ assert str(exc.exception) == ("Multiple operator is not supported.")
Review Comment:
No I wanted to simply check the case where we pass multiple operator
modules. I can add another module with different name to make it more obvious.
thoughts?
##########
python/tvm/micro/model_library_format.py:
##########
@@ -67,56 +69,75 @@ def generate_c_interface_header(
EPHEMERAL_MODULE_TYPE_KEYS = ("metadata_module",)
-def _populate_codegen_dir(mod, codegen_dir: str, module_name: str = None):
+def _populate_codegen_dir(
+ mods: Union[
+ typing.List[executor_factory.ExecutorFactoryModule],
typing.List[tvm.runtime.Module]
+ ],
+ codegen_dir: str,
+):
"""Populate the codegen sub-directory as part of a Model Library Format
export.
Parameters
----------
- mod : tvm.runtime.Module
- Module which should be written to codegen_dir.
+ mods : List[tvm.relay.backend.executor_factory.ExecutorFactoryModule],
List[tvm.runtime.Module]
+ A list of the return value of tvm.relay.build, which
+ will be exported into Model Library Format.
codegen_dir : str
Path to the codegen directory on disk.
module_name: Optional[str]
Name used to prefix the generated source files
"""
- dso_modules = mod._collect_dso_modules()
- non_dso_modules = mod._collect_from_import_tree(lambda m: m not in
dso_modules)
-
- # Filter ephemeral modules which cannot be exported.
- dso_modules = [m for m in dso_modules if m.type_key not in
EPHEMERAL_MODULE_TYPE_KEYS]
- non_dso_modules = [m for m in non_dso_modules if m.type_key not in
EPHEMERAL_MODULE_TYPE_KEYS]
+ dso_modules = []
+ for mod in mods:
+ if isinstance(mod, executor_factory.ExecutorFactoryModule):
+ lib = mod.lib
+ elif isinstance(mod, tvm.runtime.Module):
Review Comment:
done
##########
python/tvm/micro/model_library_format.py:
##########
@@ -67,56 +69,75 @@ def generate_c_interface_header(
EPHEMERAL_MODULE_TYPE_KEYS = ("metadata_module",)
-def _populate_codegen_dir(mod, codegen_dir: str, module_name: str = None):
+def _populate_codegen_dir(
+ mods: Union[
+ typing.List[executor_factory.ExecutorFactoryModule],
typing.List[tvm.runtime.Module]
Review Comment:
I changed it to `build_module.OperatorModule`
I think Union of List is correct because the input could be Union of two
different list, but not a list of two different type
##########
tests/python/unittest/test_micro_model_library_format.py:
##########
@@ -439,5 +474,140 @@ def test_export_byoc_c_module():
]
[email protected]_micro
+def test_multiple_relay_modules_same_module_name():
+ mod = get_conv2d_relay_module()
+
+ executor = Executor("graph")
+ runtime = Runtime("crt")
+ target = tvm.target.target.micro("host")
+
+ with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
+ factory1 = tvm.relay.build(mod, target, runtime=runtime,
executor=executor, mod_name="mod")
+ factory2 = tvm.relay.build(mod, target, runtime=runtime,
executor=executor, mod_name="mod")
+
+ temp_dir = utils.tempdir()
+ mlf_tar_path = temp_dir.relpath("lib.tar")
+
+ with pytest.raises(AssertionError, match="Multiple modules should have
unique names"):
+ micro.export_model_library_format([factory1, factory2], mlf_tar_path)
+
+
[email protected]_micro
+def test_multiple_relay_modules_graph():
+ mod = get_conv2d_relay_module()
+
+ executor = Executor("graph")
+ runtime = Runtime("crt")
+ target = tvm.target.target.micro("host")
+
+ with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
+ factory1 = tvm.relay.build(mod, target, runtime=runtime,
executor=executor, mod_name="mod1")
+ factory2 = tvm.relay.build(mod, target, runtime=runtime,
executor=executor, mod_name="mod2")
+
+ temp_dir = utils.tempdir()
+ mlf_tar_path = temp_dir.relpath("lib.tar")
+ micro.export_model_library_format([factory1, factory2], mlf_tar_path)
+
+ with tarfile.open(mlf_tar_path, "r:*") as tf:
+ tar_members = [ti.name for ti in tf.getmembers()]
+ print("tar members", tar_members)
+ assert "./metadata.json" in tar_members
+ assert "./codegen/host/src/mod1_lib0.c" in tar_members
+ assert "./codegen/host/src/mod2_lib0.c" in tar_members
+
+ with tf.extractfile("./metadata.json") as f:
+ metadata = json.load(f)
+ mod2_main_md =
metadata["modules"]["mod2"]["memory"]["functions"]["main"]
+ assert mod2_main_md == [
+ {
+ "constants_size_bytes": 0,
+ "device": 1,
+ "io_size_bytes": 143960,
+ "workspace_size_bytes": 158088,
+ }
+ ]
+ assert metadata["modules"]["mod1"]["model_name"] == "mod1"
+ assert metadata["modules"]["mod2"]["model_name"] == "mod2"
+
+
[email protected]_micro
+def test_multiple_relay_modules_c():
+ mod = get_conv2d_relay_module()
+
+ executor = Executor("aot", {"unpacked-api": True, "interface-api": "c"})
+ runtime = Runtime("crt")
+ target = tvm.target.target.micro("host")
+
+ with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
+ factory1 = tvm.relay.build(mod, target, runtime=runtime,
executor=executor, mod_name="mod1")
+ factory2 = tvm.relay.build(mod, target, runtime=runtime,
executor=executor, mod_name="mod2")
+
+ temp_dir = utils.tempdir()
+ mlf_tar_path = temp_dir.relpath("lib.tar")
+
+ micro.export_model_library_format([factory1, factory2], mlf_tar_path)
+
+ tf = tarfile.open(mlf_tar_path)
+
+ extract_dir = temp_dir.relpath("extract")
+ os.mkdir(extract_dir)
+ tf.extractall(extract_dir)
+
+ assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src",
"mod1_lib0.c"))
+ assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src",
"mod1_lib1.c"))
+ assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src",
"mod2_lib0.c"))
+ assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src",
"mod2_lib1.c"))
+
+ assert os.path.exists(os.path.join(extract_dir, "codegen", "host",
"include", "tvmgen_mod1.h"))
+ assert os.path.exists(os.path.join(extract_dir, "codegen", "host",
"include", "tvmgen_mod2.h"))
+
+ # check CRT runtime directory
+ assert os.path.exists(os.path.join(extract_dir, "runtime"))
+
+
[email protected]_micro
+def test_multiple_relay_modules_aot_graph():
Review Comment:
I think there could be cases where someone build a graph along with an AOT
executor. Fundamentally I don't see a reason that we cannot support it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
