内核注册¶

概述¶

在ExecuTorch 模型导出的最后阶段，我们将方言中的操作符降低为核心 ATen 操作符的out 变体。然后我们将这些操作符名称序列化到模型工件中。在运行时执行期间，对于每个操作符名称，我们需要找到实际的内核，即执行繁重计算并返回结果的 C++ 函数。

内核库¶

官方内核库：¶

便携式内核库 是内部默认的内核库，涵盖了大部分核心 ATen 操作符。它易于使用和阅读，并且是用可移植的 C++17 编写的。然而，由于它没有针对任何特定目标进行优化，因此性能并不高。因此，我们为 ExecuTorch 用户提供了内核注册 API，以便他们可以轻松注册自己的优化内核。

优化的内核库 专注于某些操作符的性能，利用现有的第三方库，如 EigenBLAS。这与便携式内核库配合使用效果最佳，在可移植性和性能之间取得了良好的平衡。结合这两个库的一个示例可以在这里找到这里。

量化内核库 实现了用于量化和反量化的操作符。这些是核心外的 ATen 操作符，但对于大多数生产用例至关重要。

自定义内核库：¶

实现核心ATen操作的自定义内核。尽管我们没有为实现核心ATen操作的自定义内核提供内部示例，但优化的内核库可以被视为一个很好的例子。我们优化了add.out和一个便携式的add.out。当用户结合这两个库时，我们提供了API来选择使用哪个内核用于add.out。为了编写和使用实现核心ATen操作的自定义内核，建议使用基于YAML的方法，因为它提供了全面的支持

组合内核库并定义回退内核；
使用选择性构建以最小化内核大小。

自定义操作符是指任何由ExecuTorch用户在PyTorch的 native_functions.yaml 之外定义的操作符。 自定义操作符 是指任何由ExecuTorch用户在PyTorch的 native_functions.yaml 之外定义的操作符。

运算符与内核契约¶

上述所有内核，无论是内部开发还是自定义的，都应满足以下要求：

匹配从算子模式推导出的调用约定。内核注册 API 将为自定义内核生成头文件作为参考。
满足边缘方言中定义的 dtype 约束。对于以特定 dtype 作为参数的张量，自定义内核的结果需要与预期的 dtype 相匹配。这些约束可在边缘方言算子中找到。
给出正确结果。我们将提供一个测试框架，用于自动测试自定义内核。

APIs¶

这些是可用于将内核/自定义内核/自定义算子注册到 ExecuTorch 的 API：

如果不清楚使用哪个API，请参阅最佳实践。

YAML 条目 API 高层架构¶

ExecuTorch 用户被要求提供：

带有 C++ 实现的自定义内核库
一个与该库关联的 YAML 文件，用于描述该库实现了哪些算子。对于部分内核，YAML 文件还包含有关该内核支持的 dtype 和维度顺序的信息。更多详情请参阅 API 部分。

YAML 条目 API 工作流¶

在构建时，与内核库关联的 yaml 文件将与模型算子信息（参见选择性构建文档）一起传递给内核解析器，其结果是算子名称组合与张量元数据到内核符号的映射。然后，代码生成工具将利用此映射生成 C++ 绑定，将内核连接到 ExecuTorch 运行时。ExecuTorch 用户需要将此生成的库链接到其应用程序中才能使用这些内核。

在静态对象初始化时，内核将被注册到 ExecuTorch 内核注册表中。

在运行时初始化阶段，ExecuTorch 将使用算子名称和参数元数据作为键来查找内核。例如，对于"aten::add.out"且输入为维度顺序为 (0, 1, 2, 3) 的浮点张量，ExecuTorch 将进入内核注册表，查找与该名称及输入元数据匹配的内核。

核心 ATen 操作输出变体的 YAML 条目 API¶

顶级属性：

op（如果操作符出现在native_functions.yaml中）或func用于自定义操作符。此键的值需要是op键的完整操作符名称（包括重载名称），或者如果我们正在描述一个自定义操作符，则需要是一个完整的操作符模式（命名空间、操作符名称、操作符重载名称和模式字符串）。有关模式语法，请参阅此说明。
kernels：定义内核信息。它由 arg_meta 和 kernel_name 组成，两者绑定在一起以描述“对于具有这些元数据的输入张量，使用此内核”。
type_alias（可选）：我们为可能的数据类型选项指定了别名。T0: [Double, Float] 表示 T0 可以是 Double 或 Float。
dim_order_alias(可选)：类似于 type_alias，我们为可能的维度顺序选项指定名称。

kernels 下的属性：

arg_meta：一个“张量参数名称”条目列表。这些键的值是数据类型和维度顺序的别名，由相应的 kernel_name 实现。由于这是 null，意味着该内核将用于所有类型的输入。
kernel_name：将实现此运算符的 C++ 函数的预期名称。您可以在此处放置任何您想要的内容，但应遵循以下约定：将重载名称中的 . 替换为下划线，并将所有字符转换为小写。在此示例中，add.out 使用名为 add_out 的 C++ 函数。add.Scalar_out 将变为 add_scalar_out，其中 S 为小写。我们支持内核的命名空间，但请注意，我们将向最内层命名空间插入一个 native::。因此，custom::add_out 中的 kernel_name 将指向 custom::native::add_out。

运算符输入的一些示例：

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

带有默认内核的核心 ATen 算子的 out 变体

带有 dtype/维度顺序专用内核的 ATen 运算符（适用于 Double 种 dtype，且维度顺序需为 (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¶

如上所述，该选项在选择性构建以及合并算子库等功能方面提供了更多支持。

首先，我们需要指定操作符模式以及一个kernel部分。因此，我们不使用op，而是使用带有操作符模式的func。例如，以下是一个自定义操作的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 算子中定义的部分相同。对于算子模式，我们重用了在此 README.md 中定义的 DSL，但有一些差异：

仅输出变体¶

ExecuTorch 仅支持 out-style 算子，其中：

调用者在最终位置提供名为 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）的操作符。它不支持返回任何其他类型，包括列表、可选类型、元组或标量（如 bool）。

支持的参数类型¶

ExecuTorch 不支持核心 PyTorch 所支持的所有参数类型。以下是我们目前支持的参数类型列表：

张量
整数
布尔
浮点型
字符串
标量
ScalarType
MemoryFormat
设备
可选的
列表
List<Optional>
Optional<List>

CMake 宏¶

我们提供构建时宏，帮助用户构建其内核注册库。该宏接受描述内核库的 YAML 文件以及模型算子元数据，并将生成的 C++ 绑定打包为 C++ 库。该宏可在 CMake 中使用。

generate_bindings_for_kernels(FUNCTIONS_YAML functions_yaml CUSTOM_OPS_YAML custom_ops_yaml) 接受一个用于核心 ATen 算子输出变体的 yaml 文件以及一个用于自定义算子的 yaml 文件，生成用于内核注册的 C++ 绑定。它还依赖于由 gen_selected_ops() 生成的选择性构建产物，有关更多信息请参阅选择性构建文档。然后 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 中的条目。

Example:

# functions.yaml
- op: add.out
  kernels:
    - arg_meta: null
      kernel_name: torch::executor::opt_add_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 运行时中。

准备自定义内核实现¶

为函数式变体（用于 AOT 编译）和 out 变体（用于 ExecuTorch 运行时）定义您的自定义算子模式。该模式需遵循 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 类型编写自定义内核，并配合用于向 ExecuTorch 运行时注册的 API：

// 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}/build/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

自定义算子 API 最佳实践¶

鉴于我们为自定义算子提供了两个内核注册 API，我们应该使用哪一个？以下是每种 API 的优缺点：

C++ API：
- Pros:
  - 仅需修改 C++ 代码
  - 类似于 PyTorch 自定义算子的 C++ API
  - 维护成本低
- Cons:
  - 不支持选择性构建
  - 无集中式账簿管理
YAML 入口 API：
- Pros:
  - 支持选择性构建
  - 提供自定义算子的集中式存放位置
    - 它展示了应用程序中正在注册哪些算子，以及哪些内核已绑定到这些算子。
- Cons:
  - 用户需要创建并维护 YAML 文件
  - 修改算子定义相对不够灵活

总体而言，如果我们正在构建一个使用自定义算子的应用程序，在开发阶段建议使用 C++ API，因为它使用成本低且易于调整。一旦应用程序进入生产阶段，此时自定义算子的定义和构建系统已相当稳定，并且需要考虑二进制文件大小，则建议使用 YAML 入口 API。

内核注册¶

概述¶

内核库¶

官方内核库：¶

自定义内核库：¶

运算符与内核契约¶

APIs¶

YAML 条目 API 高层架构¶

YAML 条目 API 工作流¶

核心 ATen 操作输出变体的 YAML 条目 API¶

自定义算子的 YAML 条目 API¶

仅输出变体¶

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

支持的参数类型¶

CMake 宏¶

自定义操作的 C++ API¶

准备自定义内核实现¶

使用 C++ 宏将其注册到 ExecuTorch¶

编译并链接自定义内核¶

自定义算子 API 最佳实践¶

文档

教程

资源

内核注册¶

概述¶

内核库¶

官方内核库：¶

自定义内核库：¶

运算符与内核契约¶

APIs¶

YAML 条目 API 高层架构¶

YAML 条目 API 工作流¶

核心 ATen 操作输出变体的 YAML 条目 API¶

自定义算子的 YAML 条目 API¶

仅输出变体¶

只能返回 Tensor 或 ()¶

支持的参数类型¶

CMake 宏¶

自定义操作的 C++ API¶

准备自定义内核实现¶

使用 C++ 宏将其注册到 ExecuTorch¶

编译并链接自定义内核¶

自定义算子 API 最佳实践¶

文档

教程

资源

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