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

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

算子名称

其他
激活函数
激活函数反向传播
张量求和
层归一化
批矩阵乘法
归一化指数函数反向传播
张量加减
张量乘除
归一化指数函数
矩阵乘法
图 3 MTTorch 核心算子在训练任务中的计算时长占比情况

3.2 MTTorch 的总体设计思路
MTTorch 本质上是对 PyTorch 官方实现的一个扩展, 通过将 MTTorch 以动态链接库的形式进行挂载, 即可在
MT-3000 上实现异构训练. MTTorch 由设备管理模块和算子库两部分组成, 其整体架构如图 4 所示.

原生 PyTorch
框架前端
用户代码
Python 包 Python 接口
计算图
计算图引擎

核心算子 CUDA 算子
虚拟算子 CPU 算子
MTTorch

MT-3000 DSP 设备管理 MT-3000 DSP 线程组管理
MT-3000 硬件层
Hthreads
MT 数学库 ftIMM
API
图 4 MTTorch 整体架构

设备管理模块保证了深度学习框架后端可以准确、高效地调用 MT-3000 执行深度学习任务. 每当任务被启
动, MTTorch 首先需要通过执行 hthread_dev_open() 初始化设备, 随后利用 hthread_dat_load() 挂载包含设备端核
函数的.dat 文件. 然而, 考虑到在一次模型训练任务中, 设备仅需开启一次. 为了避免不必要的重复操作, 本文特别
设计了基于单例模式的设备管理类来维护一个 DSP 对应的句柄. 单例模式的核心思想在于确保某个类在整个系
统中仅存在一个实例, 并提供全局访问点, 从而防止重复初始化. 当 MTTorch 需要调用设备运算时, 会首先检查
DSP 句柄是否为空. 如果句柄为空, 即设备尚未启动, 则执行相应的 API 来初始化设备并进行内存检查; 若句柄非
空, 则直接跳过启动阶段进入后续的运行流程. 这样的设计不仅提高了效率, 还确保了系统的稳定性和资源的高效
利用.
同样, 当使用 Hthreads 编程模型调用 DSP 计算时, 需要先根据对应的 DSP 句柄和要申请的线程创建线程组,
并返回一个 threadID 以唯一标识该线程组. 然而, hthreads_group_create 函数本身会产生约 1.2 ms 的开销. 为了避
免这种开销, 本文使用单例模式设计了一个线程组管理类, 如图 5 所示, 该类可以缓存已有的 threadID. 如果当前