tvm学习笔记（七）：简约操作

所谓Reduction（简约操作）其实就是说经过该操作之后，输入向量的维度会减小，例如，对输入向量沿着某个维度进行求和，先看一下python的代码：

import numpy as np

a = np.random.normal(size=(3, 4)).astype('float32')
print(a)
print(a.sum(axis=1))

打印结果为：

这里就是沿着第二个维度进行了求和操作，对应与具体实现如下：

import numpy as np

a = np.random.normal(size=(3, 4)).astype('float32')
print(a)
# print(a.sum(axis=1))

def sum_rows(a, b):
    n = len(b)
    for i in range(n):
        b[i] = np.sum(a[i,:])

b = np.empty((3,), dtype='float32')
sum_rows(a, b)
print(b)

对应输出为：

过程相当直观，就是对第一维进行循环迭代，沿着第二维对数组进行求和，在numpy中采用了:符号实现了数组沿着第二维做的切片操作，在tvm中使用tvm.reduce_axis函数创建了tvm做简约操作的维度方向，需要指明起点和终点位置：

import tvm
n, m = tvm.var('n'), tvm.var('m')
A = tvm.placeholder((n, m), name='a')
j = tvm.reduce_axis((0, m), name='j')
B = tvm.compute((n, ), lambda i: tvm.sum(A[i, j], axis=j), name='b')
s = tvm.create_schedule(B.op)
print(tvm.lower(s, [A, B], simple_mode=True))

对应操作为：

～～～未完待续～～～

参考资料：

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

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