示例工程【免费下载链接】tutorialsPyTorch tutorials.项目地址https://gitcode.com/gh_mirrors/tuto/tutorials点击查看免费下载本指南以 PyTorch tutorials 仓库中的 advanced_source/dispatcher.rst 为核心结合仓库内的 op.cpp、test.py 与 CMakeLists.txt 示例系统讲解如何用TORCH_LIBRARY/TORCH_LIBRARY_IMPLAPI 把自定义 C 算子注册进 PyTorch Dispatcher按 CPU、CUDA、Autograd、Autocast 等 dispatch key 分层组织内核实现。读完你将掌握如何拆分算子 schema 与后端实现、如何编写带自动求导的算子内核含防递归重入的要点、如何用autogradNotImplementedFallback为无需求导的算子补齐行为以及如何为自定义算子注册 autocast 包装内核。文档背景与适用范围注意重要本教程所描述的技术路线在PyTorch 2.4 起已被标记为弃用。更新、更权威的自定义算子扩展指南请参阅仓库中的 custom_ops_landing_page.rst其下还有 python_custom_ops.rst 与 cpp_custom_ops.rst 两份分册教程。尽管如此理解 Dispatcher 的分层注册机制仍然是深入 PyTorch 底层自动求导、AMP、vmap、导出追踪的必修课本教程展示的TORCH_LIBRARY/TORCH_LIBRARY_IMPL注册模式也仍然被许多第三方扩展沿用。Dispatcher 是什么一次torch::add调用背后的多重决策Dispatcher 是 PyTorch 的内部组件负责在调用像torch::add这样的函数时决定实际应该执行哪一段代码。这件事并不平凡因为 PyTorch 算子需要处理大量层层叠加layered的横切关注点例如根据输入张量的设备类型在 CPU 与 CUDA 实现之间切换根据是否需要自动求导处理在 Autograd 实现与后端实现之间切换在需要时应用 autocasting 以实现自动混合精度AMP在vmap调用下应用 batching rules批处理规则在导出追踪模型时记录算子执行轨迹tracing。核心设计思想是Dispatcher 把算子的实现拆分成多个 kernel每个 kernel 服务于一个特定的 dispatch key如 CPU、CUDA。调用算子时Dispatcher 综合张量参数与线程局部状态找出当前优先级最高的 dispatch key并把控制权转交给该 key 对应的 kernel。最终效果是调用算子时先执行 Autograd kernel再根据传入张量的设备类型 redispatch 到对应的后端 kernel。如果你在自定义算子代码里发现自己手动写了大量 if 语句去逐个处理这些情况那么 Dispatcher API 可以帮助你把代码组织得更清晰。反过来说如果你的自定义算子非常简单、只用于 CPU 推理那么大概率不需要 Dispatcher直接用基础注册 API 即可参见 torch_script_custom_ops.rst。在继续之前建议先熟悉两篇前置教程自定义算子注册 与 C 自定义自动求导函数。第一步定义算子 Schema 与后端内核仓库中的完整示例位于 advanced_source/dispatcher/op.cpp。我们先来看整个注册的骨架。1. 集中声明 Schema与简单的 pybind11 式算子注册不同这一步并不提供任何实现只提供一个声明算子类型签名的 schema 字符串后续所有 kernel 都要遵守它TORCH_LIBRARY(myops, m) { m.def(myadd(Tensor self, Tensor other) - Tensor); }TORCH_LIBRARY的第一个参数是命名空间名myops第二个参数m是用于注册的句柄。m.def传入的myadd(Tensor self, Tensor other) - Tensor是完整的函数式 schema定义了算子名、参数与返回值类型。2. 实现 CPU 内核下面是一个极其简单的 CPU 加法实现完整源码见 op.cppTensor myadd_cpu(const Tensor self_, const Tensor other_) { TORCH_CHECK(self_.sizes() other_.sizes()); TORCH_INTERNAL_ASSERT(self_.device().type() DeviceType::CPU); TORCH_INTERNAL_ASSERT(other_.device().type() DeviceType::CPU); Tensor self self_.contiguous(); Tensor other other_.contiguous(); Tensor result torch::empty(self.sizes(), self.options()); const float* self_ptr self.data_ptrfloat(); const float* other_ptr other.data_ptrfloat(); float* result_ptr result.data_ptrfloat(); for (int64_t i 0; i result.numel(); i) { result_ptr[i] self_ptr[i] other_ptr[i]; } return result; }值得注意的实现细节先用TORCH_CHECK校验两个张量形状一致再用TORCH_INTERNAL_ASSERT断言设备类型为 CPUTORCH_INTERNAL_ASSERT表示调用方保证成立的内部不变量通过contiguous()确保内存布局连续再拿data_ptrfloat()直接做逐元素循环展示了一个手写 kernel的典型形态返回用torch::empty(self.sizes(), self.options())分配的输出张量options()继承了输入张量的 dtype/device 等选项。3. 用TORCH_LIBRARY_IMPL把内核绑定到 CPU key如果直接用def(myadd, myadd_cpu)这种简单方式注册会把该 kernel 注册为catch-all兜底kernel——即任何情况下都会运行哪怕输入张量根本不是 CPU 张量。为了确保myadd_cpu只对 CPU 张量运行要使用TORCH_LIBRARY_IMPL宏TORCH_LIBRARY_IMPL(myops, CPU, m) { m.impl(myadd, myadd_cpu); }TORCH_LIBRARY_IMPL允许我们针对特定 dispatch key这里是 CPU注册算子的实现。每次调用m.impl都会把某个 CPU kernel 与之前在TORCH_LIBRARY块中定义的算子关联起来。4. 注册 CUDA 内核如果我们还有一个 CUDA 实现myadd_cuda可以在另一个独立的TORCH_LIBRARY_IMPL块中注册它Tensor myadd_cuda(const Tensor self, const Tensor other) { // Insert your CUDA implementation here TORCH_CHECK(0, CUDA not yet implemented); } TORCH_LIBRARY_IMPL(myops, CUDA, m) { m.impl(myadd, myadd_cuda); }示例里的myadd_cuda只是占位实现调用即报错实际项目中在这里插入你的 CUDA kernel 即可。可以看到CPU 与 CUDA 的注册彼此独立、互不耦合这正是 Dispatcher 分层设计的直接体现。5. 注册结构的最佳实践这些注册可以跨文件甚至跨库边界拆分。例如两个TORCH_LIBRARY_IMPL块可以被分别编译进独立的myops_cpu与myops_cuda动态库。一般推荐的注册结构如下一个集中的TORCH_LIBRARY把命名空间下所有自定义算子集中列在一处每个 dispatch key 一个TORCH_LIBRARY_IMPL注册该 key 的实现如 CPU、CUDA。如果愿意还可以进一步把TORCH_LIBRARY_IMPL细化为每个算子一个块——当每个算子实现各占一个 cpp 文件、又不想在头文件里暴露算子时直接把注册写进定义算子的 cpp 文件里即可。你知道吗你也可以为 PyTorch 现有的核心算子写TORCH_LIBRARY_IMPL块。PyTorch 的 XLA 支持正是这样实现的torch_xla库包含一个TORCH_LIBRARY_IMPL为 XLA dispatch key 上的所有基础算子提供实现。6. 构建与加载示例仓库提供了构建配置与 Python 验证脚本CMakeLists.txt一个标准的 PyTorch C 扩展构建文件通过find_package(Torch REQUIRED)找到已安装的 PyTorch用add_library(dispatcher SHARED op.cpp)把 op.cpp 编译成共享库libdispatcher.so要求 C14。test.pyPython 侧验证脚本import torch torch.ops.load_library(build/libdispatcher.so) print(torch.ops.myops.myadd(torch.randn(32, 32), torch.rand(32, 32)))torch.ops.load_library把编译出的共享库加载进进程随后即可通过torch.ops.myops.myadd从 Python 前端直接调用我们注册的算子。脚本注释还说明了一个已知限制Python 前端torch.ops目前不支持带 names 的张量参数因此 Named 相关的验证代码被注释掉了。为无需自动求导的算子注册 Autograd 内核本节仅适用于PyTorch 1.10的版本。对于不需要自动求导支持的算子注册下面这样一个 kernel可以让你的算子表现得像 PyTorch 内置算子一样好用TORCH_LIBRARY_IMPL(myops, Autograd, m) { m.impl(op, autogradNotImplementedFallback()); }上面这段代码注册了一个Autogradkernel它在 forward 阶段追加一个哑NotImplemented节点从而保留输入的requires_grad属性在 backward 阶段这个NotImplemented节点会抛错。这对调试大型模型很有帮助——之前你可能很难定位requires_grad到底是在哪一步前向传播中丢失的有了这个内核问题会在第一个不该丢的地方立刻暴露。就地in-place或视图view算子还需要ADInplaceOrViewkernel为确保正确性与最佳性能如果你的算子会就地修改输入或者返回与输入张量共享存储别名的视图张量还需要额外做两步在注册上面的Autogradkernel 之外再注册一个ADInplaceOrViewkernel。该 kernel 负责就地/视图操作所需的簿记工作保证正确性。注意这个ADInplaceOrViewkernel只能配合autogradNotImplementedFallback使用TORCH_LIBRARY_IMPL(myops, Autograd, m) { m.impl(op, autogradNotImplementedFallback()); } TORCH_LIBRARY_IMPL(myops, ADInplaceOrView, m) { m.impl(op, autogradNotImplementedInplaceOrViewFallback()); }上面注册的这些 boxed kernel 的逻辑依赖算子 schema 信息。如果你的算子会就地修改输入或返回与输入别名的张量务必确保 schema 正确反映这一点——即正确使用 schema 中的别名/变更标注如Tensor(a!)、Tensor(b!) - Tensor(b)这类 AliasInfo 注解描述输入的 mutable/aliasing 关系否则后续行为无法保证。为算子添加自动求导支持现在我们有了同时带 CPU 与 CUDA 实现的算子如何给它加上自动求导支持思路与 C 自定义自动求导函数教程 类似——注册一个 autograd kernel。但有一个关键区别与 CPU/CUDA kernel 不同autograd kernel 需要 redispatch它必须回调 Dispatcher才能到达真正的推理内核CPU 或 CUDA 实现。编写 dispatch 函数算子公开的 C API在写 autograd kernel 之前先写一个dispatch 函数它负责调用 Dispatcher 找到适合当前调用的 kernel。这个函数构成了算子的公开 C API——事实上PyTorch C API 里所有张量函数在底层都是这样调用 Dispatcher 的。完整代码见 op.cppTensor myadd(const Tensor self, const Tensor other) { static auto op torch::Dispatcher::singleton() .findSchemaOrThrow(myops::myadd, ) .typeddecltype(myadd)(); return op.call(self, other); }逐行拆解第一行从 Dispatcher 查找与待调度算子对应的类型化算子句柄。findSchemaOrThrow接收两个参数算子的带命名空间的名称与 overload 名称通常为空字符串。随后typeddecltype(myadd)()把动态类型的句柄强制转换为静态类型句柄内部会做运行时校验确保你给定的 C 类型正确这样就能对它进行普通的 C 调用。这里传入decltype(myadd)因为 dispatch 函数的类型与注册到 Dispatcher 的底层 kernel 类型一致。性能细节这个查找被放在静态变量里只做一次较慢的查找。如果你把算子名拼错了会在第一次调用该函数时报错。第二行用op.call(self, other)以全部参数调用算子句柄。这才会真正触发 Dispatcher 调度最终把控制权转交给适合本次调用的 kernel。编写 autograd kernel有了 dispatch 函数就可以写 autograd kernel 了见 op.cppclass MyAddFunction : public torch::autograd::FunctionMyAddFunction { public: static Tensor forward( AutogradContext *ctx, torch::Tensor self, torch::Tensor other) { at::AutoNonVariableTypeMode g; return myadd(self, other); } static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs) { auto grad_output grad_outputs[0]; return {grad_output, grad_output}; } }; Tensor myadd_autograd(const Tensor self, const Tensor other) { return MyAddFunction::apply(self, other)[0]; }autograd 函数按常规方式用torch::autograd::Function编写唯一的不同在于forward()里不是直接写实现而是先用at::AutoNonVariableTypeMode这个 RAII 守卫关闭 autograd 处理再调用 dispatch 函数myadd回调 Dispatcher。为什么必须有第 1 步如果不关闭 autograd调用会无限循环直至栈溢出——因为myadd会把你送回这个 autograd kernel此时最高优先级的 dispatch key 仍是 Autograd。加上AutoNonVariableTypeMode后autograd 被排除出候选 dispatch key 集合调度就会落到下一层 handler也就是 CPU 或 CUDA 实现。backward中对加法算子很简单输出对两个输入各自的梯度都是grad_output因此直接返回{grad_output, grad_output}。注册 autograd kernel与注册 CPU/CUDA 函数的方式完全一致TORCH_LIBRARY_IMPL(myops, Autograd, m) { m.impl(myadd, myadd_autograd); }提示本例注册到Autograd它会把该 kernel 安装为所有后端共用的 autograd kernel。你也可以针对特定后端注册优化 kernel方法是用对应的后端专属 dispatch key例如AutogradCPU或AutogradCUDA。想更深入探索这些 dispatch key可以研究 PyTorch 提供的PythonDispatcher调试工具torch/_python_dispatcher.py。超越自动求导更多值得注册的 Dispatch Key从某种意义上说Dispatcher 本身没做太多事——它只是实现了一个升级版 if 语句。不用 Dispatcher上面的例子大概得写成这样class MyAddFunction : ... { public: static Tensor forward( AutogradContext *ctx, torch::Tensor self, torch::Tensor other) { if (self.device().type() DeviceType::CPU) { return add_cpu(self, other); } else if (self.device().type() DeviceType::CUDA) { return add_cuda(self, other); } else { TORCH_CHECK(0, Unsupported device , self.device().type()); } } ... }既然如此为什么还要用 Dispatcher原因有三去中心化Decentralized。你不需要写一个引用所有部件的集中式 if 语句就能把算子的所有部件CPU、CUDA、Autograd组装起来。更重要的是第三方可以不修改算子的原始定义就为其他方面注册额外实现。关于如何为一个全新后端扩展 Dispatcher参见 extend_dispatcher.rst该教程详细介绍了PrivateUse1/2/3保留 dispatch key、RegistrationDeclarations.h中的必注册/可选注册判定、以及后端专属 autograd kernel 的注册方式。支持的 dispatch key 远不止 CPU、CUDA、Autograd 三个。PyTorch 当前实现的完整 dispatch key 列表可以在c10/core/DispatchKey.h中查看。这些 key 实现了各种可选的算子功能如果你希望自定义算子支持某项功能只需为对应 key 注册一个 kernel。Dispatcher 实现了 boxed fallback盒装回退函数支持——这类函数实现一次即可应用于系统中的所有算子。Boxed fallback 可以为某个 dispatch key 提供默认行为一旦用 Dispatcher 实现算子也就自动加入了所有这些操作的 fallback。下面是一些你可能需要为其定义算子的具体 dispatch key。Autocast自动混合精度包装Autocastdispatch key 实现了自动混合精度AMP支持。一个 autocast wrapper kernel 通常在运行算子前把传入的float16或float32CUDA 张量转换到某种首选精度。例如浮点 CUDA 张量上的矩阵乘与卷积通常以float16运行更快、占用内存更少且不损害收敛性。Autocast wrapper 只在 autocast 启用上下文中生效。假设我们有一个假想的自定义矩阵乘算子mymatmul它的 autocast wrapper 及注册如下// Autocast-specific helper functions #include ATen/autocast_mode.h Tensor mymatmul_autocast(const Tensor self, const Tensor other) { c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast); return mymatmul(at::autocast::cached_cast(at::kHalf, self), at::autocast::cached_cast(at::kHalf, other)); } TORCH_LIBRARY_IMPL(myops, Autocast, m) { m.impl(mymatmul, mymatmul_autocast); }要点解析cached_cast(kHalf, tensor)如果tensor是 CUDA 且为float32则转换为float16否则原样返回。这保证了当网络用任意float16/float32CUDA 张量混合调用mymatmul时算子以float16运行而对非 CUDA、整数类型或float64输入调用不受影响。遵循原生算子的 eligibility policy算子资格策略来写自己的 wrapper 是推荐做法而非强制。例如如果想强制所有输入类型都以float16执行也可以写成return mymatmul(self.half(), other.half());而不是用cached_cast。与 autograd kernel 一样在 redispatch 之前要用ExcludeDispatchKeyGuard把Autocastkey 从调度候选集合中排除。默认行为如果不提供 autocast wrapper调度会直接 fallthrough 到普通算子实现即不发生任何 autocasting。本教程不用myadd举例正是因为逐元素加法不需要 autocasting应当直接 fallthrough。什么时候该注册 autocast wrapper遗憾的是算子的首选精度没有放之四海而皆准的规则但可以参考以下通用指南做reduction归约的算子通常应执行在float32底层包含convolution 或 gemm的算子通常应执行在float16其他有多个浮点张量输入的算子应把输入统一到共同精度除非实现本身支持不同精度的输入。如果你的自定义算子属于第三类promote_type模板可以帮你求出输入张量中最宽的浮点类型作为最安全的执行类型#include ATen/autocast_mode.h Tensor my_multiple_input_op_autocast(const Tensor t0, const Tensor t1) { c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast); // The required at::kHalf argument is an optimistic initial guess. auto exec_type at::autocast::promote_type(at::kHalf, t0, t1); return my_multiple_input_op(at::autocast::cached_cast(exec_type, t0), at::autocast::cached_cast(exec_type, t1)); }如果你的自定义算子已启用 autograd即注册了 autograd wrapper只需为注册了 autograd wrapper 的同一个算子名编写并注册 autocast wrapper 即可。例如为前文 autograd 章节中的myadd添加 autocast wrapper只需要Tensor myadd_autocast(const Tensor self, const Tensor other) { c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast); return myadd(at::autocast::cached_cast(desired dtype, self), at::autocast::cached_cast(desired dtype, other)); } TORCH_LIBRARY_IMPL(myops, Autocast, m) { m.impl(myadd, myadd_autocast); }让 backward 方法兼容 autocast不需要任何额外技巧。但要注意你在自定义 autograd 函数中定义的 backward 方法会以 autocast 为 forward 设定的同一 dtype运行所以应选择一个对 forward 和 backward都合适的desired dtype。Batchedvmap批处理规则Batched 张量允许你以逐样本的方式编写代码然后在vmap调用下自动完成批处理。batching rules 的 API 目前仍在开发中一旦稳定你就可以通过在Batcheddispatch key 注册 kernel来为自己的算子添加vmap支持。Tracertorch.jit.trace追踪Tracerdispatch key 实现了在运行torch.jit.trace时把算子调用记录进 trace的支持。PyTorch 官方计划提供一个 boxed fallback为任意算子实现追踪可以在 PyTorch 官方仓库的 issue #41478 中跟踪进展。此外仓库源码 op.cpp 中还保留了一段被#if 0禁用的Nameddispatch key 注册示例myadd_named使用at::unify_from_right统一维度名并传播输出名说明像Named张量维度命名这样的扩展 key 同样可以通过TORCH_LIBRARY_IMPL接入——当你需要为自定义算子补齐维度名支持时这是一个可以直接参考的模板。小结与建议路线结构上一个命名空间一个TORCH_LIBRARY集中声明 schema每个 dispatch key 一个或多个TORCH_LIBRARY_IMPL块把 CPU/CUDA 推理内核、Autograd求导内核、Autocast包装内核、Batched/Tracer等可选内核分别注册到位。自动求导autograd kernel 内必须用at::AutoNonVariableTypeMode关闭 autograd 后调用 dispatch 函数完成 redispatch否则会无限递归不要求导的算子可注册autogradNotImplementedFallback就地/视图算子还需ADInplaceOrViewkey 与 schema 别名标注。何时不必使用 Dispatcher纯 CPU 推理、极其简单的算子直接用基础注册 API 即可只有当算子需要跨设备、求导、AMP、vmap、追踪等横切能力时才值得引入这套分层机制。后续进阶为全新后端扩展调度器请继续阅读 extend_dispatcher.rstPyTorch 2.4 之后的最新自定义算子实践请以 custom_ops_landing_page.rst 及其分册教程为准。赞分享示例工程【免费下载链接】tutorialsPyTorch tutorials.项目地址https://gitcode.com/gh_mirrors/tuto/tutorials点击查看免费下载相关推荐PyTorch torch.library 自定义算子实战指南从 opcheck 测试到注册内核的完整教程PyTorch torch.library 自定义算子实战指南从 opcheck 测试到注册内核的完整教程 torch.library 是 PyTorch 官人工智能机器学习深度学习分布式训练模型编译PyTorch tutorials高级教程自定义C扩展与CUDA算子开发PyTorch tutorials高级教程自定义C扩展与CUDA算子开发 你是否在PyTorch模型部署中遇到过性能瓶颈当Python实现的算子无法满足示例工程PyTorch 自定义算子Custom Operators实战指南torch.library 与 TORCH_LIBRARY 注册体系全解析PyTorch 自定义算子Custom Operators实战指南torch.library 与 TORCH_LIBRARY 注册体系全解析 PyTorc示例工程上一篇【亲测免费】 探秘Hero一款强大的UI动画库下一篇Preact深度解析轻量级React替代方案的架构设计创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考