开源之风吹向存算一体芯片

存算一体芯片作为一种新型架构,可以有效解决传统冯·诺依曼架构在处理人工智能算法时性能和能效存在的访存瓶颈限制。尤其是基于交叉阵列的存算一体技术,通过在存储器中原位执行矩阵向量乘,利用存储单元的固有并行度大幅提升计算性能,并消除部分访存,能够成百上千倍地提升深度神经网络推理的性能和能效。

近些年来,学术界在存算一体的各个方面都进行了大量探索,提出了众多存算一体加速器架构,中科院微电子所、清华大学、斯坦福大学等单位也制备出了存算一体芯片原型。国内也涌现出了一批存算一体初创企业,包括知存科技、后摩智能、亿铸科技、苹芯科技等等,它们研发了基于SRAM、闪存、RRAM等存储器的存算一体芯片,且已有产品问世。


(相关资料图)

然而,当前存算一体芯片的发展还处于起步阶段,面临着诸多挑战,其中最重要一点便是生态与编程框架不完善,缺乏相应的指令集与软件工具。

一方面,各单位、公司开发的存算一体芯片均基于自行定义的编程接口,缺乏统一的编程接口,造成了存算一体软件生态的分散,不同厂商开发的上层软件无法互相通用,极大的影响了存算一体芯片的大规模使用。

另一方面,除了高效的硬件设计,神经网络模型面向存算一体架构时,计算任务如何映射、调度,也是发挥神经网络加速器性能和能效的关键。

总体而言,现阶段,行业内尚缺乏神经网络面向存算一体芯片的自动化部署方法和工具。

存算一体芯片有了开源指令集和编译器

我们注意到,在今年的集成电路EDA领域顶级会议IEEE/ACMDesignAutomationConference(DAC)上,中国科学院计算技术研究所智能计算机中心陈晓明和韩银和研究员团队发表论文,公布了一项新的研究成果:PIMCOMP-NN存算一体通用编译器和PIMSIM-NN——存算一体通用模拟器,二者基于一套之前该团队开源的存算一体指令集,构成了完整的开源存算一体工具链,该项工作或将为存算一体芯片建立统一的生态做出一定的贡献。

下面让我们来分析下他们的这项研究。

该论文详细描述了一整套面向神经网络的存算一体工具链,框架如图1所示。这套工具链的核心由两部分组成:

编译器PIMCOMP-NN:编译器PIMCOMP-NN接收ONNX格式的神经网络描述和一组存算一体架构配置参数作为输入,通过编译优化生成指令序列。

模拟器PIMSIM-NN接收编译器:模拟器PIMSIM-NN接收编译器生成的指令序列和架构配置参数(与编译器所用相同)作为输入,通过行为级模拟,得到神经网络在存算一体架构上运行的性能、能耗和功耗。

该套工具基于统一的对软硬件抽象建模而开发,不针对具体的存算一体芯片,也不限制存算一体器件,因此可向上对接不同的神经网络算法,向下屏蔽存算一体芯片硬件细节。

图1面向神经网络的存算一体工具链

该工具链建立在一套面向神经网络的存算一体指令集基础上。指令集作为芯片软硬件设计的接口,对于软件生态有着举足经重的作用,为了使工具链对多种多样的存算一体芯片的底层操作具有通用性,该课题组首先抽象了存算一体架构支持的基本算子,定义了一套统一的面向神经网络的存算一体指令集。

这些指令有四类,分别是矩阵指令(例如矩阵向量乘)、向量指令(例如向量加)、标量指令(主要是传统寄存器操作)和传输指令(主要是各级存储间的数据传输指令),其中每个指令对应硬件的一个抽象操作。面对具体存算一体芯片时,可通过转化程序方便地翻译到具体的硬件指令上。

进一步来看,PIMCOMP-NN是一个面向存算阵列架构的神经网络通用编译框架,它建立在对硬件架构和算子的统一抽象基础之上。为实现对上层算法的通用性,PIMCOMP-NN采用ONNX格式的神经网络描述作为输入,因而可以处理多种多样的神经网络模型。为实现对硬件架构的通用性,PIMCOMP-NN建立在一个抽象的存算架构之上,如图2所示。该抽象架构由多个核心构成,每个核心包括存算一体矩阵单元来完成矩阵-向量乘法运算和向量计算单元来完成向量计算。课题组定义了一组统一的参数来描述该抽象架构。该架构可适配已有研究中广泛采用的多层次Crossbar/PE/Tile/Chip结构,包括发表在ISCA、ASPLOS、ISSCC等多个顶会上的多种存算一体加速器。

图2抽象存算一体架构

PIMCOMP-NN实现了从神经网络结构描述到指令流的自动化生成,并在编译过程中优化任务映射和调度,其框架如图3所示。编译器前端将读取用户提供的抽象架构配置参数(阵列尺寸、核心数目、芯片数目等),同时加载基于ONNX的神经网络模型,通过预处理得到神经网络模型的拓扑关系和参数信息。编译器后端包括四个通用编译优化阶段:节点划分、权重复制、核心映射和数据流调度,这四个阶段均在所定义的伪指令集基础上实施。通过节点划分,灵活地将权重数据进行拆分以适应阵列尺寸。通过权重复制,充分利用存算阵列资源。通过核心映射,高效分配计算任务。通过数据流调度,产生运行完整神经网络的指令流。此外,为了适应不同应用场景,论文设计了两种具有不同层间流水线粒度的编译模式,分别具有高吞吐量和低延迟的特点,在编译过程中分别优化整体吞吐量和推理延迟。

图3PIMCOMP-NN编译器框架

PIMSIM-NN是一个基于指令的存算一体系统模拟器,PIMSIM-NN的整体架构如图4所示,从上到下分别是芯片-核心-基础模块-模拟引擎SystemC,整个芯片由多个核心通过片上互连网络连接形成,核心则由多个内部模块构成,核心内部模块使用基础模块搭建,最终在SystemC的框架下运行。PIMSIM-NN接收指令序列文件和架构配置文件作为输入。指令序列文件包含每个核心的指令序列,由编译器PIMCOMP-NN生成。架构配置文件则对存算一体架构的各项参数进行配置,使模拟器能够仿真不同的存算一体架构。核心设计是PIMSIM-NN的重点,主要有4个处理单元,分别是矩阵单元、向量单元、传输单元和标量单元。矩阵单元的核心是存算阵列单元,其会预先写入权重数据,在运行时原地执行矩阵-向量乘运算,避免权重数据的搬运,降低延迟和功耗。向量单元则负责激活函数和池化等操作,完成一些非线性操作。传输模块则负责核间的数据交换,传输少量的中间结果,并负责核间同步操作。PIMSIM-NN采用了事件驱动仿真模型,基于开源的事件驱动引擎SystemC编写,拥有良好的扩展性并能得到精确的仿真结果。PIMSIM-NN基于图5所示的抽象流水线架构进行仿真,各模块通过SystemC建模,在事件驱动的引擎中,在仿真过程中高效地进行交互。

利用该套工具链,可实现深度神经网络在存算一体架构上的快速自动化部署,它不仅在部署过程中优化了任务映射和调度,还对生成的指令序列进行相关性能指标的评估。与此同时,该套工具链相较于之前的存算一体架构模拟器有了较大的改进。先前的模拟器大多采用数据流架构,能够支持的网络结构固定,网络映射方式单一,而该工具链基于抽象的存算一体指令集架构,在编译器的支持下能够处理更多的网络结构,并支持更灵活的计算任务映射方式。编译器提供了多种预设的编译优化方案可供选择,同时也可以修改编译器代码实现其它编译优化方案,通过结合模拟结果的迭代反馈可进一步实现编译空间探索,寻找针对给定存算一体架构的最优的编译优化策略,甚至是编译优化与架构设计结合的软硬件协同设计。

结语

中国科学院计算技术研究所智能计算机中心所推出的完整开源存算一体工具链,可以说是存算一体领域的一项重要突破。这一工具链不仅增强了存算一体架构的定制化和灵活性,还促进了深度神经网络在存算一体架构上的高效自动化部署。

随着开源理念在存算一体领域的拓展,将有助于行业建立统一的编程和接口标准,从而使来自不同厂商和研究机构的产品实现互通。这一标准化进程将有助于解决目前存算一体芯片领域的碎片化问题,提高生态系统的协同效率。进一步推动存算一体芯片更容易与人工智能、大数据、物联网等产业相结合,形成更加丰富和复杂的应用场景。

关键词: