内核注册¶

概述¶

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

内核库¶

第一方内核库：¶

Portable kernel library 是内部默认内核库，涵盖了大部分核心 ATen 算子。它易于使用/阅读，并用可移植的 C++17 编写。但是，它没有针对性能进行优化，因为它不是为任何特定目标专门设计的。因此，我们提供了 ExecuTorch 用户可以轻松注册自己优化内核的内核注册 API。

Optimized kernel library 专门针对某些算子的性能进行了优化，利用了现有的第三方库，例如 EigenBLAS。这与可移植内核库结合使用效果最好，在可移植性和性能之间取得了良好的平衡。结合这两个库的一个示例如此处所示。

Quantized kernel library 实现量化和反量化算子的功能。这些是核心 ATen 算子之外的功能，但对大多数生产用例至关重要。

自定义内核库：¶

实现核心 ATen Ops 的自定义内核。虽然我们没有核心 ATen Ops 自定义内核的内部示例，但优化内核库可以被视为一个很好的例子。我们已经优化了 add.out 和一个可移植的 add.out。当用户结合使用这两个库时，我们提供了 API 来选择 add.out 使用哪个内核。为了编写和使用实现核心 ATen Ops 的自定义内核，建议使用基于 YAML 的方法，因为它提供了对以下内容的全面支持：

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

Custom operator 是 ExecuTorch 用户在 PyTorch 的 native_functions.yaml 之外定义的任何算子。

算子和内核契约¶

以上提到的所有内核，无论是内部的还是自定义的，都应符合以下要求：

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

API¶

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

YAML 入口 API
C++ API
- 用于自定义 Op
- CMake 示例

如果不清楚使用哪个 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 Op Out 变体的 YAML 入口 API¶

顶级属性

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

kernels 下的属性

arg_meta：一个“张量参数名称”条目列表。这些键的值是 dtype 和维度顺序别名，由相应的 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/维度顺序专用内核的 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

自定义 Op 的 YAML 入口 API¶

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

首先，我们需要指定算子模式以及一个 kernel 部分。因此，我们使用 func 和算子模式，而不是 op。例如，这是一个自定义 Op 的 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

与核心 ATen Ops 中定义的 kernel 部分相同。对于算子模式，我们重用了此 README.md 中定义的 DSL，但有一些区别：

仅 Out 变体¶

ExecuTorch 仅支持 out 风格的算子，其中：

调用者在最后位置提供名为 out 的输出张量或张量列表。
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 支持的所有参数类型。以下是我们目前支持的参数类型列表：

张量
int
布尔值
浮点数
str
Scalar
ScalarType
MemoryFormat
设备
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 Op Out 变体的 YAML 文件以及一个用于自定义 Op 的 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 中的条目。

示例

# 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 中的条目。

自定义 Op 的 C++ API¶

与 YAML 入口 API 不同，C++ API 仅使用 C++ 宏 EXECUTORCH_LIBRARY 和 WRAP_TO_ATEN 进行内核注册，并且不支持选择性构建。这使得该 API 在开发速度方面更快，因为用户无需进行 YAML 编写和构建系统调整。

请参阅自定义 Op 最佳实践，了解应使用哪个 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);

现在我们需要为该 Op 编写一些包装器，使其在 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。

自定义 Op API 最佳实践¶

鉴于我们有两种用于自定义 Op 的内核注册 API，应该使用哪种 API？以下是每种 API 的优缺点：

C++ API
- 优点
  - 只需要 C++ 代码更改
  - 类似于 PyTorch 自定义 Op 的 C++ API
  - 低维护成本
- 缺点
  - 不支持选择性构建
  - 没有集中式簿记
YAML 入口 API
- 优点
  - 支持选择性构建
  - 为自定义 Op 提供了一个集中式位置
    - 它显示了一个应用程序中正在注册哪些 Op 以及哪些内核与这些 Op 绑定
- 缺点
  - 用户需要创建和维护 YAML 文件
  - 更改 Op 定义相对不灵活

总的来说，如果我们正在构建一个应用程序并且它使用自定义 Op，在开发阶段建议使用 C++ API，因为它使用成本低且更改灵活。一旦应用程序进入生产阶段，其中自定义 Op 定义和构建系统非常稳定且需要考虑二进制文件大小，则建议使用 Yaml 入口 API。

内核注册¶

概述¶

内核库¶

第一方内核库：¶

自定义内核库：¶

算子和内核契约¶

API¶

YAML 入口 API 高层架构¶

YAML 入口 API 工作流程¶

核心 ATen Op Out 变体的 YAML 入口 API¶

自定义 Op 的 YAML 入口 API¶

仅 Out 变体¶

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

支持的参数类型¶

CMake 宏¶

自定义 Op 的 C++ API¶

准备自定义内核实现¶

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

编译并链接自定义内核¶

在模型中使用自定义算子¶

自定义 Op API 最佳实践¶

文档

教程

资源

内核注册¶

概述¶

内核库¶

第一方内核库：¶

自定义内核库：¶

算子和内核契约¶

API¶

YAML 入口 API 高层架构¶

YAML 入口 API 工作流程¶

核心 ATen Op Out 变体的 YAML 入口 API¶

自定义 Op 的 YAML 入口 API¶

仅 Out 变体¶

只能返回 Tensor 或 ()¶

支持的参数类型¶

CMake 宏¶

自定义 Op 的 C++ API¶

准备自定义内核实现¶

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

编译并链接自定义内核¶

在模型中使用自定义算子¶

自定义 Op API 最佳实践¶

文档

教程

资源

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