[Module] Implement custom imported modules serialization

python/tvm/contrib/hexagon/session.py

@@ -399,12 +399,14 @@ def _aot_executor_from_factory(
                 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:

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

python/tvm/runtime/module.py

@@ -439,7 +439,14 @@ def _collect_dso_modules(self):
         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 @@ def export_library(
             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 @@ def export_library(
         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}"
                 )

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