存算一體芯片作為一種新型架構,可以有效解決傳統馮·諾依曼架構在處理人工智能算法時性能和能效存在的訪存瓶頸限制。尤其是基于交叉陣列的存算一體技術,通過在存儲器中原位執行矩陣向量乘,利用存儲單元的固有并行度大幅提升計算性能,并消除部分訪存,能夠成百上千倍地提升深度神經網絡推理的性能和能效。
近些年來,學術界在存算一體的各個方面都進行了大量探索,提出了眾多存算一體加速器架構,中科院微電子所、清華大學、斯坦福大學等單位也制備出了存算一體芯片原型。國內也涌現出了一批存算一體初創企業,包括知存科技、后摩智能、億鑄科技、蘋芯科技等等,它們研發了基于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建模,在事件驅動的引擎中,在仿真過程中高效地進行交互。
利用該套工具鏈,可實現深度神經網絡在存算一體架構上的快速自動化部署,它不僅在部署過程中優化了任務映射和調度,還對生成的指令序列進行相關性能指標的評估。與此同時,該套工具鏈相較于之前的存算一體架構模擬器有了較大的改進。先前的模擬器大多采用數據流架構,能夠支持的網絡結構固定,網絡映射方式單一,而該工具鏈基于抽象的存算一體指令集架構,在編譯器的支持下能夠處理更多的網絡結構,并支持更靈活的計算任務映射方式。編譯器提供了多種預設的編譯優化方案可供選擇,同時也可以修改編譯器代碼實現其它編譯優化方案,通過結合模擬結果的迭代反饋可進一步實現編譯空間探索,尋找針對給定存算一體架構的最優的編譯優化策略,甚至是編譯優化與架構設計結合的軟硬件協同設計。
結語
中國科學院計算技術研究所智能計算機中心所推出的完整開源存算一體工具鏈,可以說是存算一體領域的一項重要突破。這一工具鏈不僅增強了存算一體架構的定制化和靈活性,還促進了深度神經網絡在存算一體架構上的高效自動化部署。
隨著開源理念在存算一體領域的拓展,將有助于行業建立統一的編程和接口標準,從而使來自不同廠商和研究機構的產品實現互通。這一標準化進程將有助于解決目前存算一體芯片領域的碎片化問題,提高生態系統的協同效率。進一步推動存算一體芯片更容易與人工智能、大數據、物聯網等產業相結合,形成更加豐富和復雜的應用場景。