神經(jīng)網(wǎng)絡(luò)構(gòu)造中，算子組成了不同應(yīng)用功能的網(wǎng)絡(luò)結(jié)構(gòu)。而張量加速引擎（Tensor Boost Engine）作為算子的兵工廠，為基于昇騰AI處理器運(yùn)行的神經(jīng)網(wǎng)絡(luò)提供算子開發(fā)能力，用TBE語言編寫的TBE算子來構(gòu)建各種神經(jīng)網(wǎng)絡(luò)模型。同時(shí)，TBE對算子也提供了封裝調(diào)用能力。在TBE中有一個(gè)優(yōu)化過的神經(jīng)網(wǎng)絡(luò)TBE標(biāo)準(zhǔn)算子庫，開發(fā)者可以直接利用標(biāo)準(zhǔn)算子庫中的算子實(shí)現(xiàn)高性能的神經(jīng)網(wǎng)絡(luò)計(jì)算。除此之外，TBE也提供了TBE算子的融合能力，為神經(jīng)網(wǎng)絡(luò)的優(yōu)化開辟一條獨(dú)特的路徑。
張量加速引擎功能框架
TBE提供了基于TVM開發(fā)自定義算子的能力，通過TBE語言和自定義算子編程開發(fā)界面可以完成相應(yīng)神經(jīng)網(wǎng)絡(luò)算子的開發(fā)，TBE的結(jié)構(gòu)如圖所示，包含了特性語言（Domain-Specific Language，DSL）模塊、調(diào)度（Schedule）模塊、中間表示（Intermediate Representation，IR）模塊、編譯器傳遞（Pass）模塊以及代碼生成（CodeGen）模塊。
成.jpg)
TBE算子開發(fā)分為計(jì)算邏輯編寫和調(diào)度開發(fā)，其中特定域語言模塊提供了算子計(jì)算邏輯的編寫接口，直接基于特定域語言編寫算子的計(jì)算過程和調(diào)度過程。算子計(jì)算過程描述指明算子的計(jì)算方法和步驟，而調(diào)度過程描述完成數(shù)據(jù)切塊和數(shù)據(jù)流向的規(guī)劃。算子每次計(jì)算都按照固定數(shù)據(jù)形狀進(jìn)行處理，這就需要提前針對在昇騰AI處理器中的不同計(jì)算單元上執(zhí)行的算子進(jìn)行數(shù)據(jù)形狀切分，如矩陣計(jì)算單元、向量計(jì)算單元以及AI CPU上執(zhí)行的算子對輸入數(shù)據(jù)形狀的需求各不相同。
在完成算子的基本實(shí)現(xiàn)過程定義后，需要啟動調(diào)度模塊中分塊（Tiling）子模塊，對算子中的數(shù)據(jù)按照調(diào)度描述進(jìn)行切分，同時(shí)指定好數(shù)據(jù)的搬運(yùn)流程，確保在硬件上的執(zhí)行達(dá)到最優(yōu)。除了數(shù)據(jù)形狀切分之外，TBE的算子融合和優(yōu)化能力也是由調(diào)度模塊中的融合（Fusion）子模塊提供的。
算子編寫完成后，需要生成中間表示來進(jìn)一步優(yōu)化，而中間表示模塊通過類似于TVM的IR格式來進(jìn)行中間表示的生成。在中間表示生成后，需要將模塊針對各種應(yīng)用場景進(jìn)行編譯優(yōu)化，優(yōu)化的方式有雙緩沖（Double Buffer）、流水線（Pipeline）同步、內(nèi)存分配管理、指令映射、分塊適配矩陣計(jì)算單元等。
在算子經(jīng)過編譯優(yōu)化模塊處理后，由代碼生成模塊生成類C代碼的臨時(shí)文件，這個(gè)臨時(shí)代碼文件可以通過編譯器生成算子的實(shí)現(xiàn)文件，并可以被離線模型執(zhí)行器直接加載和執(zhí)行。
綜上所述，一個(gè)完整的自定義算子通過TBE中的子模塊完成整個(gè)開發(fā)流程，從特定域語言模塊提供算子計(jì)算邏輯和調(diào)度描述，構(gòu)成算子原型后，由調(diào)度模塊進(jìn)行數(shù)據(jù)切分和算子融合，進(jìn)入中間表示模塊，生成算子的中間表示。編譯優(yōu)化模塊以中間表示進(jìn)行內(nèi)存分配等編譯優(yōu)化，最后由代碼生成模塊產(chǎn)生類C代碼可供編譯器直接編譯。TBE在算子的定義過程不但完成了算子編寫，而且還完成了相關(guān)的優(yōu)化，提升了算子的執(zhí)行性能。
