圖引擎模塊(GE)是MindSpore的一個(gè)子模塊，其代碼由C++實(shí)現(xiàn)，位于前端模塊ME和底層硬件之間，起到承接作用。圖引擎模塊以ME下發(fā)的圖作為輸入，然后進(jìn)行一系列的深度圖優(yōu)化操作，最后輸出一張可以在底層硬件上高效運(yùn)行的圖。GE針對昇騰AI處理器的硬件結(jié)構(gòu)特點(diǎn)，做了特定的優(yōu)化工作，以此來充分發(fā)揮出昇騰AI處理器的強(qiáng)大算力。在進(jìn)行模型訓(xùn)練/推理時(shí)，GE會(huì)被自動(dòng)調(diào)用而用戶并不感知。

GE主要由GE API和GE Core兩部分組成。

GE API

GE API是連接前端模塊ME和GE Core的接口，負(fù)責(zé)GE Core中初始化、Session管理模塊的接口，支持運(yùn)行環(huán)境初始化，Session創(chuàng)建、銷毀，圖添加執(zhí)行。

GE Core

GE Core是GE的核心模塊，負(fù)責(zé)整個(gè)訓(xùn)練過程中的圖管理。GE Core中的圖處理可細(xì)分為六大步驟，分別是圖準(zhǔn)備、圖拆分、圖優(yōu)化、圖編譯、圖加載和圖執(zhí)行，對于ME下發(fā)的每一張圖都會(huì)經(jīng)過這六個(gè)步驟的操作，最終得到可以直接在底層硬件上高效執(zhí)行的圖。

1、圖準(zhǔn)備 & 整圖優(yōu)化

完成整圖級別的數(shù)據(jù)準(zhǔn)備和優(yōu)化，涉及到IR庫及算子庫。使用IR庫中算子的InferShape函數(shù)，完成整圖的Shape推導(dǎo)，以便后續(xù)申請內(nèi)存;同時(shí)根據(jù)算子的聚合屬性，完成某些算子的聚合優(yōu)化，如allreduce算子，會(huì)按照聚合參數(shù)，將若干各參數(shù)對應(yīng)梯度的allreduce算子聚合為一個(gè)，以此減少通訊耗時(shí)。

2、圖拆分

昇騰AI處理器是一種異構(gòu)芯片，含有CPU(AICPU)和向量計(jì)算部件AICORE，圖中每個(gè)算子會(huì)按照開銷模型選擇執(zhí)行的核心，此階段會(huì)對算子進(jìn)行最優(yōu)的核心分配，每種核心對應(yīng)軟件上的一個(gè)抽象引擎;按照之前對各算子的引擎分配，以引擎為邊界，將整圖拆分為若干子圖，在圖邊界算子上插入相應(yīng)的Placeholder算子以做標(biāo)識(shí)，之后的優(yōu)化、編譯、加載操作均會(huì)以子圖為單位進(jìn)行，這樣可以有效減少優(yōu)化過程的耗時(shí)。

3、子圖優(yōu)化

根據(jù)子圖所屬引擎，調(diào)用不同的優(yōu)化器接口執(zhí)行優(yōu)化。為了充分發(fā)揮昇騰AI處理器中AICORE模塊的算力，在AICORE內(nèi)CUBE單元進(jìn)行計(jì)算的算子會(huì)采用一種5D的數(shù)據(jù)格式，圖優(yōu)化階段會(huì)對相應(yīng)算子進(jìn)行4D/5D的類型轉(zhuǎn)換;為了進(jìn)一步發(fā)揮CUBE單元的算力，減少數(shù)據(jù)搬運(yùn)次數(shù)，GE會(huì)對某種范式的算子連接進(jìn)行融合操作，此步驟也在圖優(yōu)化階段進(jìn)行;對所有子圖優(yōu)化之后，需進(jìn)行算子運(yùn)行屬性計(jì)算，以計(jì)算輸入輸出內(nèi)存大小。

4、圖編譯 & 圖加載

GE采用即時(shí)算子編譯技術(shù)，即按照實(shí)際網(wǎng)絡(luò)結(jié)構(gòu)即時(shí)編譯生成算子可執(zhí)行程序，同時(shí)完成內(nèi)存復(fù)用與內(nèi)存分配、流分配、算子可執(zhí)行程序加載等。每個(gè)算子執(zhí)行任務(wù)綁定到特定的流上，同一個(gè)流的任務(wù)是串行執(zhí)行的，不同流上的任務(wù)可以并行執(zhí)行。圖加載階段按照引擎歸屬的runtime，將子圖加載到硬件上準(zhǔn)備執(zhí)行。

5、圖執(zhí)行

最終在硬件上執(zhí)行子圖，并返回相應(yīng)的輸出值。為了提高運(yùn)行效率，圖執(zhí)行階段提供了一種下沉模式，可以在底層硬件上連續(xù)運(yùn)行多輪再返回輸出值，以此減少從底層硬件拷貝數(shù)據(jù)的次數(shù)。

在訓(xùn)練/推理過程中，上述過程會(huì)自動(dòng)執(zhí)行，通過上述圖操作，GE可以將前端下發(fā)的圖轉(zhuǎn)換為一種可以在昇騰AI處理器上高效運(yùn)行的圖模式。

安裝GE

GE內(nèi)嵌在MindSpore安裝包中，MindSpore安裝完畢后，GE以三個(gè)動(dòng)態(tài)庫的方式被調(diào)用。