存算一體芯片作為一種新型架構(gòu),可以有效解決傳統(tǒng)馮·諾依曼架構(gòu)在處理人工智能算法時(shí)性能和能效存在的訪存瓶頸限制。尤其是基于交叉陣列的存算一體技術(shù),通過在存儲(chǔ)器中原位執(zhí)行矩陣向量乘,利用存儲(chǔ)單元的固有并行度大幅提升計(jì)算性能,并消除部分訪存,能夠成百上千倍地提升深度神經(jīng)網(wǎng)絡(luò)推理的性能和能效。
近些年來,學(xué)術(shù)界在存算一體的各個(gè)方面都進(jìn)行了大量探索,提出了眾多存算一體加速器架構(gòu),中科院微電子所、清華大學(xué)、斯坦福大學(xué)等單位也制備出了存算一體芯片原型。國(guó)內(nèi)也涌現(xiàn)出了一批存算一體初創(chuàng)企業(yè),包括知存科技、后摩智能、億鑄科技、蘋芯科技等等,它們研發(fā)了基于SRAM、閃存、RRAM等存儲(chǔ)器的存算一體芯片,且已有產(chǎn)品問世。
然而,當(dāng)前存算一體芯片的發(fā)展還處于起步階段,面臨著諸多挑戰(zhàn),其中最重要一點(diǎn)便是生態(tài)與編程框架不完善,缺乏相應(yīng)的指令集與軟件工具。
一方面,各單位、公司開發(fā)的存算一體芯片均基于自行定義的編程接口,缺乏統(tǒng)一的編程接口,造成了存算一體軟件生態(tài)的分散,不同廠商開發(fā)的上層軟件無法互相通用,極大的影響了存算一體芯片的大規(guī)模使用。
另一方面,除了高效的硬件設(shè)計(jì),神經(jīng)網(wǎng)絡(luò)模型面向存算一體架構(gòu)時(shí),計(jì)算任務(wù)如何映射、調(diào)度,也是發(fā)揮神經(jīng)網(wǎng)絡(luò)加速器性能和能效的關(guān)鍵。
總體而言,現(xiàn)階段,行業(yè)內(nèi)尚缺乏神經(jīng)網(wǎng)絡(luò)面向存算一體芯片的自動(dòng)化部署方法和工具。
存算一體芯片有了開源指令集和編譯器
我們注意到,在今年的集成電路EDA領(lǐng)域頂級(jí)會(huì)議IEEE/ACMDesignAutomationConference(DAC)上,中國(guó)科學(xué)院計(jì)算技術(shù)研究所智能計(jì)算機(jī)中心陳曉明和韓銀和研究員團(tuán)隊(duì)發(fā)表論文,公布了一項(xiàng)新的研究成果:PIMCOMP-NN存算一體通用編譯器和PIMSIM-NN——存算一體通用模擬器,二者基于一套之前該團(tuán)隊(duì)開源的存算一體指令集,構(gòu)成了完整的開源存算一體工具鏈,該項(xiàng)工作或?qū)榇嫠阋惑w芯片建立統(tǒng)一的生態(tài)做出一定的貢獻(xiàn)。
下面讓我們來分析下他們的這項(xiàng)研究。
該論文詳細(xì)描述了一整套面向神經(jīng)網(wǎng)絡(luò)的存算一體工具鏈,框架如圖1所示。這套工具鏈的核心由兩部分組成:
編譯器PIMCOMP-NN:編譯器PIMCOMP-NN接收ONNX格式的神經(jīng)網(wǎng)絡(luò)描述和一組存算一體架構(gòu)配置參數(shù)作為輸入,通過編譯優(yōu)化生成指令序列。
模擬器PIMSIM-NN接收編譯器:模擬器PIMSIM-NN接收編譯器生成的指令序列和架構(gòu)配置參數(shù)(與編譯器所用相同)作為輸入,通過行為級(jí)模擬,得到神經(jīng)網(wǎng)絡(luò)在存算一體架構(gòu)上運(yùn)行的性能、能耗和功耗。
該套工具基于統(tǒng)一的對(duì)軟硬件抽象建模而開發(fā),不針對(duì)具體的存算一體芯片,也不限制存算一體器件,因此可向上對(duì)接不同的神經(jīng)網(wǎng)絡(luò)算法,向下屏蔽存算一體芯片硬件細(xì)節(jié)。
圖1面向神經(jīng)網(wǎng)絡(luò)的存算一體工具鏈
該工具鏈建立在一套面向神經(jīng)網(wǎng)絡(luò)的存算一體指令集基礎(chǔ)上。指令集作為芯片軟硬件設(shè)計(jì)的接口,對(duì)于軟件生態(tài)有著舉足經(jīng)重的作用,為了使工具鏈對(duì)多種多樣的存算一體芯片的底層操作具有通用性,該課題組首先抽象了存算一體架構(gòu)支持的基本算子,定義了一套統(tǒng)一的面向神經(jīng)網(wǎng)絡(luò)的存算一體指令集。
這些指令有四類,分別是矩陣指令(例如矩陣向量乘)、向量指令(例如向量加)、標(biāo)量指令(主要是傳統(tǒng)寄存器操作)和傳輸指令(主要是各級(jí)存儲(chǔ)間的數(shù)據(jù)傳輸指令),其中每個(gè)指令對(duì)應(yīng)硬件的一個(gè)抽象操作。面對(duì)具體存算一體芯片時(shí),可通過轉(zhuǎn)化程序方便地翻譯到具體的硬件指令上。
進(jìn)一步來看,PIMCOMP-NN是一個(gè)面向存算陣列架構(gòu)的神經(jīng)網(wǎng)絡(luò)通用編譯框架,它建立在對(duì)硬件架構(gòu)和算子的統(tǒng)一抽象基礎(chǔ)之上。為實(shí)現(xiàn)對(duì)上層算法的通用性,PIMCOMP-NN采用ONNX格式的神經(jīng)網(wǎng)絡(luò)描述作為輸入,因而可以處理多種多樣的神經(jīng)網(wǎng)絡(luò)模型。為實(shí)現(xiàn)對(duì)硬件架構(gòu)的通用性,PIMCOMP-NN建立在一個(gè)抽象的存算架構(gòu)之上,如圖2所示。該抽象架構(gòu)由多個(gè)核心構(gòu)成,每個(gè)核心包括存算一體矩陣單元來完成矩陣-向量乘法運(yùn)算和向量計(jì)算單元來完成向量計(jì)算。課題組定義了一組統(tǒng)一的參數(shù)來描述該抽象架構(gòu)。該架構(gòu)可適配已有研究中廣泛采用的多層次Crossbar/PE/Tile/Chip結(jié)構(gòu),包括發(fā)表在ISCA、ASPLOS、ISSCC等多個(gè)頂會(huì)上的多種存算一體加速器。
圖2抽象存算一體架構(gòu)
PIMCOMP-NN實(shí)現(xiàn)了從神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)描述到指令流的自動(dòng)化生成,并在編譯過程中優(yōu)化任務(wù)映射和調(diào)度,其框架如圖3所示。編譯器前端將讀取用戶提供的抽象架構(gòu)配置參數(shù)(陣列尺寸、核心數(shù)目、芯片數(shù)目等),同時(shí)加載基于ONNX的神經(jīng)網(wǎng)絡(luò)模型,通過預(yù)處理得到神經(jīng)網(wǎng)絡(luò)模型的拓?fù)潢P(guān)系和參數(shù)信息。編譯器后端包括四個(gè)通用編譯優(yōu)化階段:節(jié)點(diǎn)劃分、權(quán)重復(fù)制、核心映射和數(shù)據(jù)流調(diào)度,這四個(gè)階段均在所定義的偽指令集基礎(chǔ)上實(shí)施。通過節(jié)點(diǎn)劃分,靈活地將權(quán)重?cái)?shù)據(jù)進(jìn)行拆分以適應(yīng)陣列尺寸。通過權(quán)重復(fù)制,充分利用存算陣列資源。通過核心映射,高效分配計(jì)算任務(wù)。通過數(shù)據(jù)流調(diào)度,產(chǎn)生運(yùn)行完整神經(jīng)網(wǎng)絡(luò)的指令流。此外,為了適應(yīng)不同應(yīng)用場(chǎng)景,論文設(shè)計(jì)了兩種具有不同層間流水線粒度的編譯模式,分別具有高吞吐量和低延遲的特點(diǎn),在編譯過程中分別優(yōu)化整體吞吐量和推理延遲。
圖3PIMCOMP-NN編譯器框架
PIMSIM-NN是一個(gè)基于指令的存算一體系統(tǒng)模擬器,PIMSIM-NN的整體架構(gòu)如圖4所示,從上到下分別是芯片-核心-基礎(chǔ)模塊-模擬引擎SystemC,整個(gè)芯片由多個(gè)核心通過片上互連網(wǎng)絡(luò)連接形成,核心則由多個(gè)內(nèi)部模塊構(gòu)成,核心內(nèi)部模塊使用基礎(chǔ)模塊搭建,最終在SystemC的框架下運(yùn)行。PIMSIM-NN接收指令序列文件和架構(gòu)配置文件作為輸入。指令序列文件包含每個(gè)核心的指令序列,由編譯器PIMCOMP-NN生成。架構(gòu)配置文件則對(duì)存算一體架構(gòu)的各項(xiàng)參數(shù)進(jìn)行配置,使模擬器能夠仿真不同的存算一體架構(gòu)。核心設(shè)計(jì)是PIMSIM-NN的重點(diǎn),主要有4個(gè)處理單元,分別是矩陣單元、向量單元、傳輸單元和標(biāo)量單元。矩陣單元的核心是存算陣列單元,其會(huì)預(yù)先寫入權(quán)重?cái)?shù)據(jù),在運(yùn)行時(shí)原地執(zhí)行矩陣-向量乘運(yùn)算,避免權(quán)重?cái)?shù)據(jù)的搬運(yùn),降低延遲和功耗。向量單元?jiǎng)t負(fù)責(zé)激活函數(shù)和池化等操作,完成一些非線性操作。傳輸模塊則負(fù)責(zé)核間的數(shù)據(jù)交換,傳輸少量的中間結(jié)果,并負(fù)責(zé)核間同步操作。PIMSIM-NN采用了事件驅(qū)動(dòng)仿真模型,基于開源的事件驅(qū)動(dòng)引擎SystemC編寫,擁有良好的擴(kuò)展性并能得到精確的仿真結(jié)果。PIMSIM-NN基于圖5所示的抽象流水線架構(gòu)進(jìn)行仿真,各模塊通過SystemC建模,在事件驅(qū)動(dòng)的引擎中,在仿真過程中高效地進(jìn)行交互。
利用該套工具鏈,可實(shí)現(xiàn)深度神經(jīng)網(wǎng)絡(luò)在存算一體架構(gòu)上的快速自動(dòng)化部署,它不僅在部署過程中優(yōu)化了任務(wù)映射和調(diào)度,還對(duì)生成的指令序列進(jìn)行相關(guān)性能指標(biāo)的評(píng)估。與此同時(shí),該套工具鏈相較于之前的存算一體架構(gòu)模擬器有了較大的改進(jìn)。先前的模擬器大多采用數(shù)據(jù)流架構(gòu),能夠支持的網(wǎng)絡(luò)結(jié)構(gòu)固定,網(wǎng)絡(luò)映射方式單一,而該工具鏈基于抽象的存算一體指令集架構(gòu),在編譯器的支持下能夠處理更多的網(wǎng)絡(luò)結(jié)構(gòu),并支持更靈活的計(jì)算任務(wù)映射方式。編譯器提供了多種預(yù)設(shè)的編譯優(yōu)化方案可供選擇,同時(shí)也可以修改編譯器代碼實(shí)現(xiàn)其它編譯優(yōu)化方案,通過結(jié)合模擬結(jié)果的迭代反饋可進(jìn)一步實(shí)現(xiàn)編譯空間探索,尋找針對(duì)給定存算一體架構(gòu)的最優(yōu)的編譯優(yōu)化策略,甚至是編譯優(yōu)化與架構(gòu)設(shè)計(jì)結(jié)合的軟硬件協(xié)同設(shè)計(jì)。
結(jié)語
中國(guó)科學(xué)院計(jì)算技術(shù)研究所智能計(jì)算機(jī)中心所推出的完整開源存算一體工具鏈,可以說是存算一體領(lǐng)域的一項(xiàng)重要突破。這一工具鏈不僅增強(qiáng)了存算一體架構(gòu)的定制化和靈活性,還促進(jìn)了深度神經(jīng)網(wǎng)絡(luò)在存算一體架構(gòu)上的高效自動(dòng)化部署。
隨著開源理念在存算一體領(lǐng)域的拓展,將有助于行業(yè)建立統(tǒng)一的編程和接口標(biāo)準(zhǔn),從而使來自不同廠商和研究機(jī)構(gòu)的產(chǎn)品實(shí)現(xiàn)互通。這一標(biāo)準(zhǔn)化進(jìn)程將有助于解決目前存算一體芯片領(lǐng)域的碎片化問題,提高生態(tài)系統(tǒng)的協(xié)同效率。進(jìn)一步推動(dòng)存算一體芯片更容易與人工智能、大數(shù)據(jù)、物聯(lián)網(wǎng)等產(chǎn)業(yè)相結(jié)合,形成更加豐富和復(fù)雜的應(yīng)用場(chǎng)景。