内核注册¶
概述¶
在 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.yaml
func
op
kernels
:定义内核信息。它由 和 组成,它们绑定在一起以描述 “对于具有这些元数据的输入张量,请使用此内核”。arg_meta
kernel_name
type_alias
(可选):我们正在为可能的 dtype 选项提供别名。 means 可以是 或 之一。T0: [Double, Float]
T0
Double
Float
dim_order_alias
(可选):与 类似,我们正在为可能的 dim order 选项命名。type_alias
下面的属性 :kernels
arg_meta
: “Tensor arg name” 条目列表。这些键的值是 dtypes 和 dim orders 别名,它们由相应的 .这意味着内核将用于所有类型的输入。kernel_name
null
kernel_name
:将实现此运算符的 C++ 函数的预期名称。您可以在此处放置所需的任何内容,但应遵循以下约定:将重载名称中的 the 替换为下划线,并将所有字符小写。在此示例中,使用名为 的 C++ 函数。 将变为 ,带有小写字母 。我们支持内核的 namespace,但请注意,我们将在 namespace 的最后一级插入 a。所以在将指向 。.
add.out
add_out
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/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
自定义 Ops Yaml 入口¶
对于自定义运算(不属于核心 ATen 运算集的 out 变体的运算),我们需要指定 operator 模式以及一个部分。因此,我们不是使用 Operator 模式。例如,下面是自定义操作的 yaml 条目:kernel
op
func
- 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- 时返回。out
out
void
只能退货或Tensor
()
¶
ExecuTorch 仅支持返回单个 、 或 unit 类型(映射到 )的运算符。它不支持返回任何其他类型,包括列表、可选值、元组或标量,例如 .Tensor
()
void
bool
支持的参数类型¶
ExecuTorch 不支持核心 PyTorch 支持的所有参数类型。以下是我们目前支持的参数类型列表:
张肌
int
布尔
浮
str
标量
标量类型
内存格式
装置
自选
列表
列表<可选
> 可选<List
>
构建工具宏¶
我们提供构建时宏来帮助用户构建他们的内核注册库。该宏采用描述内核库的 yaml 文件以及模型运算符元数据,并将生成的 C++ 绑定打包到 C++ 库中。该宏在 CMake 和 Buck2 上都可用。
CMake¶
generate_bindings_for_kernels(functions_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(${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_lib} executorch)
# Link "generated_lib" into the application:
target_link_libraries(executorch_binary generated_lib)
巴克2¶
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",
],
)