module 导出

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)

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

先定义一些头文件。

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
目录