module 导出

参考：tvm/tests/python/runtime/test_runtime_module_export.py

import shutil
from pathlib import Path

temp_dir = Path(".temp")
temp_dir.mkdir(exist_ok=True, parents=True)
# 创建临时目录
header_file_dir_path = temp_dir

CSourceModuleCreate

/*!
 * \brief Create a C source module for viewing and compiling GCC code.
 * \param code The code to be viewed.
 * \param fmt The code format.
 * \param func_names The name of functions inside the runtime module.
 * \param const_vars. The constant variables that the c source module needs.
 * \return The created module.
 */
runtime::Module CSourceModuleCreate(const String& code, const String& fmt,
                                    const Array<String>& func_names,
                                    const Array<String>& const_vars = {});

CSourceModuleCreate 用于创建 C 源代码模块，以便查看和编译 GCC 代码。

参数说明:

1. const String& code

   • 这是字符串类型的参数，表示需要被查看和编译的 GCC 代码。

   • const 表示这个参数在函数内部不会被修改。

   • String& 表示这是字符串的引用，避免了不必要的拷贝。

2. const String& fmt

   • 这是字符串类型的参数，表示代码的格式。

   • 可能用于指定代码的格式化方式（如缩进、换行等），或者用于指定代码的语言类型（如 C、C++ 等）。

3. const Array<String>& func_names

   • 这是字符串数组类型的参数，表示运行时模块中函数的名称。

   • Array<String> 可能是自定义的数组类型，用于存储多个字符串。

   • 这些函数名称可能是需要在生成的 C 源代码模块中导出的函数。

4. const Array<String>& const_vars = {}

   • 这是可选的字符串数组类型的参数，表示C源代码模块中需要的常量变量。

   • = {} 表示这个参数有默认值，即空数组。如果调用函数时不提供这个参数，函数会使用空数组作为默认值。

返回值:

• runtime::Module

  • 这是函数的返回值类型，表示创建的 C 源代码模块。

  • runtime::Module 可能是自定义的类或结构体，用于表示运行时模块。

函数功能:

• 该函数的主要功能是创建 C 源代码模块，该模块可以用于查看和编译 GCC 代码。

• 生成的模块可能包含指定的函数名称和常量变量，并且可以根据提供的代码格式进行格式化。

使用场景

• 这个函数可能用于编译器或代码生成工具中，用于将高级语言代码（如 GCC 代码）转换为 C 源代码模块，以便进一步编译或执行。

• 通过指定函数名称和常量变量，可以定制生成的C源代码模块的内容。

示例调用

import tvm
code = r"int main() { return 0; }";
fmt = "c"
func_names = ["main"]
const_vars = ["MAX_VALUE"]
csource_module = tvm.runtime._ffi_api.CSourceModuleCreate(code, fmt, func_names, const_vars)

下面介绍更加复杂的例子。

先定义一些头文件。

Show code cell content Hide code cell content

import textwrap
def gen_engine_header(header_file_dir_path):
    code = r"""
        #ifndef _ENGINE_H_
        #define _ENGINE_H_
        #include <cstdint>
        #include <string>
        #include <sstream>
        #include <vector>
        class Engine {
        };

        #endif
        """
    header_file_dir_path = Path(header_file_dir_path)
    header_file = header_file_dir_path/"gcc_engine.h"
    code = textwrap.dedent(code).lstrip()
    with open(header_file, "w") as f:
        f.write(code)

def generate_engine_module(header_file_dir_path):
    code = r"""
        #include <tvm/runtime/c_runtime_api.h>
        #include <dlpack/dlpack.h>
        #include "gcc_engine.h"

        extern "C" void gcc_1_(float* gcc_input4, float* gcc_input5,
                float* gcc_input6, float* gcc_input7, float* out) {
            Engine engine;
        }
        """
    import tvm.runtime._ffi_api

    gen_engine_header(header_file_dir_path)
    code = textwrap.dedent(code).lstrip()
    csource_module = tvm.runtime._ffi_api.CSourceModuleCreate(code, "cc", [], None)
    return csource_module

动态库导出

from tvm import relay, te
import tvm.relay.testing
import os
os.environ['PATH'] += ':/usr/local/cuda/bin' # 保证 nvcc 可以被找到
synthetic_mod, synthetic_params = relay.testing.synthetic.get_workload()
synthetic_llvm_mod, synthetic_llvm_params = relay.testing.synthetic.get_workload()
with tvm.transform.PassContext(opt_level=3):
    synthetic_cpu_lib = relay.build_module.build(
        synthetic_llvm_mod, "llvm", params=synthetic_llvm_params, mod_name="llvmlib"
    )
A = te.placeholder((1024,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
s = te.create_schedule(B.op)
f = tvm.build(s, [A, B], "c", name="myadd")
path_lib = temp_dir/"deploy_lib.so"
engine_module = generate_engine_module(header_file_dir_path)
synthetic_cpu_lib.module.import_module(f)
synthetic_cpu_lib.module.import_module(engine_module)
work_dir = temp_dir/"work_dir"
work_dir.mkdir(exist_ok=True)
kwargs = {"options": ["-O2", "-std=c++17", f"-I{header_file_dir_path}"]}
synthetic_cpu_lib.export_library(path_lib, fcompile=False, workspace_dir=work_dir, **kwargs)
loaded_lib = tvm.runtime.load_module(path_lib)
!ls {work_dir} # 查看生成的代码

One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.

devc.o	lib0.cc  lib1.c  lib2.o

静态库导出

生成两个 LLVM 模块：

A = te.placeholder((1024,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
s = te.create_schedule(B.op)
mod0 = tvm.build(s, [A, B], "llvm", name="myadd0")
mod1 = tvm.build(s, [A, B], "llvm", name="myadd1")

assert mod0.implements_function("myadd0")
assert mod1.implements_function("myadd1")
assert mod1.is_dso_exportable

mod1 当前是 'llvm' 模块，将其保存并重新加载为普通的 'static_library'：

mod1_o_path = f"{temp_dir}/mod1.o"
mod1.save(mod1_o_path)
mod1_o = tvm.runtime.load_static_library(mod1_o_path, ["myadd1"])
assert mod1_o.implements_function("myadd1")
assert mod1_o.is_dso_exportable

将 mod1 作为静态库导入到 mod0 中，并将其编译为独立的 DSO（动态共享对象）。

mod0.import_module(mod1_o)
mod0_dso_path = f"{temp_dir}/mod0.so"
mod0.export_library(mod0_dso_path)

导入的 mod1 被静态链接到 mod0 中。

loaded_lib = tvm.runtime.load_module(mod0_dso_path)
assert loaded_lib.type_key == "library"
assert len(loaded_lib.imported_modules) == 0
assert loaded_lib.implements_function("myadd0")
assert loaded_lib.get_function("myadd0")
assert loaded_lib.implements_function("myadd1")
assert loaded_lib.get_function("myadd1")
assert not loaded_lib.is_dso_exportable