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

本文主要是介绍tvm学习笔记（七）：简约操作，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

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

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

打印结果为：

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

import numpy as npa = 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))

对应操作为：