部署 TVM 模块

备注

TVM 提供了两种使用编译库的方法。

• 将库存储为共享库，并将库动态加载到项目中。

• 以系统模块（module）模式将编译后的库捆绑到项目中。

动态加载更加灵活，可以动态加载新模块。系统模块是更 static 的方法。可以在禁止动态库加载的地方使用系统模块。

加载 Python 模块：

import tvm
import numpy as np
from tvm import te
from tvm import relay
from pathlib import Path

下面分别以 TVM 和 Relay 为例，说明如何构建 TVM 模块。

TVM 构建

n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
s = te.create_schedule(B.op)

定义输出库的根目录：

base_dir = Path("libs")
base_dir.mkdir(exist_ok=True)

编译为动态库：

fadd_dylib = tvm.build(s, [A, B], "llvm", name="addone")
dylib_path = str(base_dir/"test_addone_dll.so")
fadd_dylib.export_library(dylib_path)

在系统库模式下编译库：

fadd_syslib = tvm.build(s, [A, B], "llvm", name="addonesys")
syslib_path = str(base_dir/"test_addone_sys.o")
fadd_syslib.save(syslib_path)

Relay 构建

定义 Relay 模块：

x = relay.var("x", shape=(2, 2), dtype="float32")
y = relay.var("y", shape=(2, 2), dtype="float32")
params = {"y": np.ones((2, 2), dtype="float32")}
mod = tvm.IRModule.from_expr(relay.Function([x, y], x + y))
mod.show()

def @main(%x: Tensor[(2, 2), float32], %y: Tensor[(2, 2), float32]) {
  add(%x, %y)
}

构建模块：

compiled_lib = relay.build(mod, tvm.target.Target("llvm"), params=params)

将其导出为共享库：

dylib_path = str(base_dir/"test_relay_add.so")
compiled_lib.export_library(dylib_path)

备注

如果正在运行交叉编译，还可以考虑导出到 tar 并稍后调用主机编译器。