内核注册¶
概述¶
在ExecuTorch 模型导出的最后阶段，我们将方言中的操作符降低为核心 ATen 操作符的out 变体。然后，我们将这些操作符名称序列化到模型工件中。在运行时执行期间，对于每个操作符名称，我们需要找到实际的内核，即执行繁重计算并返回结果的 C++ 函数。
便携式内核库是内部默认内核库，它易于使用，并且适用于大多数目标后端。然而，由于它没有针对任何特定目标进行专门化，因此它没有针对性能进行优化。因此，我们为 ExecuTorch 用户提供了内核注册 API，以便他们轻松注册自己的优化内核。
设计原则¶
我们支持什么？在操作符覆盖方面，内核注册 API 允许用户为所有核心 ATen 操作符以及自定义操作符注册内核，只要指定了自定义操作符模式即可。
请注意，我们还支持部分内核，例如内核仅支持张量数据类型和/或维度顺序的子集。
内核契约：内核需要符合以下要求
与从操作符模式派生的调用约定相匹配。内核注册 API 将生成自定义内核的标头作为参考。
满足边缘方言中定义的数据类型约束。对于具有某些数据类型作为参数的张量，自定义内核的结果需要与预期的数据类型匹配。约束在边缘方言操作符中提供。
提供正确的结果。我们将提供一个测试框架来自动测试自定义内核。
高级架构¶
要求 ExecuTorch 用户提供
带有 C++ 实现的自定义内核库
与该库关联的 yaml 文件，用于描述此库实现哪些运算符。对于部分内核，yaml 文件还包含内核支持的数据类型和维度顺序的信息。有关更多详细信息，请参阅 API 部分。
工作流¶
在构建时，与内核库关联的 yaml 文件将与模型运算符信息（请参阅选择性构建文档）一起传递给内核解析器，其结果是运算符名称和张量元数据与内核符号之间的映射。然后，代码生成工具将使用此映射生成将内核连接到 ExecuTorch 运行时的 C++ 绑定。ExecuTorch 用户需要将此生成的库链接到其应用程序中才能使用这些内核。
在静态对象初始化时，内核将注册到 ExecuTorch 内核注册表中。
在运行时初始化阶段，ExecuTorch 将使用运算符名称和参数元数据作为查找内核的键。例如，对于“aten::add.out”，输入为维度顺序为 (0, 1, 2, 3) 的浮点张量，ExecuTorch 将进入内核注册表并查找与名称和输入元数据匹配的内核。
API¶
有两组 API：描述内核 - 运算符映射的 yaml 文件和使用这些映射的代码生成工具。
核心 ATen Op Out 变体的 Yaml 条目¶
顶级属性
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::。因此，kernel_name中的custom::add_out将指向custom::native::add_out。
一些运算符条目的示例
- op: add.out
kernels:
- arg_meta: null
kernel_name: torch::executor::add_out
具有默认内核的核心 ATen 运算符的 out 变体
具有数据类型/维度顺序专用内核的 ATen 运算符（适用于 Double 数据类型，维度顺序需要为 (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
自定义操作 C++ API¶
对于实现自定义运算符的自定义内核，我们提供了 2 种方法将其注册到 ExecuTorch 运行时
使用
EXECUTORCH_LIBRARY和WRAP_TO_ATENC++ 宏。使用
functions.yaml和代码生成 C++ 库。
第一个选项需要 C++17，并且尚未支持选择性构建，但它比第二个选项更快，在第二个选项中，我们必须进行 yaml 编写和构建系统调整。
第一个选项特别适合快速原型制作，但也可以用于生产。
与 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++ 宏将其注册到 PyTorch 和 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）并将其链接。
将其链接到 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 条目¶
如上所述，此选项在选择性构建和合并运算符库等方面提供了更多支持。
首先，我们需要指定运算符架构以及 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
kernel 部分与 core ATen ops 中定义的部分相同。对于运算符架构，我们正在重复使用此 README.md 中定义的 DSL，但有一些区别
仅输出变量¶
ExecuTorch 仅支持输出样式运算符，其中
调用方在最终位置提供输出张量或张量列表，名称为
out。C++ 函数修改并返回相同的
out参数。如果 YAML 文件中的返回类型为
()（映射到 void），则 C++ 函数仍应修改out，但不需要返回任何内容。
out参数必须是仅关键字的，这意味着它需要跟随一个名为*的参数，如下面的add.out示例所示。按照惯例，这些输出运算符使用模式
<name>.out或<name>.<overload>_out命名。
由于所有输出值都是通过 out 参数返回的，因此 ExecuTorch 忽略实际的 C++ 函数返回值。但是，为了保持一致性，当返回类型为非 void 时，函数应始终返回 out。
只能返回 Tensor 或 ()¶
ExecuTorch 仅支持返回单个 Tensor 或单元类型 ()（映射到 void）的运算符。它不支持返回任何其他类型，包括列表、可选类型、元组或标量，如 bool。
支持的参数类型¶
ExecuTorch 并不支持核心 PyTorch 支持的所有参数类型。以下是我们当前支持的参数类型列表
Tensor
int
bool
float
str
Scalar
ScalarType
MemoryFormat
Device
Optional
List
List<Optional
> Optional<List
>
构建工具宏¶
我们提供构建时宏以帮助用户构建其内核注册库。该宏采用描述内核库的 yaml 文件以及模型运算符元数据，并将生成的 C++ 绑定打包到 C++ 库中。该宏可在 CMake 和 Buck2 上使用。
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
以及我们的后备
# fallback.yaml
- op: add.out
kernels:
- arg_meta: null
kernel_name: torch::executor::add_out
合并后的 yaml 将包含 functions.yaml 中的条目。
Buck2¶
executorch_generated_lib 是采用 yaml 文件并依赖于选择性构建宏 et_operator_library 的宏。例如
# Yaml file for kernel library
export_file(
name = "functions.yaml"
)
# Kernel library
cxx_library(
name = "add_kernel",
srcs = ["add.cpp"],
)
# Selective build artifact, it allows all operators to be registered
et_operator_library(
name = "all_ops",
include_all_ops = True, # Select all ops in functions.yaml
)
# Prepare a generated_lib
executorch_generated_lib(
name = "generated_lib",
functions_yaml_target = ":functions.yaml",
deps = [
":all_ops",
":add_kernel",
],
)
# Link generated_lib to ExecuTorch binary
cxx_binary(
name = "executorch_bin",
deps = [
":generated_lib",
],
)
