tvm学习笔记（六）：数据类型及形状

本文主要是介绍tvm学习笔记（六）：数据类型及形状，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

据说tvm会支持训练，然后沐神他们为了推广tvm，专门对tvm重写了文档d2l-tvm，具体详见资料[1]，下面就是对着沐神他们写的文档做的学习记录：

1.tvm的数据类型

我们在声明placeholder的时候，可以显式的指定数据的类型，如：'float16', 'float64', 'int8','int16', 'int64'等

import tvm
A = tvm.placeholder((n,), dtype='float32')
print(A.dtype)
B = tvm.placeholder((n,), dtype='int8')
print(B.dtype)

结果如下：

下面来看一下，如何使用tvm来定义两个向量相加的运算：

首先声明两个占位符A和B，并指定他们的shape：

A = tvm.placeholder((n,), dtype=dtype)
A = tvm.placeholder((n,), dtype=dtype)

然后使用使用tvm.compute定义C的计算，并使用lambda函数来指定输出C中每个元素的计算方式：

C = tvm.compute(A.shape, lambda i: A[i] + B[i])

其次，我们需要指定如何来执行程序，如，数据访问顺序及如何进行多线程并行计算，这样一个执行设计称之为一个schedule，C为输出，下面使用tvm.create_schedule就对C的操作创建一个默认的schedule：

s = tvm.create_schedule(C.op)

后面可以看到如何通过改变执行设计来更好的利用硬件资源，进而提升执行效率。这里，我们可以使用tvm.lower操作，打印一下当前默认执行设计：

tvm.lower(s, [A, B, C], simple_mode=True)

可以得到：

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

module = tvm.build(s, [A, B, C])

将上面的代码连在一起，就是使用tvm来创建两个向量相加的函数：

import tvm
import numpy as npn = 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 arraystvm_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.placeholder((n,), dtype=dtype)B = tvm.placeholder((n,), dtype=dtype)C = tvm.compute(A.shape, lambda i: A[i] + B[i])s = tvm.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 typea, 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)

运行效果如下：

2.变化形状

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

n = tvm.var(name='n')
print(type(n), n.dtype)A = tvm.placeholder((n,), name='a')
B = tvm.placeholder((n,), name='b')
C = tvm.compute(A.shape, lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
mod = tvm.build(s, [A, B, C])

在调用时，传入数组长度：

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)

完整代码如下：

import numpy as np
import tvm# 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 arraystvm_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())n = tvm.var(name='n')
print(type(n), n.dtype)A = tvm.placeholder((n,), name='a')
B = tvm.placeholder((n,), name='b')
C = tvm.compute(A.shape, lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
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.placeholder([tvm.var() for _ in range(ndim)])B = tvm.placeholder(A.shape)C = tvm.compute(A.shape, lambda *i: A[i] + B[i])s = tvm.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))

运行结果如下：