内核注册¶

概述¶

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

可移植内核库是内部默认的内核库，它易于使用且可移植，适用于大多数目标后端。但是，它没有针对性能进行优化，因为它不是专门针对任何特定目标的。因此，我们为 ExecuTorch 用户提供了内核注册 API，以便他们轻松注册自己的优化内核。

设计原则¶

我们支持什么？在操作员覆盖率方面，内核注册 API 允许用户为所有核心 ATen 操作以及自定义操作注册内核，只要指定了自定义操作架构即可。

请注意，我们还支持 _partial 内核，_for示例内核仅支持张量 dtypes 和/或 dim orders 的子集。

内核合约：内核需要符合以下要求：

匹配从 operator 架构派生的调用约定。内核注册 API 将为自定义内核生成标头作为引用。
满足 edge dialect 中定义的 dtype 约束。对于具有某些 dtype 作为参数的张量，自定义内核的结果需要与预期的 dtype 匹配。这些约束在 edge dialect 操作中可用。
给出正确的结果。我们将提供一个测试框架来自动测试自定义内核。

高级架构¶

ExecuTorch 用户需要提供：

具有 C++ 实现的自定义内核库
与库关联的 YAML 文件，用于描述此库正在实现的运算符。对于部分内核，yaml 文件还包含有关内核支持的 dtypes 和 dim 顺序的信息。API 部分提供了更多详细信息。

工作流¶

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

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

在运行时初始化阶段，ExecuTorch 将使用 operator 名称和参数元数据作为 key 来查找内核。例如，当 “aten：：add.out” 和输入是具有 dim 顺序（0， 1， 2， 3）的浮点张量时，ExecuTorch 将进入内核注册表并查找与名称和输入元数据匹配的内核。

蜜蜂属¶

有两组 API：描述内核 - 操作员映射的 yaml 文件和使用这些映射的 codegen 工具。

核心 ATen op out 变体的 yaml entry¶

顶级属性：

op（如果运算符显示在）或 for custom 运算符。此键的值需要是 key 的完整运算符名称（包括重载名称），或者如果我们描述的是自定义运算符，则需要是完整的运算符架构（命名空间、运算符名称、运算符重载名称和架构字符串）。有关 schema 语法，请参阅此说明。native_functions.yamlfuncop
kernels：定义内核信息。它由和组成，它们绑定在一起以描述 “对于具有这些元数据的输入张量，请使用此内核”。arg_metakernel_name
type_alias（可选）：我们正在为可能的 dtype 选项提供别名。 means 可以是或之一。T0: [Double, Float]T0DoubleFloat
dim_order_alias（可选）：与类似，我们正在为可能的 dim order 选项命名。type_alias

下面的属性：kernels

arg_meta： “Tensor arg name” 条目列表。这些键的值是 dtypes 和 dim orders 别名，它们由相应的 .这意味着内核将用于所有类型的输入。kernel_namenull
kernel_name：将实现此运算符的 C++ 函数的预期名称。您可以在此处放置所需的任何内容，但应遵循以下约定：将重载名称中的 the 替换为下划线，并将所有字符小写。在此示例中，使用名为的 C++ 函数。将变为，带有小写字母。我们支持内核的 namespace，但请注意，我们将在 namespace 的最后一级插入 a。所以在将指向。.add.outadd_outadd.Scalar_outadd_scalar_outSnative::custom::add_outkernel_namecustom::native::add_out

运算符输入的一些示例：

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

具有默认内核的核心 ATen 运算符的 out 变体

具有 dtype/dim order 专用内核的 ATen 运算符（适用于 dtype 和 dim order 需要为（0， 1， 2， 3））Double

- 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

自定义操作 C++ API¶

对于实现自定义算子的自定义内核，我们提供了 2 种方式将其注册到 ExecuTorch 运行时中：

Using 和 C++ 宏，本节将介绍。EXECUTORCH_LIBRARYWRAP_TO_ATEN
Using 和 codegen 的 C++ 库，将在下一节中介绍。functions.yaml

请参阅自定义操作最佳实践了解使用哪个 API。

第一个选项需要 C++17，并且还没有选择性构建支持，但它比第二个选项更快，我们必须完成 yaml 创作和构建系统调整。

第一种选项特别适用于快速原型制作，但也可用于生产。

与类似，获取运算符名称和 C++ 函数名称，并将它们注册到 ExecuTorch 运行时中。TORCH_LIBRARYEXECUTORCH_LIBRARY

准备自定义内核实现¶

为功能变体（用于 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++宏将其注册到PyTorch和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）并链接它。

将其链接到 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

自定义 Ops Yaml 入口¶

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

首先，我们需要指定 operator 架构以及一个部分。因此，我们不是使用 Operator 模式。例如，下面是自定义操作的 yaml 条目：kernelopfunc

- 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 操作中定义的部分相同。对于运算符架构，我们将重用此 README.md 中定义的 DSL，但有一些区别：kernel

仅输出变体¶

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

调用方在最终位置提供名为 .out
C++ 函数修改并返回相同的参数。out
- 如果 YAML 文件中的返回类型为（映射到 void），则 C++ 函数仍应修改，但不需要返回任何内容。()out
该参数必须是仅关键字，这意味着它需要遵循如下例中命名的参数。out*add.out
通常，这些 out 运算符使用模式或 .<name>.out<name>.<overload>_out

由于所有输出值都是通过参数返回的，因此 ExecuTorch 会忽略实际的 C++ 函数返回值。但是，为了保持一致，函数应该总是在返回类型为 non- 时返回。outoutvoid

只能退货或`Tensor()`¶

ExecuTorch 仅支持返回单个、或 unit 类型（映射到）的运算符。它不支持返回任何其他类型，包括列表、可选值、元组或标量，例如 .Tensor()voidbool

支持的参数类型¶

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

张肌
int
布尔
浮
str
标量
标量类型
内存格式
装置
自选
列表
列表<可选>
可选<List>

构建工具宏¶

我们提供构建时宏来帮助用户构建他们的内核注册库。该宏采用描述内核库的 yaml 文件以及模型运算符元数据，并将生成的 C++ 绑定打包到 C++ 库中。该宏在 CMake 上可用。

CMake¶

generate_bindings_for_kernels(FUNCTIONS_YAML functions_yaml CUSTOM_OPS_YAML custom_ops_yaml)为核心 ATen op out 变体获取 yaml 文件，为自定义操作获取 yaml 文件，为内核注册生成 C++ 绑定。它还取决于生成的选择性构建工件，有关更多信息，请参阅选择性构建文档。然后将这些绑定打包为 C++ 库。例如：gen_selected_ops()gen_operators_lib

# 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 文件的功能，但有优先权。将 functions_yaml 和 fallback_yaml 合并为单个 YAML，如果 functions_yaml 和 fallback_yaml 中存在重复的条目，则此宏将始终采用 functions_yaml 中的条目。merge_yaml(FUNCTIONS_YAML functions_yaml FALLBACK_YAML fallback_yaml OUTPUT_DIR out_dir)

例：

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

自定义操作 API 最佳实践¶

鉴于我们有 2 个用于自定义操作的内核注册 API，我们应该使用哪个 API？以下是每个 API 的一些优缺点：

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

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

内核注册¶

概述¶

设计原则¶

高级架构¶

工作流¶

蜜蜂属¶

核心 ATen op out 变体的 yaml entry¶

自定义操作 C++ API¶

准备自定义内核实现¶

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

编译并链接自定义内核¶

自定义 Ops Yaml 入口¶

仅输出变体¶

只能退货或`Tensor()`¶

支持的参数类型¶

构建工具宏¶

CMake¶

自定义操作 API 最佳实践¶

文档

教程

资源

内核注册¶

概述¶

设计原则¶

高级架构¶

工作流¶

蜜蜂属¶

核心 ATen op out 变体的 yaml entry¶

自定义操作 C++ API¶

准备自定义内核实现¶

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

编译并链接自定义内核¶

自定义 Ops Yaml 入口¶

仅输出变体¶

只能退货或Tensor()¶

支持的参数类型¶

构建工具宏¶

CMake¶

自定义操作 API 最佳实践¶

文档

教程

资源

只能退货或`Tensor()`¶