You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2023/09/06 12:51:53 UTC
[tvm] branch main updated: [Module] Implement custom imported modules serialization (#15666)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 35164b379f [Module] Implement custom imported modules serialization (#15666)
35164b379f is described below
commit 35164b379f2bb64eef2f4c5d2a16b085ee75ca9a
Author: Krzysztof Parzyszek <kp...@quicinc.com>
AuthorDate: Wed Sep 6 07:51:46 2023 -0500
[Module] Implement custom imported modules serialization (#15666)
* [Module] Implement custom imported modules serialization
When a module with imported modules is exported into a shared library,
the imported modules are serialized and embedded inside of that library.
This is done by generating a raw binary from the imported modules, which
is then assigned to a symbol `__tvm_dev_mblob` in the final shared
library.
The way it happens for targets that are not "llvm", is by creating a C
source file, and defining a statically-initialized array `__tvm_dev_mblob`
in it. The static initializer is the byte-by-byte hexadecimal represen-
tation of the serialized modules. While working with Hexagon, this has
presented us with two issues:
1. For most models, the embedded data is very large, taking significant
amout of time to compile the auto-generated C source.
2. There are some models for which the C source size has exceeded clang's
limits on the size of the input file, making it impossible to export the
corresponding module.
This PR allows users to provide a custom serialization routine to
`Module.export_library`. We then apply it in Hexagon to build the ELF
object file with `__tvm_dev_mblob` by using the objcopy tool from the
Hexagon toolchain. This bypasses the C compilation altogether, avoiding
both of the issues.
Using the same custom mechanism, similar method can be implemented for
many other targets which have utilities to manipulate object files
directly.
* Allow toolchain version 8.5.x as well: objcopy has necessary options
It seems like older toolchains may work as well, but the oldest supported
SDK is 4.5.0.3, which contains toolchain version 8.5.08.
---
python/tvm/contrib/hexagon/session.py | 2 +
python/tvm/contrib/hexagon/tools.py | 130 +++++++++++++++++++++++++++++++++-
python/tvm/runtime/module.py | 24 ++++++-
src/target/codegen.cc | 75 +++++++++++---------
4 files changed, 196 insertions(+), 35 deletions(-)
diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py
index 70bbedbf6f..4d8f97d91d 100644
--- a/python/tvm/contrib/hexagon/session.py
+++ b/python/tvm/contrib/hexagon/session.py
@@ -399,12 +399,14 @@ class Session:
module.export_library(
str(binary_path),
fcompile=hexagon.create_aot_shared,
+ fpack_imports=hexagon.pack_imports,
hexagon_arch=hexagon_arch,
)
elif target_type == "llvm":
module.export_library(
str(binary_path),
fcompile=hexagon.create_shared,
+ fpack_imports=hexagon.pack_imports,
cc=hexagon.hexagon_clang_plus(),
)
else:
diff --git a/python/tvm/contrib/hexagon/tools.py b/python/tvm/contrib/hexagon/tools.py
index 175311294a..3b129b0332 100644
--- a/python/tvm/contrib/hexagon/tools.py
+++ b/python/tvm/contrib/hexagon/tools.py
@@ -19,7 +19,9 @@
import os
import pathlib
-from typing import Union
+import re
+from typing import List, Union
+import subprocess
import sys
import tarfile
import io
@@ -79,6 +81,37 @@ def hexagon_clang_plus() -> str:
return str(HEXAGON_CLANG_PLUS)
+def toolchain_version(toolchain=None) -> List[int]:
+ """Return the version of the Hexagon toolchain.
+
+ Parameters
+ ----------
+ toolchain: str, optional
+ Path to the Hexagon toolchain. If not provided, the environment
+ variable HEXAGON_TOOLCHAIN is used.
+
+ Returns
+ -------
+ version: List[int]
+ List of numerical components of the version number. E.g. for version
+ "8.5.06" it will be [8, 5, 6].
+ """
+
+ if toolchain is None:
+ toolchain = HEXAGON_TOOLCHAIN
+ assert toolchain is not None, "Please specify toolchain, or set HEXAGON_TOOLCHAIN variable"
+ result = subprocess.run(
+ [f"{toolchain}/bin/hexagon-clang", "-v"], capture_output=True, check=True
+ )
+ output = result.stderr.decode()
+ for line in output.splitlines():
+ m = re.match(r".* [Cc]lang version ([0-9\.]+)", line)
+ if m:
+ assert len(m.groups()) == 1
+ return [int(v) for v in m.group(1).split(".")]
+ raise RuntimeError("Cannot establish toolchain version")
+
+
@register_func("tvm.contrib.hexagon.link_shared")
def link_shared(so_name, objs, extra_args=None):
"""Link shared library on Hexagon using the registered Hexagon linker.
@@ -98,6 +131,7 @@ def link_shared(so_name, objs, extra_args=None):
ret_val : int
This function returns 0 at the moment.
"""
+
# The list of object files can be passed as built-in Python strings,
# or as tvm.tir.StringImm's.
def to_str(s):
@@ -168,6 +202,7 @@ def link_shared_macos(so_name, objs, extra_args=None):
ret_val : int
This function returns 0 at the moment.
"""
+
# The list of object files can be passed as built-in Python strings,
# or as tvm.tir.StringImm's.
def to_str(s):
@@ -273,6 +308,99 @@ def create_aot_shared(so_name: Union[str, pathlib.Path], files, hexagon_arch: st
cross_compile(str(so_name), c_files, options=compile_options + options)
+def pack_imports(
+ module: tvm.runtime.Module,
+ is_system_lib: bool, # pylint: disable=unused-argument
+ c_symbol_prefix: str,
+ workspace_dir: str,
+):
+ """Create an ELF object file that contains the binary data for the modules
+ imported in `module`. This is a callback function for use as `fpack_imports`
+ in `export_library`.
+
+ Parameters
+ ----------
+ module: tvm.runtime.Module
+ Module whose imported modules need to be serialized.
+ is_system_lib: bool
+ Flag whether the exported module will be used as a system library.
+ c_symbol_prefix: str
+ Prefix to prepend to the blob symbol.
+ workspace_dir: str
+ Location for created files.
+
+ Returns
+ -------
+ file_name: str
+ The name of the created object file.
+ """
+
+ path_bin = os.path.join(workspace_dir, "imports.bin")
+ pack_to_bin_f_name = "runtime.ModulePackImportsToNDArray"
+ fpack_to_bin = tvm.get_global_func(pack_to_bin_f_name)
+ assert fpack_to_bin, f"Expecting {pack_to_bin_f_name} in registry"
+
+ fpack_to_bin(module).numpy().tofile(path_bin)
+
+ mblob_symbol = c_symbol_prefix + tvm.get_global_func("runtime.ModuleImportsBlobName")()
+
+ binary_size = os.path.getsize(path_bin)
+ hexagon_toolchain = os.environ.get("HEXAGON_TOOLCHAIN")
+ assert hexagon_toolchain, "Please set HEXAGON_TOOLCHAIN variable"
+ version = toolchain_version(hexagon_toolchain)
+ assert (
+ version[0] == 8 and version[1] >= 5
+ ), "Please use Hexagon toolchain version 8.5.x or later"
+ if version[1] <= 6:
+ path_o = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.o")
+ subprocess.run(
+ [
+ f"{hexagon_toolchain}/bin/hexagon-clang",
+ "-x",
+ "c",
+ "-c",
+ "/dev/null",
+ "-o",
+ path_o,
+ ],
+ check=True,
+ )
+ subprocess.run(
+ [
+ f"{hexagon_toolchain}/bin/hexagon-llvm-objcopy",
+ path_o,
+ "--add-section",
+ f".rodata={path_bin}",
+ "--add-symbol",
+ f"{mblob_symbol}=.rodata:0,object",
+ ],
+ check=True,
+ )
+ return path_o
+
+ else: # 8.6.07+
+ path_c = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.c")
+ path_o = os.path.join(workspace_dir, f"{c_symbol_prefix}devc.o")
+ with open(path_c, "w") as f:
+ f.write(
+ f"const unsigned char {mblob_symbol}[{binary_size}] "
+ f'__attribute__((section(".rodata"))) = {{0x1}};'
+ )
+ subprocess.run(
+ [f"{hexagon_toolchain}/bin/hexagon-clang", "-c", path_c, "-o", path_o], check=True
+ )
+ subprocess.run(
+ [
+ f"{hexagon_toolchain}/bin/hexagon-llvm-objcopy",
+ path_o,
+ "--update-section",
+ f".rodata={path_bin}",
+ ],
+ check=True,
+ )
+ return path_o
+
+
def export_module(module, out_dir, binary_name="test_binary.so"):
"""Export Hexagon shared object to a file."""
binary_path = pathlib.Path(out_dir) / binary_name
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 15c2a5a258..de53017ca4 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -439,7 +439,14 @@ class Module(object):
return self._collect_from_import_tree(lambda m: m.is_dso_exportable)
def export_library(
- self, file_name, *, fcompile=None, addons=None, workspace_dir=None, **kwargs
+ self,
+ file_name,
+ *,
+ fcompile=None,
+ fpack_imports=None,
+ addons=None,
+ workspace_dir=None,
+ **kwargs,
):
"""
Export the module and all imported modules into a single device library.
@@ -467,6 +474,16 @@ class Module(object):
If fcompile has attribute object_format, will compile host library
to that format. Otherwise, will use default format "o".
+ fpack_imports: function(mod: runtime.Module, is_system_lib: bool, symbol_prefix: str,
+ workspace_dir: str) -> str
+ Function used to pack imported modules from `mod` into a file suitable for passing
+ to fcompile as an input file. The result can be a C source, or an .o object file,
+ or any other file that the fcompile function can handle. The function returns the
+ name of the created file.
+
+ If not provided, the imported modules will be serialized either via packing to an
+ LLVM module, or to a C source file.
+
workspace_dir : str, optional
The path of the directory used to create the intermediate
artifacts when exporting the module.
@@ -569,7 +586,10 @@ class Module(object):
if self.imported_modules:
pack_lib_prefix = system_lib_prefix if system_lib_prefix else ""
- if enabled("llvm") and llvm_target_string:
+ if fpack_imports is not None:
+ path_out = fpack_imports(self, is_system_lib, pack_lib_prefix, workspace_dir)
+ files.append(path_out)
+ elif enabled("llvm") and llvm_target_string:
path_obj = os.path.join(
workspace_dir, f"{pack_lib_prefix}devc.{global_object_format}"
)
diff --git a/src/target/codegen.cc b/src/target/codegen.cc
index d1f2d4a479..a221fa60e6 100644
--- a/src/target/codegen.cc
+++ b/src/target/codegen.cc
@@ -306,16 +306,27 @@ runtime::Module DeserializeModuleFromBytes(std::string blob) {
return root_mod;
}
-std::string PackImportsToC(const runtime::Module& mod, bool system_lib,
- const std::string& c_symbol_prefix) {
+std::string PackImportsToBytes(const runtime::Module& mod) {
std::string bin = SerializeModuleToBytes(mod);
- std::string mdev_blob_name = c_symbol_prefix + runtime::symbol::tvm_dev_mblob;
+ uint64_t nbytes = bin.length();
+ std::string header;
+ for (size_t i = 0; i < sizeof(nbytes); ++i) {
+ header.push_back(((nbytes >> (i * 8)) & 0xffUL));
+ }
+ return header + bin;
+}
+
+std::string PackImportsToC(const runtime::Module& mod, bool system_lib,
+ const std::string& c_symbol_prefix) {
if (c_symbol_prefix.length() != 0) {
CHECK(system_lib)
<< "c_symbol_prefix advanced option should be used in conjuction with system-lib";
}
+ std::string mdev_blob_name = c_symbol_prefix + runtime::symbol::tvm_dev_mblob;
+ std::string blob = PackImportsToBytes(mod);
+
// translate to C program
std::ostringstream os;
os << "#ifdef _WIN32\n"
@@ -327,27 +338,15 @@ std::string PackImportsToC(const runtime::Module& mod, bool system_lib,
<< "extern \"C\" {\n"
<< "#endif\n";
os << "TVM_EXPORT extern const unsigned char " << mdev_blob_name << "[];\n";
- uint64_t nbytes = bin.length();
- os << "const unsigned char " << mdev_blob_name << "[" << bin.length() + sizeof(nbytes)
- << "] = {\n ";
+ os << "const unsigned char " << mdev_blob_name << "[" << blob.length() << "] = {";
os << std::hex;
- size_t nunit = 80 / 4;
- for (size_t i = 0; i < sizeof(nbytes); ++i) {
- // sperators
- if (i != 0) {
- os << ",";
+ size_t nunit = 100 / 5; // 100 columns, 5 chars per "0xab,"
+ for (size_t i = 0; i < blob.length(); ++i) {
+ if (i % nunit == 0) {
+ os << "\n ";
}
- os << "0x" << ((nbytes >> (i * 8)) & 0xffUL);
- }
- for (size_t i = 0; i < bin.length(); ++i) {
- // sperators
- if ((i + sizeof(nbytes)) % nunit == 0) {
- os << ",\n ";
- } else {
- os << ",";
- }
- int c = bin[i];
- os << "0x" << (c & 0xff);
+ int c = blob[i];
+ os << "0x" << std::setw(2) << std::setfill('0') << (c & 0xff) << ',';
}
os << "\n};\n";
if (system_lib) {
@@ -370,14 +369,7 @@ runtime::Module PackImportsToLLVM(const runtime::Module& mod, bool system_lib,
<< "c_symbol_prefix advanced option should be used in conjuction with system-lib";
}
- std::string bin = SerializeModuleToBytes(mod);
-
- uint64_t nbytes = bin.length();
- std::string header;
- for (size_t i = 0; i < sizeof(nbytes); ++i) {
- header.push_back(((nbytes >> (i * 8)) & 0xffUL));
- }
- std::string blob = header + bin;
+ std::string blob = PackImportsToBytes(mod);
TVMByteArray blob_byte_array;
blob_byte_array.size = blob.length();
blob_byte_array.data = blob.data();
@@ -392,9 +384,28 @@ runtime::Module PackImportsToLLVM(const runtime::Module& mod, bool system_lib,
TVM_REGISTER_GLOBAL("target.Build").set_body_typed(Build);
-// Export two auxiliary function to the runtime namespace.
-TVM_REGISTER_GLOBAL("runtime.ModulePackImportsToC").set_body_typed(PackImportsToC);
+// Export a few auxiliary function to the runtime namespace.
+TVM_REGISTER_GLOBAL("runtime.ModuleImportsBlobName").set_body_typed([]() -> std::string {
+ return runtime::symbol::tvm_dev_mblob;
+});
+
+TVM_REGISTER_GLOBAL("runtime.ModulePackImportsToNDArray")
+ .set_body_typed([](const runtime::Module& mod) {
+ std::string buffer = PackImportsToBytes(mod);
+ ShapeTuple::index_type size = buffer.size();
+ DLDataType uchar;
+ uchar.code = kDLUInt;
+ uchar.bits = 8;
+ uchar.lanes = 1;
+ DLDevice dev;
+ dev.device_type = kDLCPU;
+ dev.device_id = 0;
+ auto array = runtime::NDArray::Empty({size}, uchar, dev);
+ array.CopyFromBytes(buffer.data(), size);
+ return array;
+ });
+TVM_REGISTER_GLOBAL("runtime.ModulePackImportsToC").set_body_typed(PackImportsToC);
TVM_REGISTER_GLOBAL("runtime.ModulePackImportsToLLVM").set_body_typed(PackImportsToLLVM);
} // namespace codegen