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

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

研究中探索基于 Hthreads 的高性能矩阵乘法的实现方式.
MTTorch 的设计具有很好的通用性. MTTorch 中所有算子均使用 PyTorch 原生的 TORCH_LIBRARY_IMPL
方法在 PyTorch 后端进行注册, 以供图引擎调用. 其实现的算子以扩展库的形式编译和原生 PyTorch 框架解耦. 用
户只需通过 torch.ops.load_library 方法在 Python 端将编译好的 MTTorch 动态链接库进行挂载, 即可在 MT-3000
的 DSP 簇上调用算子进行高效训练. 本文的设计思路几乎未改变 PyTorch 的源码和框架结构, 因此广泛适用于各
版本的 PyTorch 框架.

3.3 虚拟算子

在实现虚拟算子时, 一种直观的实现方法是在算子内将输入的 Tensor 搬移到 CPU 控制的 DDR 区域执行
CPU 分支的算子, 再在计算结束之后将结果搬移到 DSP 可读写的 DDR 区域. 然而不同设备之间的切换需要执行
memcpy 函数进行内存拷贝. 经过测试, MT-3000 上 DMA 对 DDR 的访存带宽只有 16.22 GB/s, 这导致内存拷贝所
占用的时间是不可忽视的. 针对这类算子, 本文结合 MT-3000 的设备特性, 提出了一种无需内存拷贝的虚拟算子
实现方法. 在 PyTorch 中, Tensor 在设备上的储存分为头信息区和存储区两个部分, 头信息区保存 Tensor 的形状
(size)、步长 (stride)、设备 (device)、数据的索引等信息; 存储区 (storage) 保存该 Tensor 的真实物理地址. 得益于
MT-3000 片上的 CPU 可以读写所有 DSP 簇所占用的 DDR 内存, 开发人员只需要修改头信息区的 device 属性便
可以创建一个设备为 CPU 的 Tensor, 由于该 Tensor 存储区指针并没有变, 从而避免了在切换不同后端时候的频
繁内存拷贝.

3.4 计算密集型算子向量化
本文采用向量化优化来处理点对点计算操作 (pointwise), 如加法 (add)、减法 (sub)、除法 (div) 和乘法 (mul).
MT-3000 的 DSP 性能主要依赖于其矢量处理单元 (VPU). 每个 VPU 包含 16 个加速核心 (VPE), 它们以同步的方
式并行工作, 每个核心能够处理 64 位宽的双精度数据, 以充分发挥 VLIW 微架构的计算潜力. 对于单精度数据, 一
个 VPE 能够将 2 个单精度数打包到一个 64 位宽的数据中, 因此一个向量可以同时支持 32 个单精度数的计算. 这
意味着 DSP 具有出色的向量化支持. 在算子向量化实现时, 应该注意算子原生实现应满足以下特点.
(1) 访存地址无冲突或程序中不存在原子操作;
(2) 向量循环体的循环次数可计算且循环体只有一个出口和入口.
算子实现中包含了上述特点即可使用向量化进行优化. 此外, 一个 DSP 簇上的 AM 缓存大小为 768 KB, 因此
在设计实现时应该根据输入大小为 Tensor 分段处理, 防止 AM 内存溢出. 具体实现过程见算法 1.

算法 1. 自适应大小的向量化核函数.

输入: 矩阵 A, 总长度 L;
输出: 矩阵 C.

1. 计算均分给每个核的 Tensor 片段长度 L_c;
2. FOR m=0; L_C; L DO in parallel
3. coreID = get_thread_id();
4. offset=min(coreID×L_c, L-coreID×L_c);
5. cacheLen = MaxVectorCount×32 //单次 load 到 AM 大小;
6. FOR i=0; tmpLen; len – 1 DO
7. tmpLen = (L – i) > cacheLen ? cacheLen : (L – i);
8. vector_load(A[offset+i]) //子矩阵传入 AM 空间;
9. FOR j=0;1; L_c/32 DO
10. Vector C++ execute;