在带有NVIDIA GPU的环境中使用tvm

在带有NVIDIA GPU的环境中使用tvm

• 环境配置
• 使用官方提供的docker image
• 编译安装tvm[^1]
• 使用一下子

环境配置

用一个小表格表示一下搭建环境吧。

环境项	参数
cpu	Intel® Xeon® Gold 6142 cpu @ 2.60GHz
GPU	NVIDIA A100-PCI
OS	Ubuntu 18.04

使用官方提供的docker image

在clone过code后，code带有一些别人写好的docker images，直接用他们好了。

git clone --recursive https://github.com/shaojiewang/incubator-tvm tvm #仅第一次需要（自己fork的文件夹）
tvm/docker/build.sh Dockerfile.conda_cuda100
sudo tvm/docker/bash.sh tvm.conda_cuda100:latest

这个脚本干了如下几件事：

• 把当前目录挂载在/workspace目录下；

• 切换用户到root权限；

• 使用宿主机的网络；

第一次使用需要build image，以后就不用了。

编译安装tvm¹

这里我就不记录了，还是主要靠tvm官网的方式：编译安装方式

# 编译动态库文件
sudo apt-get update
sudo apt-get install -y python3 python3-dev python3-setuptools gcc libtinfo-dev zlib1g-dev build-essential cmake libedit-dev libxml2-dev
mkdir build
cp cmake/config.cmake build # 随后在config.cmake里面修改选项，打开想要使用的内容。在此时发现tvm已经支持hip和rocm了
cd build
cmake ..
make -j4
./install_python.sh

使用一下子

在Python里面import 一下子，证明装上去了：

import tvm

---

1. 更新版本后会有变动，主要需要查阅link ↩︎