Page 482 - 《软件学报》2025年第8期

王昊天 等: MTTorch: 面向 MT-3000 芯片和 Transformer 模型的 PyTorch 算子库实现与优化 3905

8. von = not(vo64); //逻辑非
9. tmp1 = and(x64, vo64); //逻辑与
10. tmp2 = and(y64, von); //逻辑与
11. rett = or(tmp1, tmp2); //逻辑或
12. res = (lvector float)(rett); //强制类型转换

GELU 也是一种常用的激活函数, 相较于传统的激活函数如 Sigmoid 和 ReLU, GELU 在神经网络中具有更好
的性能, 同时也更易于优化. 它在训练时收敛速度更快, 并且比 ReLU 更能适应各种深度神经网络结构, 因此被广
泛应用于自然语言处理、计算机视觉等领域, 其公式为:

( ( ))
x
GELU (x) = 0.5× x× 1+erf √ (2)
2
其中, erf 为高斯误差函数. 其反向传播可由下式得到.

( ( ))
dy x x −x 2
= 0.5× 1+erf √ + √ exp (3)
dx 2 2π 2
MT-3000 提供了一系列高性能数学库 (HPML), 其中的数值计算库 libvm 包含了向量的 erf 和 exp 操作. 利用
libvm 和算法 1 防止 AM 内存溢出的思路, 本文使用向量指令实现了高效的 GELU 前向传播以及反向传播算子.

3.6 基于规约的算子实现与优化
归约操作是指将一组数据中的所有元素进行某种计算 (例如求和、求最大值、最小值等), 并得到单一的结
果. 在 PyTorch 中, 存在着大量基于规约的算子. 本文结合 MT-3000 DSP 的硬件特性, 设计了一个高效的多核规约
流程, 并借助其实现了高效的 softmax 和 layer_norm 算子.

3.6.1 MT-3000 多核规约
本节首先介绍了实现上述算子最关键的步骤——基于多核的并行规约算法. 归约的串行操作非常简单, 但是
由于每次循环都需要利用上一次循环的结果, 因此无法直接向量化. 二分法并行规约是规约算法的常见思路. 其将
源数组划分为多个子数组分别规约, 再对子数组的结果进行规约, 如此递归得到结果. 然而, 当前的规约算法大多
是针对 CUDA 平台或特定的指令集, 并不适用于 MT-3000 的异构众核架构. 本文结合 MT-3000 的硬件架构, 实现
了一个 DSP 上的多核向量化规约算法, 流程如图 8 所示.

均分为 N 块
···

核 1 DDR to AM 核 N DDR to AM VPU
···
向量计算
···
AM 向量内规约 (SVR)
GSM
···
N 个核间规约
图 8 MT-3000 DSP 多核规约流程

3.6.1.1 核内分段计算
算法首先根据参与规约的核数 N 将数组分为 N 份. 值得注意的是, 由于 VPE 位宽限制, 每一份的长度应该为
偶数. 本文使用 DMA 指令将分好的子数组搬移到 AM 存储器中用于接下来的计算. 使用 DMA 一方面可以充分
利用各级内存间的带宽, 另一方面也方便开发人员针对访存进行更精细的优化, 这将会在第 3.6.3 节中展开. 搬移