如何設計一個高內聚低耦合的模組——MegEngine 中自定義 Op 系統的實踐經驗

MegEngine發表於2021-11-18

作者：褚超群 | 曠視科技 MegEngine 架構師

背景介紹

在演算法研究的過程中，演算法同學們可能經常會嘗試定義各種新的神經網路層（neural network layer），比如 Layer Norm，Deformable Conv 等。為了實現這些層以進行實驗，演算法同學可以使用神經網路框架或者 numpy 中提供的基礎操作（如張量/標量的加減乘除等）去組合出所需的層的功能。然而這通常會造成這些層的效能斷崖式的下跌，大大影響了演算法同學們嘗試新演算法的效率。所以很多情況下，演算法同學們會選擇為自己定義的層實現高效能的 kernel，並希望可以將之整合入框架作為框架中的 Op（Operator） 使用。

然而在一般情況下，演算法同學必須要對框架本身十分了解，才可以靈活自由的將我們的 kernel 接入框架中去使用。但這並不是一件簡單的事，神經網路框架作為一個規模龐大的系統，其結構之複雜遠超常規的軟體專案。為了保證這樣的機器學習系統的可維護性以及可擴充性，系統中往往會做各種各樣的層次，模組設計，並對各種概念（比如 Op）進行抽象，各個層次各個模組間的互動又非常複雜。

以 MegEngine 中的 Op 系統為例，圖 1 中展示 Op 這一最基本的概念在 MegEngine 中不同層的各種抽象。

圖 1：MegEngine 中不同層次的 Operator 的抽象

  • 在底層的 MegDNN 運算元庫中，Op 被抽象成了 MegDNNOpr 類，其中封裝了各個 Op 在 x86，nvidia gpu 等硬體平臺上的具體 kernel 實現以及相關硬體的 context 管理。
  • 在靜態圖（graph runtime）中，Op 被抽象成了 OpNode 類，其主要目的並非是用於計算而是圖優化，故而這個資料結構的設計上又有了相當多這方面的考量。
  • 在動態圖（imperative runtime）中，Op 又會被抽象為 OpDef 類，配合動態圖直譯器進行任務的排程。
  • 在 python 中，Op 會被封裝成 functional 和 module，這才是符合一般演算法同學認知的 Op。

而從 python 中執行一個操作時，這些 Op 會逐層向下呼叫，分別在每一層完成一部分工作，直到最後才呼叫了 MegDNN 運算元庫中具體的 kernel。這個過程中，任何一個層次的 Op 概念都是缺一不可的。其實不只是 Op，包括 Tensor 在內的很多其他概念，在 MegEngine 系統中都存在著類似的多種抽象。學習瞭解這樣的一個框架設計需要花費大量的時間和精力，這個代價往往是演算法同學難以接受的。

然而即使演算法同學犧牲了大把時間和頭髮，學習瞭解了 MegEngine 的系統設計，完成了自己 kernel 到 MegEngine 的整合，事情還遠遠沒有結束。kernel 的整合過程通常與框架本身是高度耦合的。其構建 Op 的過程需要獲取框架的所有原始碼，修改編譯框架中的絕大多數模組。如果之後框架內部的相關抽象發生變化，則之前構建的 Op 則又變為不可用的狀態。

為了允許把演算法同學的 kernel 快速的整合入框架去進行使用，並且整合出來的 Op 既可以與框架內的原生 Op 有著一致的行為，同時其又與框架本身相解耦，MegEngine 提出了一套工具 Custom Op。其可以很簡單便捷的將演算法同學自己編寫的 c++/cuda kernel 封裝成 Op 並自動化的編譯成動態連結庫並整合入 MegEngine 中。

然而，編寫高效能的 c++/cuda kernel 對於一般沒有體系結構/平行計算背景的演算法研究人員依然是一件很困難的事情。所以為了避免演算法同學自行編寫 kernel 的種種問題，MegEngine 基於 Custom Op 進一步提出了 Custom Op Generator，嘗試利用神經網路編譯器程式碼生成的方式去自動化端到端的生成 kernel 和 Custom Op 程式碼，並將之整合入 MegEngine，使演算法同學無需編寫任何 c++ 程式碼，即可在 MegEngine 中新增高效能的 kernel 並使用。

正常的 Op 整合與 Custom Op

為了便於理解 Custom Op 設計，我們首先對傳統 Op 的整合過程和 Custom Op 的整合過程進行對比，建立其 Custom Op 的初步印象。

一般而言，我們的演算法同學想將自己編寫的 c++/cuda kernel 整合為 MegEngine 的 Op，那麼他首先必須瞭解：

  • MegEngine 整個的系統結構。
  • MegEngine 中各種層次模組的功能。
  • 諸如 Op，Tensor 等概念在不同層次模組的設計目的以及實質含義。

在對這樣一個系統有了十足瞭解之後，其需要：

  • 在 MegDNN 運算元層中對自己的 kernel 進行封裝，將之封裝成 MegDNNOpr 類。
  • 基於 MegEngine 中靜態圖中的相關元件將自己的 Op 封裝成 OpNode 類。
  • 基於 MegEngine 中動態圖中的相關元件將自己的 Op 封裝成 OpDef 類。
  • 編寫 python 和 C++ 的互動程式碼，將自己的 Op 暴露到 python 環境中。

而為了簡化這個過程，Custom Op 中提供了一套框架無關的十分簡潔的 Op 的模型，演算法使用者在新增 Op 時無需對框架本身做任何瞭解。其唯一需要做的就是基於這套模型去設定一些 Op 的基本資訊，比如 Op 有幾個輸入輸出，呼叫哪個 kernel 等，從而建立起關於自己 Op 的描述，其畫風一般如下程式碼所示。

CUSTOM_OP_REG(MatMulScale)              // 定義一個名為 MatMulScale 的 Op
     .add_inputs(2)                  // 兩個輸入 Tensor
     .add_outputs(1)                  // 一個輸出 Tensor
     .add_param("scale", 1.0f)           // 一個名為 scale 的 Parameter，預設值為 1.0f
     .set_shape_infer(shape_infer)         // 設定這個 Op 的 shape 推導函式
     .set_compute("cuda", compute);        // 設定這個 Op 的 計算函式

而這些設定過程一般使用幾行至十幾行的程式碼就可以表達，大大簡化了使用者整合 kernel 時的工作量。更多的關於 Custom Op 的使用可以參考 MegEngine Custom Op 使用說明。

然後 Custom Op 會自動的將使用者的 Op 封裝成 MegEngine 中靜態圖與動態圖中的 OpNode 與 OpDef，併為之生成和原生運算元一致的 python 介面，從而使使用者的 Custom Op 可以與系統中的原生 Op 在介面和底層行為上保持統一。

Custom Op 的設計

Custom Op 同時面向使用者（即上述需要編寫 kernel 的演算法同學）與 MegEngine 系統，使兩者可以簡單便捷的進行互動。而為了達到這個目的，Custom Op 需要具備以下特性：

  • 面向使用者，Custom Op 需要提供一套簡潔統一的，與框架無關的，且編寫 Op 所必須使用的抽象概念。使用者可以使用這些抽象介面去將自己的 c++/cuda kernel 封裝成 Op。
  • 面向系統，Custom Op 需要給 MegEngine 提供一套完備的 Op 的適配與管理工具，從而允許系統對 Custom Op 進行呼叫及維護管理。

基於此，我們設計出了 Custom Op 的整體架構，如圖 2 所示。

圖 2：Custom Op 的整體結構

在一般演算法使用者的認知裡，Op 就是個計算函式，接受一些輸入，完成計算，然後得到一些輸出。而這裡的輸入輸出又分為張量型資料（即一般的輸入輸出 Tensor）和非張量型資料（即 Param，比如卷積中的 padding，stride 等等），如圖 3 所示。

圖 3：使用者視角的 Op

為了契合演算法使用者對於 Op 的認知，Custom Op 主要向使用者提供了三個抽象 Tensor，Param 以及 Op。

  • Tensor 是 Op kernel 主要計算和操作的物件，與 MegEngine python 中的 Tensor 有著基本一致的抽象和行為。
  • Param 是用於記錄傳輸一些 Op 的非張量的輸入（比如卷積中的 padding，stride 等等）。
  • Op 是對使用者的 c++/cuda kernel 計算函式的一個封裝，同時記錄著這個 Op 的輸入輸出 Tensor 以及 Param 的資訊。

而面向 MegEngine 系統，Custom Op 一方面給 MegEngine 提供了一套完備的 Adaptors，可以根據 MegEngine 系統中不同層次的需要，將使用者的 Op 和 Tensor 適配成 MegEngine 的 runtime 可以處理的 Op 和 Tensor。另外一方面，Custom Op 同時還向 MegEngine 提供一套使用者 Op 的 Managers，從而允許 MegEngine 對 Custom Op 進行維護管理。

下面我們將分別介紹 Custom Op 的這些模組。

Tensor

在演算法使用者的視角中，Tensor 是一個多維陣列，同時其還有著一些如形狀，量化資訊等屬性，故而 Custom Op 中的 Tensor 也被設計為資料以及資料的相關屬性的集合。 其中資料由一個指向資料儲存空間的指標管理。而這些屬性則告訴我們該如何去解析資料空間中的資料。如圖2中 Tensor 部分所示，這些屬性主要包含 Device，DType，Shape 這三者：

  • Shape 代表的是 Tensor 維度資訊。
  • DType 對應 Tensor 中元素的資料型別，如 float32，uint8 等。
  • Device 則表示這個 Tensor 在什麼裝置（cpu/gpu）上。

通過這些屬性可以建立起對 Tensor 的完備的描述。

事實上 Tensor 及其附屬的這些屬性在 MegEngine 系統都會有另一套功能豐富，但對使用者而言略顯冗餘的表達。為了降低使用難度，Custom Op 簡化了這些概念，只留下編寫 Op 時需要的功能。

具體來說，Shape 提供給使用者的行為類似於 c++ 的原生陣列，我們可以使用如下的程式碼來構建和使用它：

Shape shape = {16, 3, 224, 224};    // 構建 shape
bool equal = (shape[3] == 224);     // 獲取 shape 特定維度的值
shape = {128, 100};                 // 修改 shape 的值

至於 Device 以及 DType，使用者並不需要知曉其背後實現，只需要通過這些屬性知道資料在什麼裝置上，是什麼型別就夠了。故而 Custom Op 中的 Device 和 DType 的行為均類似於 string，使用者可以以直接設定字串值的形式去設定具體的 Device 和 DType 型別。

Device device = "x86";                 // 建立一個 x86 這種裝置型別
device = "cuda";                     // 裝置型別改為 cuda
bool equal = (device == "cuda");            // 判斷某個 device 是否是 cuda

DType dtype = "float32";                       // 構建 dtype
bool equal = (dtype == "int8");           // 判斷 dtype 是否相等

而為了使這些型別的介面與 MegEngine 解耦，其實現時均使用了 pimpl 手法，隱藏了這些型別的記憶體佈局，而使用者通過一系列的介面去 set/get 其中的資料。

Param

Param 主要用於表達 Op 的非 Tensor 的輸入（比如卷積中的 padding，stride 等等），然而不同的 Op 的這些非 Tensor 輸入的差別往往非常大。可能 Op A 的 Param 是一系列的 string 而 Op B 的 Param 是一個 int 型別。所以我們需要設計一套機制以將這些彼此差異很大的 Param 統一起來，供使用者和 MegEngine 系統進行使用。為了實現這個目的，Custom Op 中設計了 ParamVal。ParamVal 中會擦除各個 Op Param 的靜態型別，使這些 Param 靜態型別統一，以解決不同 Op 的 Param 型別不一致的問題，然後另外定義一套執行時的動態型別系統去進行 Param 實際型別的管理。

說起來可能比較複雜，實際上其可以簡化成下面的這個資料結構。其中使用 void* 進行型別擦除，並將擦除後的資料放在 data 中進行儲存，而 type 中則記錄著這個資料所對應的動態型別。

class ParamVal {
    void *data;             // 型別擦除後的資料
   DynamicDataType type;   // 資料的動態型別
};

這樣的設計不僅能解決不同 Param 型別不統一的問題，同時這個動態型別的存在同時也大大緩解了 c++ 中沒有反射所帶來的一些不便。Custom Op 根據此動態型別系統設計了一套統一的引數解析 （Param Parse） 和序列化機制，而無需使用者為自己的 Param 去編寫這部分程式碼。

同時為了進一步方便使用者去使用，Custom Op 中還為 ParamVal 定義了非常多的運算子，以及其與靜態型別之間相互轉換的函式。最終展現在使用者視角，ParamVal 的行為與 python 中的變數是很接近的。

ParamVal a = 1.0, b = 2, c = false, d = {1, 2}; // 表達不同型別的資料
ParamVal e = a + b;                             // ParamVal 彼此間的計算
ParamVal f = e - 2;                             // ParamVal 與靜態型別資料的計算
d = c;

Op

在一般演算法使用者的視角里，Op 是對一個計算過程的描述而不記錄儲存任何資料資訊。為了與演算法使用者的認知相統一，Custom Op 中的 Op，也被設計為無狀態的，即 Op 中只儲存相關的函式以及輸入輸出的一些佈局資訊，而不記錄輸入輸出的具體值。

圖 4：Op 及其元件

具體來說，在 Custom Op 中，Op 是輸入輸出 Tensor 資訊（TensorInfo），Param 資訊（ParamInfo），以及 Op 相關函式的集合（Functions），如圖4所示。TensorInfo 中記錄了這個 Op 的輸入輸出 Tensor 的數量，名字，合法型別，維度以及記憶體分配策略等資訊。ParamInfo 則記錄著這個 Op 的各個 Param 的名字，預設型別以及預設值等。至於 Functions 其實際包含兩部分，kernel 計算函式和 Tensor 屬性推導函式。

  • kernel 計算函式主要負責將 Tensor，Param 的值傳遞給使用者的 c++/cuda kernel，並將計算結果返回。
  • Tensor 屬性推導函式則是根據輸入 Tensor/Param 的一些屬性在 kernel 執行前完成輸出 Tensor 的資料佈局的推導，從而將運算元執行和運算元記憶體分配解耦，以進行記憶體規劃。

其中大部分函式 Custom Op 均提供了預設的實現，使用者可以通過 override 這些函式的預設行為去定製化自己的 Op。

Manager 與 Adaptor

考慮到 Custom Op 是一套與 MegEngine 系統解耦的 Op 抽象，根據 Custom Op 定義出來的 Op，MegEngine 並不能直接與之進行互動。為了解決這個問題，Custom Op 中額外設計了一組 Adaptors 和 Managers。其中 Adaptor 允許 MegEngine 使用 Custom Op，而 Manager 允許 MegEngine 感知和管理 Custom Op。

Adaptor 的設計目的主要包含兩個方面：一方面其需要允許 MegEngine 去操作使用 Custom Op 去構建網路等，另一方面其需要允許 MegEngine 與 Custom Op 進行資料互動完成計算。

  • 對於前者，Adaptors 可以將 Custom Op 包裝成 MegEngine 中動態圖與靜態圖中的 Op 抽象，使之行為與 MegEngine 中內建 Op 保持一致。
  • 對於後者，Adaptors 可以將 Custom Op 中面向使用者的 Tensor 抽象與 MegEngine 中的 Tensor 結合起來，使兩者間可以互相轉換，允許資料可以自由的在 MegEngine 與 Custom Op 之間流動。

而對於 Manager 模組，其提供了對 Custom Op 所編譯成的動態連結庫以及 Custom Op 本身的管理。具體來說，Custom Op 在使用時會以動態連結庫的形式被載入入 MegEngine 系統中，因此我們基於 RAII 機制去管理這些動態連結庫。動態庫的載入與解除安裝和 Lib 類的構造與析構繫結起來，從而避免資源洩漏。而對於 Op 本身，Manager 提供了一些基本的增刪改查的操作去允許 MegEngine 對 Custom Op 進行管理。

Custom Op 的編寫

使用者在使用 Custom Op 編寫 Op 時，使用者可以使用上述這些概念去將自己寫好的 kernel 封裝成 Custom Op，並使用 Custom Op 提供的構建工具將之編譯成動態連結庫，在執行時將之載入入 MegEngine 進行使用。

具體來說，假如現在我們需要為 MegEngine 新增一個名為 MatmulScale 的運算元，這個運算元在計算時首先會對兩個輸入 Tensor，lhs 和 rhs 執行矩陣乘，然後再將這個矩陣乘的結果再乘以標量 Scale。

該運算元數學上的執行過程的虛擬碼如下：

def MatMulScale(lhs, rhs, scale):
    result = lhs.dot(rhs)
    result = result * scale
    return result

對於這樣的一個操作，假設我們已經為之寫好了一份 cuda kernel 程式碼，並提供如下的介面函式用於呼叫：

void matmul_scale(const float *lhs, const float *rhs, float *result, size_t M, size_t K, size_t N, float scale);

這些的引數中，lhs，rhs，以及 result 是三個 float 型別的指標， 分別代表這個 Op 的兩個輸入 Tensor 和一個輸出 Tensor，其均需要指向一片已經分配好的 cuda memory。 而 M，K，N 是矩陣的維度資訊，表示一個 M*K 的矩陣乘以一個 K*N 的矩陣。 而 scale 則代表著矩陣乘的結果需要乘以的那個係數。

對於這種情況我們可以編寫如下的 c++ 程式碼，就可以將之封裝成 MegEngine 的 Op。

void shape_infer(const std::vector<Shape> &inputs, const Param &params, std::vector<Shape> &outputs) {
    outputs[0] = {inputs[0][0], inputs[1][1]};
}

void compute(const std::vector<Tensor> &inputs, const Param &params, std::vector<Tensor> &outputs) {
    matmul_scale(inputs[0].data<float>(), inputs[1].data<float>(), outputs[0].data<float>(), ...); 
}

 CUSTOM_OP_REG(MatMulScale)              // 定義一個名為 MatMulScale 的 Op
     .add_inputs(2)                  // 兩個輸入 Tensor
     .add_outputs(1)                  // 一個輸出 Tensor
     .add_param("scale", 1.0f)           // 一個名為 scale 的 Parameter，預設值為 1.0f
     .set_shape_infer(shape_infer)         // 設定這個 Op 的 shape 推導函式
     .set_compute("cuda", compute);        // 設定這個 Op 的 計算函式

這段程式碼主要包含兩個部分，第一個部分是一些函式的定義，包括輸出 Tensor 屬性推斷函式和計算函式。 其中前者會根據輸入 Tensor 的屬性（比如 shape）去推導輸出 Tensor 的對應屬性，而後者則是在其中呼叫 cuda kernel，完成計算。 第二部分是 Op 的註冊，主要用於定義 Op 有幾個輸入輸出 Tensor，有幾個 Param，並將上面定義的屬性推斷函式和計算函式的指標也註冊給 Op。到此就完成了 Custom Op 的構建。

Custom Op Generator

通過 Custom Op 可以將使用者編寫好的 c++/cuda kernel 簡單方便的整合入 MegEngine。然而，讓使用者自己去編寫 kernel 總是最後的選擇，如何能夠將使用者編寫 kernel 的這一步的工作也給省掉呢？MegEngine 正在嘗試基於 AI 編譯器提出一個工具 Custom Op Generator 去解決這個問題。

在 Custom Op Generator 中使用者可以直接使用 AI 編譯器提供的簡單的 python 原語去建立其 Op 的表達，而不需要寫任何 c++/cuda 的程式碼。然後 AI 編譯器會自動生成 Op 所對應的 kernel 以及 Custom Op 的裝飾程式碼將這個 kernel 封裝成 MegEngine 的 Custom Op，並自動化的進行構建並將之整合入 MegEngine中。全過程使用者只需編寫一些 python 程式碼即可，避免了使用者自己編寫 kernel 的問題。

一般而言，框架與編譯器結合的 workflow 都是框架在前編譯器在後，由框架去構建一個模型用於訓練，然後將訓練好的模型送給編譯器進行優化部署。然而在 Custom Op Generator 中兩者的位置恰恰相反，而在這樣的一個 workflow 中，編譯器在前而框架在後，編譯器利用其程式碼生成的能力為框架提供擴充支援。從某種意義上來說，這是一種 AI 編譯器與框架結合的一種新思路。

總結

Custom Op 作為溝通使用者 kernel 和系統 Op 的橋樑，其面向使用者提供了一套簡潔統一且與框架無關的 Op 抽象，面向系統提供一套完備的 Op 適配與管理工具。為了這個目的，在設計實現 Custom Op 的過程中，我們分析使用者編寫 Op 時所必須的概念並進行抽象，設計出了一套框架無關的 Op 模型並提供簡潔的介面給使用者，從而實現了使用者側與系統側的解耦，使使用者編寫的 Op 無需隨著系統的更新迭代而做對應變化。而同時為了允許 MegEngine 系統去管理使用這些 Op，Custom Op 中又設計了相關管理與適配模組，其會自動化的將使用者的 Op 封裝適配成 MegEngine 動態圖與靜態圖中的 Op，使使用者 Op 在系統中行為與原生 Op 保持一致，從而便於系統的使用管理。通過這樣的設計，使用者在整合 Op 時無需瞭解 MegEngine 框架本身，只需二十行程式碼即可快速的將 kernel 整合入 MegEngine 並使用，可以大大降低演算法使用者整合 kernel 時的難度與工作量。

原文地址：https://megengine.org.cn/blog/design-custom-operator-system

相關文章