​作者​昇腾实战派​知识地图​https://blog.csdn.net/Lumos_Lovegood/article/details/161601003背景概述PyTorch 作为高度抽象化的深度学习框架用户在使用时更关注计算逻辑的正确性与计算结果的准确性。然而实际计算在硬件上的执行是由一系列算子Operator完成的这些算子是神经网络计算的实际执行者负责与底层硬件交互。由于硬件环境和使用方式的不同同一计算逻辑往往需要对应不同的算子实现。例如PyTorch 中的nn.Conv2D在 CPU、GPU 或 NPU 上运行时会衍生出不同的实现。PyTorch 负责管理这一调度过程的机制称为Dispatch其核心任务是将计算任务分发给合适的算子实现。Dispatch 的整体流程可概括为输入一个 ATen 函数名如aten::conv2D表示抽象的计算逻辑根据系统配置如native_functions.yaml和 Tensor 的具体属性如数据类型、设备类型找到实际的计算函数。该过程分为三个阶段类似于一个决策树Dispatcher 阶段根据 Tensor 的稀疏性及DispatchKeySet选择计算入口函数分流至 Sparse 或非 CPU/GPU 设备的计算函数。DispatchStub 阶段根据 Tensor 所处的设备类型选择实际计算函数。Dispatch 数据类型阶段根据 Tensor 的数据类型实例化实际计算逻辑。在实际开发中有时需要了解 PyTorch API 底层调用的算子逻辑以便进行性能分析、问题排查或自定义扩展。以下介绍几种常用的检测机制。检测机制一通过环境变量打印算子日志1. 打印 aclnn 算子信息设置环境变量TORCH_NPU_LOGSop_plugin可以在算子调用过程中获取并打印 aclnn 算子相关信息。相关实现可参考以下文件third_party/op-plugin/op_plugin/utils/op_log.hthird_party/op-plugin/op_plugin/utils/op_api_common.h使用方式exportTORCH_NPU_LOGSop_plugin2. 打印 TaskQueue 任务信息设置环境变量TORCH_NPU_LOGSdispatch可以在 TaskQueue 入队和出队时获取任务信息。相关实现可参考torch_npu/csrc/core/npu/NPUQueue.cpp该机制可以打印所有经过 TaskQueue 的任务信息可通过搜索以下关键字进行过滤WriteQueue: write successReadQueue: read success使用方式exportTORCH_NPU_LOGSdispatch检测机制二通过__torch_dispatch__拦截 ATen 级别算子__torch_dispatch__是 PyTorch 提供的一个底层拦截机制允许用户在算子最终 dispatch 之前获取对应的算子及其输入。PyTorch 提供了TorchDispatchMode来优雅地实现这一功能用户只需实现其内部的__torch_dispatch__方法即可在算子执行前进行拦截从而对输入和输出进行自定义处理。以下示例展示了如何使用TorchDispatchMode拦截并打印nn.Linear实际调用的算子importtorchimporttorch.nnasnnfromtorch.utils._python_dispatchimportTorchDispatchModeclassOpTracer(TorchDispatchMode):def__torch_dispatch__(self,func,types,args(),kwargsNone):print(f[DISPATCH]{func})returnfunc(*args,**(kwargsor{}))modelnn.Linear(128,64).npu()xtorch.randn(32,128).npu()withOpTracer():ymodel(x)运行后会打印出实际经过 dispatch 的每一个算子例如[DISPATCH] aten::t [DISPATCH] aten::addmm这是最直接的方式能够清晰地看到nn.Linear实际走了哪条路径。检测机制三通过torch.jit.trace Dispatch Key 追踪使用torch.jit.trace结合 dispatch key 追踪可以精确查看算子的注册情况。以下示例展示了如何查看aten::linear的 dispatch 表importtorchimporttorch.nnasnn modelnn.Linear(128,64).npu()xtorch.randn(32,128).npu()# 查看 aten::linear 的 dispatch 表print(torch._C._dispatch_dump(aten::linear))从输出可以看到aten::linear注册在CompositeImplicitAutograd别名 key意味着它会被分解为更底层的算子aten::addmm或aten::mm。aten::mm在PrivateUse1即 NPU上注册了实现位于torch_npu/csrc/aten/RegisterNPU.cpp。因此nn.Linear的调用链路为nn.Linear→aten::linear→aten::addmm/aten::mm→ NPU kernel并不经过npu_linear。检测机制四AOTAutograd 追踪正向与反向传播计算图AOTAutograd 以 Ahead-of-Time 的方式同时追踪正向传播和反向传播从而在函数真正执行之前获取正向和反向传播的计算图。其工作流程如下通过__torch_dispatch__机制以 AOT 方式追踪正向传播和反向传播生成联合计算图joint forward and backward graph该图是包含 ATen/Prim 算子的 FX Graph。使用partition_fn将联合计算图划分为正向传播计算图和反向传播计算图。可选通过decompositions将高层算子分解、下沉到粒度更小的算子。调用fw_compiler和bw_compiler分别编译正向传播计算图和反向传播计算图通过 TorchFX 生成编译后的 Python 代码并整合为一个torch.autograd.Function。以下是一个典型的输出示例展示了正向传播和反向传播计算图对应的 Python 代码add_2 None neg_1 torch.ops.aten.neg.default(sin_1); sin_1 None mul_1 torch.ops.aten.mul.Tensor(mul, neg_1); mul neg_1 None return [mul_1, mul_1, mul_1, mul_1]自定义的编译器compiler_fn()会被调用两次分别打印正向传播和反向传播计算图对应的 Python 代码。其中的primals和tangents是微分几何中的概念primals可理解为用户函数的输入正向传播的输入tangents可理解为用户函数输出的梯度反向传播的输入。两张计算图均为 FX Graph其中包含的是 ATen 算子属于 low-level 算子而非 Torch 级别的算子如Linear。总结本文介绍了四种检测 PyTorch API 底层算子逻辑的机制机制适用场景特点环境变量日志快速查看算子调用信息简单易用无需修改代码__torch_dispatch__拦截并自定义处理算子灵活可打印每个算子torch.jit.trace Dispatch Key精确查看算子注册情况适合分析算子调度链路AOTAutograd同时追踪正向和反向传播适合高级分析和自定义编译开发者可根据实际需求选择合适的机制以深入了解 PyTorch 算子的底层执行逻辑。