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

作者：褚超群 | 曠視科技 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