Apache TVM 是一个端到端的深度学习编译框架，适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 → https://tvm.hyper.ai/ 

作者：Tianqi Chen

本文介绍如何在 TVM 中规约（reduce）。关联规约算子（如 sum/max/min）是线性代数运算的典型构造块。

描述行的总和

在 NumPy 语法中，计算行的总和可以写成 B = numpy.sum(A, axis=1)

下面几行描述了行求和操作。为创建一个规约公式，用 te.reduce_axis 声明了一个 reduction 轴，它接收规约的范围。 te.sum 接收要规约的表达式以及 reduction 轴，并计算声明范围内所有 k 值的总和。

等效的 C 代码如下：

Schedule 规约

有几种方法可以 Schedule Reduce，先打印出默认 Schedule 的 IR 代码。

输出结果：

IR 代码与 C 代码非常相似，reduction 轴类似于普通轴，可以拆分。

以下代码按不同的因子将 B 的行轴和轴进行拆分，得到一个嵌套 reduction。

输出结果：

把 B 的行绑定到 GPU 线程，从而构建一个 GPU 内核。

输出结果：

规约因式分解和并行化

构建规约时不能简单地在 reduction 轴上并行化，需要划分规约，将局部规约结果存储在数组中，然后再对临时数组进行规约。

rfactor 原语对计算进行了上述重写，在下面的调度中，B 的结果被写入一个临时结果 B.rf，分解后的维度成为 B.rf 的第一个维度。

输出结果：

B 的调度算子被重写为 B.f 的规约结果在第一个轴上的和。

输出结果：

跨线程规约

接下来可以在因子轴上进行并行化，这里 B 的 reduction 轴被标记为线程，如果唯一的 reduction 轴在设备中可以进行跨线程规约，则 TVM 允许将 reduction 轴标记为 thread。

也可以直接在规约轴上计算 BF。最终生成的内核会将行除以 blockIdx.x，将 threadIdx.y 列除以 threadIdx.x，最后对 threadIdx.x 进行跨线程规约。

输出结果：

结果内核与 NumPy 进行比较来验证结果内核的正确性。

用二维规约描述卷积

在 TVM 中，用简单的二维规约来描述卷积（过滤器大小 = [3, 3]，步长 = [1, 1]）。

输出结果：

定义一般交换规约运算

除了 te.sum, tvm.te.min 和 tvm.te.max 等内置规约操作外，还可以通过 te.comm_reducer 定义交换规约操作。

总结

本教程演示了如何规约 schedule。

• 用 reduce_axis 描述规约。

• 如需并行性（parallelism），用 rfactor 来分解轴。

• 通过 te.comm_reducer 定义新的规约操作。