TVM1 介绍/Ubuntu安装/向量加法转置举例

Posted 2022-11-25

文章目录

• ​​什么是TVM​​
• ​​一.TVM安装​​

• ​​1.1 需要安装的软件​​
• ​​1.2 安装​​
• ​​1.3 python环境​​
• ​​1.4 C++ test​​
• ​​1.5 python链接tvm​​
• ​​1.6 测试(显示下面表示成功了,因为这里只需要CPU代码优化, 因此只使用了LLVM)​​

• ​​二. TVM 例子​​

• ​​2.1 测试可用​​
• ​​2.2 compute 和 schedule分离​​
• ​​2.3 tvm的数据类型​​
• ​​2.4 变化的形状​​
• ​​2.5 矩阵转置​​

什么是TVM

网上搜索一大堆一个看一遍, 把共性的东西了解一下就明白了

• 该工具使用的是C++的后端,python作为前端直接调用TVM包进行代码优化
• TVM是一套完整的stack，包括神经网络图优化（比如op fusion）和单个operation优化等部分
• 不同后端设备上实现现有的operation
• TVM可以优化的训练好的模型，并将你的模型打包好，然后你可以将这个优化好的模型放在任何平台去运行

一.TVM安装

环境
win10子系统
Ubuntu18
​llvm安装链接​​

1.1 需要安装的软件

llvm-as -version # 6.0.1, llvm安装需要用到博客内容
clang -v # 6.0.1
cmake -version # 3.10.2
gcc/g++ -v # 7.5.0

环境中包含python2, python3, 同时安装conda自带python, 目前使用的是conda自带的python

1.2 安装

• 下载

git clone --recursive https://github.com/apache/incubator-tvm tvm && cd tvm
# 上面的仓库已经保存到了国内仓库gitee中, 可用下面的代替
git clone https://gitee.com/nwu_zjq/tvm.git && cd tvm
# 下面的命令式添加子模块
git submodule init
git submodule update

• cmake编译

mkdir build
cp cmake/config.cmake build/
cd build/
vim config.cmake # 这里设置打开llvm,debug等重要工具
# set(USE_MICRO ON)
# set(USE_CPP_RPC ON)
# set(USE_GRAPH_RUNTIME_DEBUG ON)
# set(USE_MICRO_STANDALONE_RUNTIME ON)
# set(USE_LLVM ON)
cmake ..
make -j4
# 这里完成后会生成libtvm.so

1.3 python环境

• 创建名称为tvm,python版本为3.6的环境

conda create -n tvm python=3.6
source activate tvm # 进入搭建的环境中
pip install tornado psutil xgboost mxnet tensorflow==2.0.0 # 安装包
sudo vi ~/.bashrc # 将刚编译好的tvm放到系统环境变量中
# 插入下面
# export TVM_HOME="/mnt/e/00_TVM/tvm"
# export PYTHONPATH=$TVM_HOME/python:$PYTHONPATH

1.4 C++ test

git clone https://github.com/google/googletest
cd googletest
mkdir build
cd build
cmake ..
make
sudo make install

1.5 python链接tvm

# cd /mnt/E:\\00_TVM\\tvm\\python
python setup.py install

1.6 测试(显示下面表示成功了,因为这里只需要CPU代码优化, 因此只使用了LLVM)

二. TVM 例子

2.1 测试可用

import tvm
n = 1024
A = tvm.te.placeholder((n,), name=A)
B = tvm.te.placeholder((n,), name=B)
C = tvm.te.compute(A.shape, lambda i: A[i] + B[i], name="C")
print(type(C))  #<class tvm.te.tensor.Tensor>
print(C)

<class tvm.te.tensor.Tensor>
Tensor(shape=[1024], op.name=C)

2.2 compute 和 schedule分离

• Halide项目将所有operation的程序做了拆解，认为所有程序由两部分组成:compute和schedule。而将compute和schedule相分离也是一个里程碑式的重要想法，为自动代码生成提供了可能。

# 代码1
for(int i = 0; i < 10; i ++) 
  for(int j = 0; j < 10; j++) 
    b[i][j] = a[i][j];
  

# 代码2
for(int i = 0; i < 10; i ++) 
  for(int j = 0; j < 10; j++) 
    b[j][i] = a[j][i];
  

• 代码1和代码2的compute一样
• 显然性能是不一样的（考虑cache的命中率等等），这些涉及具体实现的因素就是schedule：嵌套循环之间的顺序就是schedule的一种，除此之外还有更复杂的schedule：比如在GPU上，哪些数组要放到share memory，某个block到底负责处理那些数据，某个循环是否要展开等等。

# 这是一个向量加法的实例
# placeholder是一种特殊的tensor可以理解为任意维度的数组, 和pytorch一样

import tvm 
# https://blog.csdn.net/Tao_758/article/details/108046094
# 发现新版本的tvm里，var、placeholder等都移到了 tvm.te下，因此改为tvm.te.*即可

# 这部分定义了compute
n = 1024
A = tvm.te.placeholder((n,), name=A) # 声明维度为(1024,),命名为A 
B = tvm.te.placeholder((n,), name=B)
C = tvm.te.compute((n,), lambda i: A[i] + B[i], name="C")
print(type(C))  #<class tvm.te.tensor.Tensor>
print(C)

<class tvm.te.tensor.Tensor>
Tensor(shape=[1024], op.name=C)

# 这部分定义了schedule部分
s = tvm.te.create_schedule(C.op)
print(s)

schedule(0x2698500)

# 默认执行设计
code_default = tvm.lower(s, [A, B, C], simple_mode=True)
print("默认生成的代码是:\\n",code_default)

默认生成的代码是:
 #[version = "0.0.5"]
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = "global_symbol": "main", "tir.noalias": True
  buffers = A: Buffer(A_2: Pointer(float32), float32, [1024], []),
             C: Buffer(C_2: Pointer(float32), float32, [1024], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024], [])
  buffer_map = A_1: A, B_1: B, C_1: C 
  for (i: int32, 0, 1024) 
    C_2[i] = ((float32*)A_2[i] + (float32*)B_2[i])
  

2.3 tvm的数据类型

# 使用tvm.build操作将定义好的计算和设计编译成可执行模块

import tvm
import numpy as np
 
n = 100
 
# Defined in file: ./chapter_expression/vector_add.md
def eval_mod(mod, *inputs, out):
    """Evaluate a TVM module, and save results in out.
    """
    # Convert all numpy arrays to tvm arrays
    tvm_args = [tvm.nd.array(x) if isinstance(x, np.ndarray) 
                else x for x in inputs + (out,)]
    mod(*tvm_args)
    # If out is a tvm array, then its value has already been inplaced. 
    # Otherwise, explicitly copy the results. 
    if isinstance(out, np.ndarray):
        np.copyto(out, tvm_args[-1].asnumpy())
 
 
def tvm_vector_add(dtype):
    A = tvm.te.placeholder((n,), dtype=dtype)
    B = tvm.te.placeholder((n,), dtype=dtype)
    C = tvm.te.compute(A.shape, lambda i: A[i] + B[i])
    s = tvm.te.create_schedule(C.op)
    return tvm.build(s, [A, B, C])
 
mod = tvm_vector_add(int32)
 
def test_mod(mod, dtype):
    # you can use astype to convert data type
    a, b = (np.random.normal(size=100).astype(dtype) for _ in range(2))
    c = np.empty(100, dtype=dtype)
    eval_mod(mod, a, b, out=c)
    print("data type of c: ".format(c.dtype))
    np.testing.assert_equal(c, a + b)
 
test_mod(mod, int32)
 
for dtype in [float16, float64, int8, int16, int64]:
    mod = tvm_vector_add(dtype)
    test_mod(mod, dtype)

data type of c: int32
data type of c: float16
data type of c: float64
data type of c: int8
data type of c: int16
data type of c: int64

2.4 变化的形状

在定义计算时，可能对于输入的形状是未知的，可以通过tvm.var定义一个变量来指定形状，然后在具体调用时，传入具体值即可，如对于定义A、B、C三个占位符时，如果不知道输入的维度，可以使用变量n来创建任意长度数组：

import numpy as np
import tvm
np.random.seed(0)
def eval_mod(mod, *inputs, out):
    tvm_args = [tvm.nd.array(x) if isinstance(x, np.ndarray) else x for x in inputs + (out,)]
    mod(*tvm_args)
    if isinstance(out, np.ndarray):
        np.copyto(out, tvm_args[-1].asnumpy())
        
n = tvm.te.var(name=n)
print(type(n), n.dtype)
 
A = tvm.te.placeholder((n,), name=a)
B = tvm.te.placeholder((n,), name=b)
C = tvm.te.compute(A.shape, lambda i: A[i] + B[i])
s = tvm.te.create_schedule(C.op)
print("原始生成\\n",tvm.lower(s, [A, B, C], simple_mode=True))

def test_mod(mod, size):
    a, b = (np.random.normal(size=size).astype(float32) for _ in range(2))
    c = np.empty(size, dtype=float32)
    print("c shape: ".format(c.shape))
    eval_mod(mod, a, b, out=c)
    np.testing.assert_equal(c, a + b)
 
mod = tvm.build(s, [A, B, C])
test_mod(mod, 5)
test_mod(mod, 1000)
 
def tvm_vector_add(ndim):
    A = tvm.te.placeholder([tvm.te.var() for _ in range(ndim)])
    B = tvm.te.placeholder(A.shape)
    C = tvm.te.compute(A.shape, lambda *i: A[i] + B[i])
    s = tvm.te.create_schedule(C.op)
    return tvm.build(s, [A, B, C])
 
mod = tvm_vector_add(2)
test_mod(mod, (2, 2))
 
mod = tvm_vector_add(4)
test_mod(mod, (2, 3, 4, 5))

<class tvm.tir.expr.Var> int32
原始生成
 #[version = "0.0.5"]
primfn(a_1: handle, b_1: handle, compute_1: handle) -> ()
  attr = "global_symbol": "main", "tir.noalias": True
  buffers = compute: Buffer(compute_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             a: Buffer(a_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),
             b: Buffer(b_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")
  buffer_map = a_1: a, b_1: b, compute_1: compute 
  for (i: int32, 0, n) 
    compute_2[(i*stride)] = ((float32*)a_2[(i*stride_1)] + (float32*)b_2[(i*stride_2)])
  

2.5 矩阵转置

# 矩阵转置
import tvm
n = tvm.te.var("n")
m = tvm.te.var("n")

A = tvm.te.placeholder((n,m), name=A)
B = tvm.te.compute((m,n), lambda i,j:A[j,i], "B")
s = tvm.te.create_schedule(B.op)
tvm.lower(s,[A,B], simple_mode=True)

#[version = "0.0.5"]
primfn(A_1: handle, B_1: handle) -> ()
  attr = "global_symbol": "main", "tir.noalias": True
  buffers = B: Buffer(B_2: Pointer(float32), float32, [n: int32, n_1: int32], [stride: int32, stride_1: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n_1, n], [stride_2: int32, stride_3: int32], type="auto")
  buffer_map = A_1: A, B_1: B 
  for (i: int32, 0, n) 
    for (j: int32, 0, n_1) 
      B_2[((i*stride) + (j*stride_1))] = (float32*)A_2[((j*stride_2) + (i*stride_3))]
    
  

ref:

​​利用Pytorch导出Onnx模型进行模型优化例子​​

​​利用tutorials子文件测试图像识别例子​​

​​TVM代码库演练示例​​

​​知乎例子​​