核心註冊¶

概覽¶

在 ExecuTorch 模型匯出的最後階段，我們將方言中的運算元降級為核心 ATen 運算元的 out 變體。然後我們將這些運算元名稱序列化到模型 artifact 中。在執行時執行期間，對於每個運算元名稱，我們需要找到實際的核心，即執行繁重計算並返回結果的 C++ 函式。

核心庫¶

第一方核心庫：¶

可移植核心庫 是內部預設的核心庫，涵蓋了大多數核心 ATen 運算元。它易於使用/閱讀，並用可移植的 C++17 編寫。然而，它並未針對性能進行最佳化，因為它沒有為任何特定目標進行專門化。因此，我們提供了核心註冊 API，以便 ExecuTorch 使用者輕鬆註冊他們自己的最佳化核心。

最佳化核心庫 針對某些運算元專門最佳化效能，利用了 EigenBLAS 等現有第三方庫。這與可移植核心庫配合使用效果最佳，在可移植性和效能之間取得了良好的平衡。在這裡可以找到一個結合這兩個庫的示例。

量化核心庫 實現了量化和反量化的運算元。這些不是核心 ATen 運算元，但對於大多數生產用例至關重要。

自定義核心庫：¶

實現核心 ATen 運算元的自定義核心。儘管我們沒有關於實現核心 ATen 運算元的自定義核心的內部示例，但最佳化核心庫可以被視為一個很好的例子。我們有最佳化的 add.out 和可移植的 add.out。當用戶組合這兩個庫時，我們提供了 API 來選擇對 add.out 使用哪個核心。為了編寫和使用實現核心 ATen 運算元的自定義核心，建議使用基於 YAML 的方法，因為它在以下方面提供了全面的支援

組合核心庫和定義回退核心；
使用選擇性構建來最小化核心大小。

一個 自定義運算元 是 ExecuTorch 使用者在 PyTorch 的 native_functions.yaml 之外定義的任何運算元。

運算元與核心契約¶

上述所有核心，無論內部實現還是自定義，都應符合以下要求

匹配源自運算元 schema 的呼叫約定。核心註冊 API 將為自定義核心生成標頭檔案作為參考。
滿足邊緣方言中定義的 dtype 約束。對於以特定 dtype 作為引數的張量，自定義核心的結果需要與預期的 dtype 匹配。這些約束在邊緣方言運算元中可用。
給出正確結果。我們將提供一個測試框架來自動測試自定義核心。

API¶

這些是將核心/自定義核心/自定義運算元註冊到 ExecuTorch 中的可用 API

YAML 入口 API
C++ API
- 用於自定義運算元
- CMake 示例

如果不清楚使用哪個 API，請參閱最佳實踐。

YAML 入口 API 高階架構¶

要求 ExecuTorch 使用者提供

帶有 C++ 實現的自定義核心庫
與庫關聯的 YAML 檔案，描述該庫實現了哪些運算元。對於部分核心，yaml 檔案還包含核心支援的 dtypes 和 dim order 資訊。更多詳情請參閱 API 部分。

YAML 入口 API 工作流程¶

在構建時，與核心庫相關的 yaml 檔案將連同模型運算元資訊（參見選擇性構建文件）一起傳遞給 核心解析器，結果是將運算元名稱和張量元資料的組合對映到核心符號。然後程式碼生成工具將使用此對映生成連線核心到 ExecuTorch 執行時的 C++ 繫結。ExecuTorch 使用者需要將此生成的庫連結到其應用程式中才能使用這些核心。

在靜態物件初始化時，核心將被註冊到 ExecuTorch 核心登錄檔中。

在執行時初始化階段，ExecuTorch 將使用運算元名稱和引數元資料作為鍵來查詢核心。例如，對於“aten::add.out”和輸入為 dim order (0, 1, 2, 3) 的 float 張量，ExecuTorch 將進入核心登錄檔查詢與名稱和輸入元資料匹配的核心。

核心 ATen 運算元 out 變體的 YAML 入口 API¶

頂層屬性

op（如果運算元出現在 native_functions.yaml 中）或用於自定義運算元的 func。對於 op 鍵，此鍵的值需要是完整的運算元名稱（包括過載名稱）；如果描述自定義運算元，則為完整的運算元 schema（名稱空間、運算元名稱、運算元過載名稱和 schema 字串）。關於 schema 語法，請參考此說明。
kernels：定義核心資訊。它由 arg_meta 和 kernel_name 組成，它們繫結在一起描述“對於具有這些元資料的輸入張量，使用此核心”。
type_alias（可選）：我們為可能的 dtype 選項賦予別名。T0: [Double, Float] 意味著 T0 可以是 Double 或 Float 中的一個。
dim_order_alias（可選）：類似於 type_alias，我們為可能的 dim order 選項賦予名稱。

kernels 下的屬性

arg_meta：“張量引數名稱”條目列表。這些鍵的值是 dtypes 和 dim orders 別名，由相應的 kernel_name 實現。如果此值為 null，則表示核心將用於所有型別的輸入。
kernel_name：實現此運算元的 C++ 函式的預期名稱。您可以在此處放置任何名稱，但應遵循以下約定：將過載名稱中的 . 替換為下劃線，並將所有字元轉換為小寫。在此示例中，add.out 使用名為 add_out 的 C++ 函式。add.Scalar_out 將變為 add_scalar_out，其中 S 為小寫。我們支援核心的名稱空間，但請注意，我們將在最後一級名稱空間中插入 native::。因此，kernel_name 中的 custom::add_out 將指向 custom::native::add_out。

運算元條目的一些示例

- op: add.out
  kernels:
    - arg_meta: null
      kernel_name: torch::executor::add_out

具有預設核心的核心 ATen 運算元的 out 變體

具有 dtype/dim order 專用核心的 ATen 運算元（適用於 Double dtype，dim order 需為 (0, 1, 2, 3)）

- op: add.out
  type_alias:
    T0: [Double]
  dim_order_alias:
    D0: [[0, 1, 2, 3]]
  kernels:
    - arg_meta:
        self: [T0, D0]
        other: [T0 , D0]
        out: [T0, D0]
      kernel_name: torch::executor::add_out

自定義運算元的 YAML 入口 API¶

如上所述，此選項在選擇性構建和合並運算元庫等功能方面提供了更多支援。

首先我們需要指定運算元 schema 以及一個 kernel 部分。因此，我們使用 func 替代 op 並帶有運算元 schema。例如，這是一個自定義運算元的 yaml 條目

- func: allclose.out(Tensor self, Tensor other, float rtol=1e-05, float atol=1e-08, bool equal_nan=False, bool dummy_param=False, *, Tensor(a!) out) -> Tensor(a!)
  kernels:
    - arg_meta: null
      kernel_name: torch::executor::allclose_out

的 kernel 部分與核心 ATen 運算元中定義的相同。對於運算元 schema，我們重用此 README.md 中定義的 DSL，但有一些差異

僅 Out 變體¶

ExecuTorch 僅支援 out 風格的運算元，其中

呼叫者在最後位置提供名為 out 的輸出 Tensor 或 Tensor 列表。
C++ 函式修改並返回相同的 out 引數。
- 如果 YAML 檔案中的返回型別是 ()（對應於 void），C++ 函式仍應修改 out，但無需返回任何內容。
的 out 引數必須是僅限關鍵字的，這意味著它需要跟在名為 * 的引數後面，就像下面的 add.out 示例中那樣。
按照慣例，這些 out 運算元使用模式 <name>.out 或 <name>.<overload>_out 命名。

由於所有輸出值都透過一個 out 引數返回，ExecuTorch 會忽略實際的 C++ 函式返回值。但是，為了保持一致性，當返回型別非 void 時，函式應始終返回 out。

只能返回 `Tensor` 或 `()`¶

ExecuTorch 僅支援返回單個 Tensor 或單元型別 ()（對應於 void）的運算元。它不支援返回任何其他型別，包括列表、optional、元組或布林值等標量。

支援的引數型別¶

ExecuTorch 不支援核心 PyTorch 支援的所有引數型別。以下是我們當前支援的引數型別列表

Tensor
int
bool
float
str
Scalar
ScalarType
MemoryFormat
Device
Optional
List
List<Optional>
Optional<List>

CMake 宏¶

我們提供了構建時宏來幫助使用者構建他們的核心註冊庫。該宏接受描述核心庫的 yaml 檔案以及模型運算元元資料，並將生成的 C++ 繫結打包到一個 C++ 庫中。該宏在 CMake 中可用。

generate_bindings_for_kernels(FUNCTIONS_YAML functions_yaml CUSTOM_OPS_YAML custom_ops_yaml) 接受一個用於核心 ATen 運算元 out 變體的 yaml 檔案和一個用於自定義運算元的 yaml 檔案，生成用於核心註冊的 C++ 繫結。它還依賴於 gen_selected_ops() 生成的選擇性構建 artifact，更多資訊請參閱選擇性構建文件。然後 gen_operators_lib 將把這些繫結打包成一個 C++ 庫。例如

# SELECT_OPS_LIST: aten::add.out,aten::mm.out
gen_selected_ops("" "${SELECT_OPS_LIST}" "")

# Look for functions.yaml associated with portable libs and generate C++ bindings
generate_bindings_for_kernels(FUNCTIONS_YAML ${EXECUTORCH_ROOT}/kernels/portable/functions.yaml)

# Prepare a C++ library called "generated_lib" with _kernel_lib being the portable library, executorch is a dependency of it.
gen_operators_lib("generated_lib" KERNEL_LIBS ${_kernel_lib} DEPS executorch)

# Link "generated_lib" into the application:
target_link_libraries(executorch_binary generated_lib)

我們還提供了合併兩個 yaml 檔案（給定優先順序）的能力。merge_yaml(FUNCTIONS_YAML functions_yaml FALLBACK_YAML fallback_yaml OUTPUT_DIR out_dir) 將 functions_yaml 和 fallback_yaml 合併為一個 yaml 檔案，如果 functions_yaml 和 fallback_yaml 中存在重複條目，此宏將始終採用 functions_yaml 中的條目。

示例

# functions.yaml
- op: add.out
  kernels:
    - arg_meta: null
      kernel_name: torch::executor::opt_add_out

以及 out 回退

# fallback.yaml
- op: add.out
  kernels:
    - arg_meta: null
      kernel_name: torch::executor::add_out

合併後的 yaml 將包含 functions.yaml 中的條目。

自定義運算元的 C++ API¶

與 YAML 入口 API 不同，C++ API 僅使用 C++ 宏 EXECUTORCH_LIBRARY 和 WRAP_TO_ATEN 進行核心註冊，也不支援選擇性構建。這使得此 API 在開發速度方面更快，因為使用者不必進行 YAML 編寫和構建系統調整。

關於使用哪個 API，請參閱自定義運算元最佳實踐。

類似於 PyTorch 中的 TORCH_LIBRARY，EXECUTORCH_LIBRARY 接受運算元名稱和 C++ 函式名稱，並將它們註冊到 ExecuTorch 執行時。

準備自定義核心實現¶

定義自定義運算元的 schema，包括函式式變體（用於 AOT 編譯）和 out 變體（用於 ExecuTorch 執行時）。該 schema 需要遵循 PyTorch ATen 約定（參見 native_functions.yaml）。例如

custom_linear(Tensor weight, Tensor input, Tensor(?) bias) -> Tensor
custom_linear.out(Tensor weight, Tensor input, Tensor(?) bias, *, Tensor(a!) out) -> Tensor(a!)

然後使用 ExecuTorch 型別並根據 schema 編寫自定義核心，同時使用 API 註冊到 ExecuTorch 執行時

// custom_linear.h/custom_linear.cpp
#include <executorch/runtime/kernel/kernel_includes.h>
Tensor& custom_linear_out(const Tensor& weight, const Tensor& input, optional<Tensor> bias, Tensor& out) {
   // calculation
   return out;
}

使用 C++ 宏將其註冊到 ExecuTorch 中¶

在上面的示例中追加以下行

// custom_linear.h/custom_linear.cpp
// opset namespace myop
EXECUTORCH_LIBRARY(myop, "custom_linear.out", custom_linear_out);

現在我們需要為此運算元編寫一些包裝器，以便其在 PyTorch 中出現，但別擔心，我們不需要重寫核心。為此建立一個單獨的 .cpp 檔案

// custom_linear_pytorch.cpp
#include "custom_linear.h"
#include <torch/library.h>

at::Tensor custom_linear(const at::Tensor& weight, const at::Tensor& input, std::optional<at::Tensor> bias) {
    // initialize out
    at::Tensor out = at::empty({weight.size(1), input.size(1)});
    // wrap kernel in custom_linear.cpp into ATen kernel
    WRAP_TO_ATEN(custom_linear_out, 3)(weight, input, bias, out);
    return out;
}
// standard API to register ops into PyTorch
TORCH_LIBRARY(myop, m) {
    m.def("custom_linear(Tensor weight, Tensor input, Tensor(?) bias) -> Tensor", custom_linear);
    m.def("custom_linear.out(Tensor weight, Tensor input, Tensor(?) bias, *, Tensor(a!) out) -> Tensor(a!)", WRAP_TO_ATEN(custom_linear_out, 3));
}

編譯並連結自定義核心¶

將其連結到 ExecuTorch 執行時：在用於構建二進位制/應用程式的 CMakeLists.txt 中，我們需要將 custom_linear.h/cpp 新增到二進位制目標中。我們也可以構建並連結一個動態載入庫（.so 或 .dylib）。

這是一個示例

# For target_link_options_shared_lib
include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)

# Add a custom op library
add_library(custom_op_lib SHARED ${CMAKE_CURRENT_SOURCE_DIR}/custom_op.cpp)

# Include the header
target_include_directory(custom_op_lib PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)

# Link ExecuTorch library
target_link_libraries(custom_op_lib PUBLIC executorch)

# Define a binary target
add_executable(custom_op_runner PUBLIC main.cpp)

# Link this library with --whole-archive !! IMPORTANT !! this is to avoid the operators being stripped by linker
target_link_options_shared_lib(custom_op_lib)

# Link custom op lib
target_link_libraries(custom_op_runner PUBLIC custom_op_lib)

將其連結到 PyTorch 執行時：我們需要將 custom_linear.h、custom_linear.cpp 和 custom_linear_pytorch.cpp 打包成一個動態載入庫（.so 或 .dylib），並將其載入到我們的 python 環境中。一種方法是

import torch
torch.ops.load_library("libcustom_linear.so/dylib")

# Now we have access to the custom op, backed by kernel implemented in custom_linear.cpp.
op = torch.ops.myop.custom_linear.default

在模型中使用自定義運算元¶

自定義運算元可以在 PyTorch 模型中顯式使用，或者您可以編寫一個轉換來用自定義變體替換核心運算元的例項。對於此示例，您可以找到所有 torch.nn.Linear 的例項並將其替換為 CustomLinear。

def  replace_linear_with_custom_linear(module):
    for name, child in module.named_children():
        if isinstance(child, nn.Linear):
            setattr(
                module,
                name,
                CustomLinear(child.in_features,  child.out_features, child.bias),
        )
        else:
            replace_linear_with_custom_linear(child)

剩餘步驟與正常流程相同。現在您可以在 eager 模式下執行此模組，也可以匯出到 ExecuTorch。

自定義運算元 API 最佳實踐¶

考慮到我們有兩個自定義運算元核心註冊 API，應該使用哪個 API？以下是每個 API 的一些優點和缺點

C++ API
- 優點
  - 只需修改 C++ 程式碼
  - 類似於 PyTorch 自定義運算元 C++ API
  - 維護成本低
- 缺點
  - 不支援選擇性構建
  - 沒有集中的記錄
YAML 入口 API
- 優點
  - 支援選擇性構建
  - 為自定義運算元提供集中位置
    - 對於一個應用程式，它顯示了哪些運算元正在註冊以及哪些核心繫結到這些運算元
- 缺點
  - 使用者需要建立和維護 yaml 檔案
  - 更改運算元定義相對不靈活

總的來說，如果我們正在構建一個使用自定義運算元的應用程式，在開發階段建議使用 C++ API，因為它使用成本低且靈活易變。一旦應用程式進入生產階段，自定義運算元定義和構建系統相當穩定，且需要考慮二進位制大小時，建議使用 Yaml 入口 API。

核心註冊¶

概覽¶

核心庫¶

第一方核心庫：¶

自定義核心庫：¶

運算元與核心契約¶

API¶

YAML 入口 API 高階架構¶

YAML 入口 API 工作流程¶

核心 ATen 運算元 out 變體的 YAML 入口 API¶

自定義運算元的 YAML 入口 API¶

僅 Out 變體¶

只能返回 `Tensor` 或 `()`¶

支援的引數型別¶

CMake 宏¶

自定義運算元的 C++ API¶

準備自定義核心實現¶

使用 C++ 宏將其註冊到 ExecuTorch 中¶

編譯並連結自定義核心¶

在模型中使用自定義運算元¶

自定義運算元 API 最佳實踐¶

文件

教程

資源

核心註冊¶

概覽¶

核心庫¶

第一方核心庫：¶

自定義核心庫：¶

運算元與核心契約¶

API¶

YAML 入口 API 高階架構¶

YAML 入口 API 工作流程¶

核心 ATen 運算元 out 變體的 YAML 入口 API¶

自定義運算元的 YAML 入口 API¶

僅 Out 變體¶

只能返回 Tensor 或 ()¶

支援的引數型別¶

CMake 宏¶

自定義運算元的 C++ API¶

準備自定義核心實現¶

使用 C++ 宏將其註冊到 ExecuTorch 中¶

編譯並連結自定義核心¶

在模型中使用自定義運算元¶

自定義運算元 API 最佳實踐¶

文件

教程

資源

只能返回 `Tensor` 或 `()`¶