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 nnclass 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 xmod, param_spec = MLPModel().export_tvm(spec={"forward": {"x": nn.spec.Tensor((1, 784), "float32")}}
)
mod.show()import numpy as nptarget = 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.functiondef 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") = add1R.output(gv)return gv[[27072.736 24800.156 25007.713 25144.564 23554.26 25691.121 25755.08223728.863 26370.531 25379.229]]
