存算一體架構(gòu)將數(shù)據(jù)存儲單元和計算單元融合為一體,能顯著減少數(shù)據(jù)搬運,極大地提高計算并行度和能效。本次試驗采用知存科技的40nm制程WTM2101存算一體芯片作為計算載體(圖1),該芯片于2022年3月在業(yè)內(nèi)率先實現(xiàn)商用量產(chǎn),采用NOR-Flash非易失存儲器件實現(xiàn)AI權(quán)重存儲和矩陣乘加運算,支持卷積、全連接、Relu等深度神經(jīng)網(wǎng)絡(luò)算子,可以為端側(cè)AI計算提供高能效的算力。
面向WTM2101存算一體芯片計算特性,項目團隊通過算子優(yōu)化技術(shù),將超分模型中的AI算子轉(zhuǎn)換為存算一體芯片支持的算子類型,更好地發(fā)揮存內(nèi)計算優(yōu)勢。針對陣列規(guī)模有限的問題,基于結(jié)構(gòu)重參數(shù)化思想,將帶有局部特征提取算子的多分支卷積結(jié)構(gòu)融合轉(zhuǎn)換為一個3×3卷積層(圖2-a),實現(xiàn)近5倍的參數(shù)量壓縮,得到輕量化超分模型骨干網(wǎng)絡(luò)(圖2-b)。在此基礎(chǔ)上,利用權(quán)重量化技術(shù),將 FP32權(quán)重轉(zhuǎn)換成INT8整數(shù),實現(xiàn)超分模型在存算一體芯片的適配和高效運行,計算能效相比基于傳統(tǒng)馮·諾依曼計算架構(gòu)的12nm制程GPU提升2倍以上。
為了支持視頻超分模型在WTM2101芯片的編譯、部署和推理,項目團隊研發(fā)面向存算一體芯片的軟件計算引擎(圖3),兼容Pytorch、Tensorflow等多種AI框架,提供AI模型編排、部署、推理、管理、驗證、優(yōu)化等全流程服務(wù),有效降低用戶的開發(fā)門檻,提升開發(fā)調(diào)試效率。另外,軟件計算引擎提供了一系列的模型誤差補償技術(shù),有效解決了存算一體芯片模擬計算存在誤差、器件非理性特性等問題,實測顯示視頻超分模型在存算一體芯片上計算的特征圖(feature map)和CPU上計算的特征圖余弦相似度為91.8%,在提升計算能效的同時確保了足夠高的計算精度。圖4為基于存算一體芯片的4倍圖像超分效果。
下一步,中國移動研究院將不斷深耕存算一體領(lǐng)域,一方面發(fā)揮應(yīng)用牽引作用,推動存算一體芯片在算力機頂盒、AR/VR終端等場景落地應(yīng)用;另一方面持續(xù)完善軟件計算引擎功能,助力存算一體軟件生態(tài)構(gòu)建。