理論+實踐，揭秘昇騰CANN運算元開發

華為雲開發者聯盟發表於2023-02-17
摘要： 本文介紹了CANN自定義運算元開發的幾種開發方式和運算元的編譯執行流程。然後以開發一個DSL Add運算元為例，講解運算元開發的基本流程。

本文分享自華為雲社群《昇騰CANN運算元開發揭秘》，作者：昇騰CANN 。

開發者在利用昇騰硬體進行神經網路模型訓練或者推理的過程中，可能會遇到以下場景：

  • 訓練場景下，將第三方框架（例如TensorFlow、PyTorch等）的網路訓練指令碼遷移到昇騰AI處理器時遇到了不支援的運算元。
  • 推理場景下，將第三方框架模型（例如TensorFlow、Caffe、ONNX等）使用ATC工具轉換為適配昇騰AI處理器的離線模型時遇到了不支援的運算元。
  • 網路調優時，發現某運算元效能較低，影響網路效能，需要重新開發一個高效能運算元替換效能較低的運算元。
  • 推理場景下，應用程式中的某些邏輯涉及到數學運算（例如查詢最大值，進行資料型別轉換等），希望透過自定義運算元的方式實現這些邏輯，從而獲得效能提升。

此時我們就需要考慮進行自定義運算元的開發，本期我們主要帶您瞭解CANN自定義運算元的幾種開發方式和基本開發流程，讓您對CANN運算元有宏觀的瞭解。

一、運算元基本概念

相信大家對運算元的概念並不陌生，這裡我們來做簡單回顧。深度學習演算法由一個個計算單元組成，我們稱這些計算單元為運算元（Operator，簡稱OP）。

在網路模型中，運算元對應層中的計算邏輯，例如：卷積層（Convolution Layer）是一個運算元；全連線層（Fully-connected Layer， FC layer）中的權值求和過程，是一個運算元。

再例如：tanh、ReLU等，為在網路模型中被用做啟用函式的運算元。

理論+實踐，揭秘昇騰CANN運算元開發

二、CANN自定義運算元開發方式

學習CANN自定義運算元開發方式之前，我們先來了解一下CANN運算元的執行位置：包括AI Core和AI CPU。

  • AI Core是昇騰AI處理器的計算核心，負責執行矩陣、向量、標量計算密集的運算元任務。
  • AI CPU負責執行不適合跑在AI Core上的運算元，是AI Core運算元的補充，主要承擔非矩陣類、邏輯比較複雜的分支密集型計算。

CANN支援使用者使用多種方式來開發自定義運算元，包括TBE DSL、TBE TIK、AICPU三種開發方式。其中TBE DSL、TBE TIK運算元執行在AI Core上，AI CPU運算元執行在AI CPU上。

理論+實踐，揭秘昇騰CANN運算元開發

1. 基於TBE開發框架的運算元開發

TBE（Tensor Boost Engine：張量加速引擎）是CANN提供的運算元開發框架，開發者可以基於此框架使用Python語言開發自定義運算元，透過TBE進行運算元開發有TBE DSL、TBE TIK兩種方式。

  • TBE DSL（Domain-Specific Language ，基於特性域語言）開發方式

為了方便開發者進行自定義運算元開發，CANN預先提供一些常用運算的排程，封裝成一個個運算介面，稱為基於TBE DSL開發。DSL介面已高度封裝，使用者僅需要使用DSL介面完成計算過程的表達，後續的運算元排程、運算元最佳化及編譯都可透過已有的介面一鍵式完成，適合初級開發使用者。

  • TBE TIK（Tensor Iterator Kernel）開發方式

TIK（Tensor Iterator Kernel）是一種基於Python語言的動態程式設計框架，呈現為一個Python模組，執行於Host CPU上。開發者可以透過呼叫TIK提供的API基於Python語言編寫自定義運算元，TIK編譯器會將其編譯為昇騰AI處理器應用程式的二進位制檔案。

TIK需要使用者手工控制資料搬運和計算流程，入門較高，但開發方式比較靈活，能夠充分挖掘硬體能力，在效能上有一定的優勢。

2. AI CPU運算元開發方式

AI CPU運算元的開發介面即為原生C++介面，具備C++程式開發能力的開發者能夠較容易的開發出AI CPU運算元。AI CPU運算元在AI CPU上執行。

下面的開發方式一覽表，對上述幾種開發方式作對比說明，您可以根據各種開發方式的適用場景選擇適合的開發方式。

理論+實踐，揭秘昇騰CANN運算元開發

三、CANN運算元編譯執行

運算元構成

一個完整的CANN運算元包含四部分：運算元原型定義、對應開源框架的運算元適配外掛、運算元資訊庫和運算元實現。這四個組成部分會在運算元編譯執行的過程中使用。

理論+實踐，揭秘昇騰CANN運算元開發

運算元編譯

推理場景下，進行模型推理前，我們需要使用ATC模型轉換工具將原始網路模型轉換為適配昇騰AI處理器的離線模型，該過程中會對網路中的運算元進行編譯。

訓練場景下，當我們跑訓練指令碼時，CANN內部實現邏輯會先將開源框架網路模型下發給Graph Engine進行圖編譯，該過程中會對網路中的運算元進行編譯。

CANN運算元的編譯邏輯架構如下：

理論+實踐，揭秘昇騰CANN運算元開發

具體的CANN運算元編譯流程如下，在編譯流程中會用到上文提到的運算元的四個組成部分。

  1. Graph Engine呼叫運算元外掛，將原始網路模型中的運算元對映為適配昇騰AI處理器的運算元，從而將原始開源框架圖解析為適配昇騰AI處理器的圖。
  2. 呼叫運算元原型庫校驗介面進行基本引數的校驗，校驗透過後，會根據原型庫中的推導函式推導每個節點的輸出shape與dtype，進行輸出tensor的靜態記憶體的分配。
  3. Graph Engine根據圖中資料將圖拆分為子圖並下發給FE。FE在處理過程中根據運算元資訊庫中運算元資訊找到運算元實現，將其編譯成運算元kernel，最後將最佳化後子圖返回給Graph Engine。
  4. Graph Engine進行圖編譯，包含記憶體分配、流資源分配等，並向FE傳送tasking請求，FE返回運算元的taskinfo資訊給Graph Engine，圖編譯完成後生成適配昇騰AI處理器的模型。

運算元執行

推理場景下，使用ATC模型轉換工具將原始網路模型轉換為適配昇騰AI處理器的離線模型後，開發AscendCL應用程式，載入轉換好的離線模型檔案進行模型推理，該過程中會進行運算元的呼叫執行。

訓練場景下，當我們跑訓練指令碼時，內部實現邏輯將開源框架網路模型下發給Graph Engine進行圖編譯後，後續的訓練流程會進行運算元的呼叫執行。

CANN運算元的執行邏輯架構如下：

理論+實踐，揭秘昇騰CANN運算元開發

​具體流程如下，首先Graph Engine下發運算元執行請求給Runtime，然後Runtime會判斷運算元的Task型別，若是TBE運算元，則將運算元執行請求下發到AI Core上執行；若是AI CPU運算元，則將運算元執行請求下發到AI CPU上執行。

四、運算元開發流程

本章節以透過DSL開發方式開發一個Add運算元為例，帶您快速體驗CANN運算元開發的流程。流程圖如下：

理論+實踐，揭秘昇騰CANN運算元開發

運算元開發準備

環境準備：準備運算元開發及執行驗證所依賴的開發環境與執行環境。工程建立：建立運算元開發工程，有以下幾種實現方式：

  • 基於MindStudio工具進行運算元開發，直接使用MindStudio工具建立運算元工程，會自動生成運算元工程及程式碼模板。
  • 基於msopgen工具進行開發，會自動生成運算元工程及程式碼模板。
  • 基於自定義運算元樣例工程進行開發，開發者需要自己建立運算元相關實現檔案，或者基於已有樣例進行修改。

下面以msopgen工具建立運算元開發工程為例進行介紹：

定義AddDSL運算元的原型定義json檔案，用於生成AddDSL的運算元開發工程。例如，定義的json檔案的名字為add_dsl.json，儲存路徑為:$HOME/sample，檔案內容如下：

[
{
 "op":"AddDSL", 
 "input_desc":[
 { 
 "name":"x1", 
 "param_type":"required", 
 "format":[ 
 "NCHW"
 ],
 "type":[ 
 "fp16"
 ]
 },
 { 
 "name":"x2",
 "param_type":"required",
 "format":[
 "NCHW"
 ],
 "type":[
 "fp16"
 ]
 }
 ],
 "output_desc":[ 
 { 
 "name":"y",
 "param_type":"required",
 "format":[
 "NCHW"
 ],
 "type":[
 "fp16"
 ]
 }
 ]
}
]

使用msopgen工具生成AddDSL運算元的開發工程。

$HOME/Ascend/ascend-toolkit/latest/python/site-packages/bin/msopgen gen -i $HOME/sample/add_dsl.json -f tf -c ai_core-Ascend310 -out $HOME/sample/AddDsl

“$HOME/Ascend”為CANN軟體安裝目錄；

“-f tf”引數代表選擇的原始框架為TensorFlow；

“ai_core-<soc_version>”代表運算元在AI Core上執行，<soc_version>為昇騰AI處理器的型號。

此命令執行完後，會在$HOME/sample/AddDsl目錄下生成運算元工程，工程中包含各交付件的模板檔案，編譯指令碼等，如下所示：

AddDsl
├── build.sh                           // 編譯入口指令碼
├── cmake // 編譯解析指令碼存放目錄
├── CMakeLists.txt                      
├── framework                          // 運算元適配外掛相關檔案存放目錄
│   ├── CMakeLists.txt
│   └── tf_plugin
│       ├── CMakeLists.txt
│       └── tensorflow_add_dsl_plugin.cc   // 運算元適配外掛實現檔案
├── op_proto // 運算元原型定義相關檔案存放目錄
│   ├── add_dsl.cc 
│   ├── add_dsl.h 
│   └── CMakeLists.txt
├── op_tiling
│   └── CMakeLists.txt
├── scripts               // 自定義運算元工程打包指令碼存放目錄
└── tbe
    ├── CMakeLists.txt
    ├── impl // 運算元程式碼實現
    │   └── add_dsl.py 
    └── op_info_cfg // 運算元資訊庫存放目錄
        └── ai_core
            └── <soc_version> 
                └── add_dsl.ini

運算元開發過程

實現AddDSL運算元的原型定義。

運算元原型定義檔案包含運算元註冊程式碼的標頭檔案（*.h）以及實現基本校驗、Shape推導的實現檔案（*.cc）。

  • msopgen工具根據add_dsl.json檔案在“op_proto/add_dsl.h”中生成了運算元註冊程式碼，開發者需要檢查自動生成的程式碼邏輯是否正確，一般無需修改。
  • 改“op_proto/add_dsl.cc”檔案，實現運算元的輸出描述推導函式及校驗函式。

在IMPLEMT_COMMON_INFERFUNC(AddDSLInferShape)函式中，填充推導輸出描述的程式碼，針對AddDSL運算元，輸出Tensor的描述資訊與輸入Tensor的描述資訊相同，所以直接將任意一個輸入Tensor的描述賦給輸出Tensor即可。

IMPLEMT_COMMON_INFERFUNC(AddDSLInferShape)
{
 // 獲取輸出資料描述
 TensorDesc tensordesc_output = op.GetOutputDescByName("y");
    tensordesc_output.SetShape(op.GetInputDescByName("x1").GetShape());
    tensordesc_output.SetDataType(op.GetInputDescByName("x1").GetDataType());
    tensordesc_output.SetFormat(op.GetInputDescByName("x1").GetFormat());
 //直接將輸入x1的Tensor描述資訊賦給輸出
 (void)op.UpdateOutputDesc("y", tensordesc_output);
 return GRAPH_SUCCESS;
}

​​在IMPLEMT_VERIFIER(AddDSL, AddDSLVerify)函式中，填充運算元引數校驗程式碼。

IMPLEMT_VERIFIER(AddDSL, AddDSLVerify)
{
 // 校驗運算元的兩個輸入的資料型別是否一致，若不一致，則返回失敗。
 if (op.GetInputDescByName("x1").GetDataType() != op.GetInputDescByName("x2").GetDataType()) {
 return GRAPH_FAILED;
 }
 return GRAPH_SUCCESS;
}

實現AddDSL運算元的計算邏輯。

“tbe/impl/add_dsl.py”檔案中已經自動生成了運算元程式碼的框架，開發者需要在此檔案中修改add_dsl_compute函式，實現此運算元的計算邏輯。

add_dsl_compute函式的實現程式碼如下：

@register_op_compute("add_dsl")
def add_dsl_compute(x1, x2, y, kernel_name="add_dsl"):
 # 呼叫dsl的vadd計算介面
    res = tbe.vadd(x1, x2)
 return res

配置運算元資訊庫。

運算元資訊庫的路徑為“tbe/op_info_cfg/ai_core/<soc_version>/add_dsl.ini”，包含了運算元的型別，輸入輸出的名稱、資料型別、資料排布格式等資訊，msopgen工具已經根據add_dsl.json檔案將上述內容自動填充，開發者無需修改。

AddDSL運算元的資訊庫如下：

[AddDSL] // 運算元的型別 
input0.name=x1              // 第一個輸入的名稱
input0.dtype=float16        // 第一個輸入的資料型別
input0.paramType=required   // 代表此輸入必選，且僅有一個
input0.format=NCHW          // 第一個輸入的資料排布格式
input1.name=x2              // 第二個輸入的名稱
input1.dtype=float16        // 第二個輸入的資料型別
input1.paramType=required   // 代表此輸入必選，且僅有一個
input1.format=NCHW          // 第二個輸入的資料排布格式
output0.name=y              // 運算元輸出的名稱
output0.dtype=float16       // 運算元輸出的資料型別
output0.paramType=required  // 代表此輸出必選，有且僅有一個
output0.format=NCHW         // 運算元輸出的資料排布格式
opFile.value=add_dsl // 運算元實現檔案的名稱
opInterface.value=add_dsl // 運算元實現函式的名稱

實現運算元適配外掛。

運算元適配外掛實現檔案的路徑為“framework/tf_plugin/tensorflow_add_dsl_plugin.cc”，針對原始框架為TensorFlow的運算元，CANN提供了自動解析對映介面“AutoMappingByOpFn”，如下所示：

#include "register/register.h"
namespace domi {
// register op info to GE
REGISTER_CUSTOM_OP("AddDSL") // CANN運算元的型別
 .FrameworkType(TENSORFLOW) // type: CAFFE, TENSORFLOW
 .OriginOpType("AddDSL") // 原始框架模型中的運算元型別
 .ParseParamsByOperatorFn(AutoMappingByOpFn); //解析對映函式
} // namespace domi

以上為工程自動生成的程式碼，開發者僅需要修改.OriginOpType("AddDSL")中的運算元型別即可。此處我們僅展示運算元開發流程，不涉及原始模型，我們不做任何修改。

至此，AddDSL運算元的所有交付件都已開發完畢。

運算元工程編譯及運算元包部署

運算元開發過程完成後，需要編譯自定義運算元工程，生成自定義運算元安裝包並進行自定義運算元包的安裝，將自定義運算元部署到運算元庫。

運算元工程編譯

1. 修改build.sh指令碼，配置運算元編譯所需環境變數。

將build.sh中環境變數ASCEND_TENSOR_COMPILER_INCLUDE配置為CANN軟體標頭檔案所在路徑。

修改前，環境變數配置的原有程式碼行如下：

# export ASCEND_TENSOR_COMPILER_INCLUDE=/usr/local/Ascend/ascend-toolkit/latest/compiler/include

修改後，新的程式碼行如下：

export ASCEND_TENSOR_COMPILER_INCLUDE=${INSTALL_DIR}/include

${INSTALL_DIR}請替換為CANN軟體安裝後檔案儲存路徑。例如，若安裝的Ascend-cann-toolkit軟體包，則安裝後檔案儲存路徑為：$HOME/Ascend/ascend-toolkit/latest。

2. 在運算元工程目錄下執行如下命令，進行運算元工程編譯。

./build.sh

編譯成功後，會在當前目錄下建立build_out目錄，並在build_out目錄下生成自定義運算元安裝包custom_opp_<target os>_<target architecture>.run。

自定義運算元安裝包部署

以執行使用者執行如下命令，安裝自定義運算元包。

./custom_opp_<target os>_<target architecture>.run

命令執行成功後，自定義運算元包中的相關檔案部署到CANN運算元庫中。

運算元執行驗證

運算元包部署完成後，可以進行ST測試（System Test）和網路測試，對運算元進行執行驗證。

ST測試

ST測試的主要功能是：基於運算元測試用例定義檔案*.json生成單運算元的om檔案；使用AscendCL介面載入並執行單運算元om檔案，驗證運算元執行結果的正確性。ST測試會覆蓋運算元實現檔案，運算元原型定義與運算元資訊庫，不會對運算元適配外掛進行測試。

網路測試

你可以將運算元載入到網路模型中進行整網的推理驗證，驗證自定義運算元在網路中執行結果是否正確。網路測試會覆蓋運算元開發的所有交付件，包含實現檔案，運算元原型定義、運算元資訊庫以及運算元適配外掛。

具體的驗證過程請參考“昇騰文件中心”。

以上就是CANN自定義運算元開發的相關知識點，您也可以在“昇騰社群線上課程”板塊學習影片課程，學習過程中的任何疑問，都可以在“昇騰論壇”互動交流！

相關參考：

[1]昇騰文件中心

[2]昇騰社群線上課程

[3]昇騰論壇

 

點選關注，第一時間瞭解華為雲新鮮技術~

相關文章