【TVM】——ubuntu18.04源码编译TVM

tvm, ubuntu18.04

1.创建conda环境

# make sure to start with a fresh environment
conda env remove -n tvm-build-venv
# create the conda environment with build dependency
conda create -n tvm-build-venv -c conda-forge \
    "llvmdev>=15" \
    "cmake>=3.24" \
    git \
    python=3.11
# enter the build environment
conda activate tvm-build-venv

注意

• llvmdev要从conda-forge安装，我从清华源安装会报错“tvm LLVM ERROR: Only small, tiny and large code models are allowed on AArch64”
• llvmdev版本有可能报错：OSError: /root/miniconda3/lib/libstdc++.so.6: version `GLIBCXX_3.4.30’ not found (required by /root/miniconda3/envs/snpe/lib/libLLVM-16.so)，应该是预编译的libllvm和系统的c++不匹配。conda install llvmdev==14.0.0换个版本试一下

2.获得源码

git clone --recursive https://github.com/apache/tvm tvm

3.编译

打开llvm后端，
set(USE_LLVM ON)

cd tvm
rm -rf build && mkdir build && cd build
# Specify the build configuration via CMake options
# 
cp ../cmake/config.cmake .
cmake .. && cmake --build . --parallel $(nproc)

编译之后会有libtvm.so和libtvm_runtime.so在build/文件夹下

将tvm python接口添加到python路径，打开~/.bashrc

export TVM_HOME=/path-to-tvm
export PYTHONPATH=$TVM_HOME/python:$PYTHONPATH

4.安装测试

>>> python -c "import tvm; print(tvm.__file__)"
/some-path/lib/python3.11/site-packages/tvm/__init__.py

5.编译/runtime测试

创建以下python文件并执行

import tvm
from tvm import relax
from tvm.relax.frontend import nn


class MLPModel(nn.Module):
    def __init__(self):
        super(MLPModel, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        return x
    

mod, param_spec = MLPModel().export_tvm(
    spec={"forward": {"x": nn.spec.Tensor((1, 784), "float32")}}
)
mod.show()

import numpy as np

target = tvm.target.Target("llvm")
ex = relax.build(mod, target)
device = tvm.cpu()
vm = relax.VirtualMachine(ex, device)
data = np.random.rand(1, 784).astype("float32")
tvm_data = tvm.nd.array(data, device=device)
params = [np.random.rand(*param.shape).astype("float32") for _, param in param_spec]
params = [tvm.nd.array(param, device=device) for param in params]
print(vm["forward"](tvm_data, *params).numpy())

结果：

# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def forward(x: R.Tensor((1, 784), dtype="float32"), fc1_weight: R.Tensor((256, 784), dtype="float32"), fc1_bias: R.Tensor((256,), dtype="float32"), fc2_weight: R.Tensor((10, 256), dtype="float32"), fc2_bias: R.Tensor((10,), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"):
        R.func_attr({"num_input": 1})
        with R.dataflow():
            permute_dims: R.Tensor((784, 256), dtype="float32") = R.permute_dims(fc1_weight, axes=None)
            matmul: R.Tensor((1, 256), dtype="float32") = R.matmul(x, permute_dims, out_dtype="void")
            add: R.Tensor((1, 256), dtype="float32") = R.add(matmul, fc1_bias)
            relu: R.Tensor((1, 256), dtype="float32") = R.nn.relu(add)
            permute_dims1: R.Tensor((256, 10), dtype="float32") = R.permute_dims(fc2_weight, axes=None)
            matmul1: R.Tensor((1, 10), dtype="float32") = R.matmul(relu, permute_dims1, out_dtype="void")
            add1: R.Tensor((1, 10), dtype="float32") = R.add(matmul1, fc2_bias)
            gv: R.Tensor((1, 10), dtype="float32") = add1
            R.output(gv)
        return gv

[[27072.736 24800.156 25007.713 25144.564 23554.26  25691.121 25755.082
  23728.863 26370.531 25379.229]]

参考

1. tvm中文文档
2. tvm官方文档