diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py index 70bbedbf6ff3..4d8f97d91d9c 100644 --- a/python/tvm/contrib/hexagon/session.py +++ b/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: diff --git a/python/tvm/contrib/hexagon/tools.py b/python/tvm/contrib/hexagon/tools.py index 175311294aa3..3b129b03323b 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 15c2a5a258af..de53017ca4a6 100644 --- a/python/tvm/runtime/module.py +++ b/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}" ) diff --git a/src/target/codegen.cc b/src/target/codegen.cc index d1f2d4a4795e..a221fa60e63c 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