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 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.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 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)

运行效果如下：

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 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())


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))

运行结果如下：

3.矩阵转置

 

参考资料：

[1] https://github.com/d2l-ai/d2l-tvm

[2] http://tvm.d2l.ai.s3-website-us-west-2.amazonaws.com/