快捷方式

后端与委托

目标受众:对将自己的编译器和硬件集成到 ExecuTorch 中感兴趣的供应商、后端委托开发者

后端委托是后端处理和执行 PyTorch 程序的入口点,以便利用专用后端和硬件的性能及效率优势,同时仍为 PyTorch 用户提供接近 PyTorch 运行时的体验。

后端接口:概述

从高层来看,后端的入口点由两个组件定义

  • 用于表示程序的 IR:Edge 方言(通过 to_edge API 生成)

  • 后端需要实现的几个接口

    • 预先 (Ahead-of-Time, AOT)

      • 程序预处理(例如,预先编译、转换、优化等)。

    • 运行时

      • 程序初始化(例如,运行时编译)。

      • 程序执行。

      • (可选)程序销毁(例如,释放后端拥有的资源)。

委托后端实现由以下部分组成:

  1. 预先预处理接口

  2. 运行时初始化和执行接口

图示如下

drawing

图 1. 后端接口的高层入口点,包括预先和运行时。

后端接口:预先预处理

后端主要需要实现两个预先入口点:partitionpreprocess

partitioner 是由后端实现的算法,用于标记要下层处理到后端的节点。to_backend API 将应用分区算法,并将每个由相互连接的标记节点组成的子图下层处理到目标后端。每个子图将被发送到后端提供的 preprocess 部分,编译成二进制块(binary blob)。

在分区过程中,不允许 exported_program 修改程序,它应该为每个节点应用标签。PartitionResult 包含带有标记的导出程序和分区标签字典,供 to_backend 查找标签并链接到 backend_idcompile_spec

def partition(
    exported_program: ExportedProgram,
) -> PartitionResult:

在预处理过程中,后端会接收一个 Edge 方言程序和一份指定编译所需值的编译规范列表,并预期返回一个已编译的二进制块(compiled blob),即包含期望在后端运行的程序的二进制文件。在序列化期间,该已编译的二进制块将作为 .pte 文件的一部分被序列化,并直接加载到设备上。此过程的 API 是

def preprocess(
    edge_program: ExportedProgram,
    compile_specs: List[CompileSpec],
) -> PreprocessResult:

此处 实现了一个预处理函数演示。该演示循环遍历 edge_program 图模块中的节点,并将 addmulsin 指令序列化为字符串,该字符串稍后在运行时解析并执行。

图示如下

drawing

图 2. 图经过分区,每个子图将发送到预处理部分。

后端接口:运行时初始化和执行

在运行时,来自 preprocess 函数的已编译二进制块将被加载并直接传递给后端的自定义 init 函数。该函数负责进一步处理已编译单元,以及执行任何后端初始化。然后将调用后端的自定义 execute 函数来执行由 init 生成的句柄。最后,如果某些后端需要销毁,则后端可以实现一个 destroy 函数,该函数将在程序生命周期结束时被调用。

// Runtime check
ET_NODISCARD bool is_available();

// Runtime initialization
ET_NODISCARD virtual Result<DelegateHandle*> init(
    BackendInitContext& context,
    FreeableBuffer* processed,
    ArrayRef<CompileSpec> compile_specs);

// Runtime execution
ET_NODISCARD virtual Error execute(
    BackendExecutionContext& context,
    DelegateHandle* handle,
    EValue** args);

// [optional] Runtime destroy. Destroy the resource held by the backend
virtual void destroy(ET_UNUSED DelegateHandle* handle);

图示如下

drawing

图 3. 标准 ExecuTorch 运行时与后端入口点的关系。

为了使后端对 ExecuTorch 运行时可用,必须通过 register_backend API 进行注册

ET_NODISCARD Error register_backend(const Backend& backend);

后端的静态注册,即在库初始化或加载时,可以通过以下方式实现

namespace {
auto cls = BackendWithCompiler();
Backend backend{"BackendWithCompilerDemo", &cls};
static auto success_with_compiler = register_backend(backend);
} // namespace

开发者工具集成:可调试性

提供一致的调试体验,无论是运行时故障还是性能分析,都非常重要。ExecuTorch 为此目的采用了原生的开发者工具,该工具通过调试句柄(debug handles)实现程序指令与原始 PyTorch 代码的关联。您可以在此处阅读更多信息。

委托程序或子图对于 ExecuTorch 运行时来说是不可见的,它们表现为一个特殊的 call_delegate 指令,该指令要求相应的后端处理子图或程序的执行。由于后端委托的不透明性,原生的开发者工具无法查看委托的程序。因此,与非委托执行相比,委托执行的调试(功能性或性能)体验会显著受损。

为了为用户提供一致的调试体验,无论模型是否使用委托,开发者工具都提供了一个接口来关联委托(子)图与原始(子)图。开发者工具通过调试句柄映射(debug handles map)来实现这一点,该映射允许委托生成可以与委托消费的原始(子)图相关联的内部句柄。然后在运行时,后端开发者可以使用内部句柄报告错误或性能分析信息,这些信息将通过调试句柄映射映射回原始(子)图。更多信息请参阅委托调试

通过利用调试标识符(debug identifier),后端开发者可以将调试信息嵌入到委托的二进制块(delegated blob)中

drawing

这样,在执行阶段,借助调试标识符,后端开发者可以将委托内部的失败指令关联回原始 Python 代码的确切行。

drawing

常见问题

1. 如何在 backend.preprocess 中获取数据?

正在预处理的图模块是一个提升后的图(lifted graph),这意味着静态数据(如权重和偏差)作为输入提供给图。但是,我们可以通过导出程序(exported program)预先访问这些权重和偏差。要从给定节点访问这些参数,可以使用 torch/_export/utils.py 中提供的函数 get_params

2. 如何将数据(如权重/偏差)嵌入到后端?

后端通常有一些方法来优化常量数据。在这种情况下,我们需要在分区器中标记作为状态的占位符节点,并且在 backend.preprocess 期间,我们可以按照第一个问题中的描述获取权重。

3. 如何使用特定后端在 Python 中运行下层处理后的模块?

我们尚未添加此功能,但这在计划中!

4. 我们是否应该期望在 edge 方言程序中看到 get_attr 节点?

get_attr 节点仅用于控制流或委托的子模块。它不包含任何数据。

5. 我们可以委托给多个后端吗?

可以!有两种方法可以实现这一点

选项 1: 对不同的后端多次运行 to_backend

如果我们有两个后端:backend_1 和 backend_2,并且它们有各自的分区器:backend_1_parititioner 和 backend_2_partitioner,我们可以这样运行:

# Will first lower nodes to backend_1 depending on the backend_1_parititioner depending on partitioner algorithm
exported_program_backend_1 = to_backend(exported_program, backend_1_parititioner())
# For the rest of nodes, they will be lowered to backend_2 depending on backend_2_parititioner
exported_program_backend_1_and_2 = to_backend(exported_program_backend_1, backend_2_parititioner())

更具体的示例可以在此处找到。在此示例中,qnnpack 是一个后端,xnnpack 是另一个后端。我们尚未开源这两个后端委托,因此此示例无法直接运行。它可以作为参考,了解如何实现。

此选项易于尝试,因为通常所有后端都会实现自己的分区器。然而,如果我们改变 to_backend 调用的顺序,此选项可能会得到不同的结果。如果我们想更好地控制节点(例如它们应该去哪个后端),选项 2 更好。

选项 2: 拥有一个为不同后端分区的分区器

另一个选项是创建一个自定义分区器,例如分区器 backend_1_2_partitioner,并在分区器逻辑内部,

class Backend_1_2_Partitioner(Partitioner):
    """
    Partitions all add/mul nodes regardless of order for Backend2
    """

    def __init__(self) -> None:
        self.delegation_spec_1 = DelegationSpec("Backend1", [])
        self.delegation_spec_2 = DelegationSpec("Backend2", [])
        self.partition_tags = {}

    def partition(
        self, exported_program: ExportedProgram
    ) -> ExportedProgram:

        # Tag all nodes in the first partiton to backend 1
        node_to_backend_1 = ... # some logic to select the nodes from the graph
        delegation_tag = f"backend2_tag{partitioner_1.id}"
        node.meta["delegation_tag"] = delegation_tag
        self.partition_tags[delegation_tag] = self.delegation_spec_1

        # Tag all nodes in the first partiton to backend 2
        node_to_backend_2 = ... # some logic to select the nodes from the graph
        delegation_tag = f"backend2_tag{partitioner_2.id}"
        node.meta["delegation_tag"] = delegation_tag
        self.partition_tags[delegation_tag] = self.delegation_spec_2
        return exported_program

6. 有没有简单的方法来编写分区器?

我们此处提供了一些辅助分区器,以便轻松地从分解的算子中查找节点。

7. 如何将节点链接回源代码? 我们提供了一个辅助函数

from executorch.exir.print_program import inspect_node

print(inspect_node(graph, node))

它将突出显示图中的节点并指向源代码,示例输出如下所示

_param_constant1 error_msg:  Here is the node in the graph module:
graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %_param_constant0 : [num_users=1] = get_attr[target=_param_constant0]
--> %_param_constant1 : [num_users=1] = get_attr[target=_param_constant1]
    %aten_convolution_default : [num_users=2] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%arg0_1, %_param_constant0, %_param_constant1, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
    %_param_constant2 : [num_users=1] = get_attr[target=_param_constant2]
    %_param_constant3 : [num_users=1] = get_attr[target=_param_constant3]
    %aten_convolution_default_1 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%aten_convolution_default, %_param_constant2, %_param_constant3, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
    %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%aten_convolution_default, %aten_convolution_default_1), kwargs = {})
    %_param_constant4 : [num_users=1] = get_attr[target=_param_constant4]
    %_param_constant5 : [num_users=1] = get_attr[target=_param_constant5]
    %aten_convolution_default_2 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%aten_add_tensor, %_param_constant4, %_param_constant5, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
    %aten_gelu_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.gelu.default](args = (%aten_convolution_default_2,), kwargs = {})
    return [aten_gelu_default]
This node _param_constant1 has metadata of:
The node stacktrace:
Traceback (most recent call last):
    File "/tmp/ipykernel_1204253/3382880687.py", line 7, in forward
return self.test_model(x)
    File "/mnt/xarfuse/uid-25337/7b86ad0c-seed-nspid4026532987_cgpid2707357-ns-4026532984/torch/nn/modules/module.py", line 1528, in _call_impl
return forward_call(*args, **kwargs)
    File "/tmp/ipykernel_1204253/712280972.py", line 10, in forward
a = self.conv1(x)

文档

查阅 PyTorch 全面的开发者文档

查看文档

教程

获取适用于初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获取问题解答

查看资源