【PTA】昇腾 TorchNPU 技术解析：从工程到数据的全链路实现

某某咸鱼

12人浏览 · 2026-10-24 14:37:46

某某咸鱼 · 2026-10-24 14:37:46 发布

背景概述

随着深度学习框架与异构计算平台的深度融合，PyTorch 作为主流训练框架，其对多样化硬件后端（如 NPU）的支持需求日益增长。torch_npu 作为连接 PyTorch 与 NPU（神经网络处理器）的桥梁，承担着算子分发、设备管理、数据流转和资源调度等核心职责。本文从工程实现、数据流转和资源管理三个视角，系统梳理 torch_npu 的架构设计与关键实现路径，帮助开发者理解其工作原理与接入方式。

一、工程视角：torch_npu 如何接入 PyTorch 社区

1. 接入社区的核心机制：注册与 Patch

torch_npu 接入 PyTorch 社区的核心思路是“注册”与“Patch”，主要分为两个层面：

算子层面：计算算子与通信算子通过 PyTorch 的 dispatch 机制分发到 NPU 实现。
设备层面：device、stream、event 等设备抽象通过 Patch 方式扩展。

2. 基于 PrivateUse1 后端的接入流程

torch_npu 利用 PyTorch 的 PrivateUse1 后端机制完成注册：

在 torch_npu/__init__.py 中将 PrivateUse1 命名为 npu，并将 torch_npu.npu 注册到 torch.npu。
为 tensor、module、storage 绑定 .npu() 方法，实现设备迁移接口。

在 torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp 中完成以下注册：

注册 NPUGuardImpl，负责 device 与 stream 的 guard 管理。
注册 storage 创建函数。
注册 hooks，相关接口见 torch_npu/csrc/core/npu/NPUHooksInterface.h。
注册序列化与反序列化函数。

此外，通过注册 HCCL 后端，init_process_group 可创建 ProcessGroupHCCL，使社区通信后端能够分发到 NPU 的 HCCL 实现。

3. Python 到 C++ 的调用链路

torch_npu 的 Python 层通过 CPython 绑定机制调用 C++ 实现，核心构建配置位于 setup.py，C++ 源码位于 torch_npu/csrc/npu。

4. C++ 层接入 CANN 的方式

C++ 层通过“打桩”与 dlopen 动态加载机制接入 CANN（华为 AI 计算框架）：

打桩脚本：third_party/acl/libs/build_stub.sh
动态加载实现：torch_npu/csrc/core/npu/register/FunctionLoader.cpp
接口定义：torch_npu/csrc/core/npu/interface

5. 其他工程配置

环境变量管理：torch_npu/csrc/core/npu/register/OptionsManager.cpp
Option 配置：torch_npu/csrc/core/npu/register/OptionRegister.cpp
环境变量定义：torch_npu/csrc/framework/interface/EnvVariables.cpp

二、数据视角：数据如何在 NPU 上流转与操作

1. 数据流转到 NPU 的路径

模型训练与推理的数据载体是 tensor，其内部包含 storage，而 storage 中的 data_ptr 指向实际数据内存。

数据流转到 NPU 的典型路径如下：

当 tensor 被创建或通过 .to('npu:0')、.npu() 转移到 NPU 时，会携带 device 信息。
通过 dispatch 分发到 NPU 实现，进入 tensor 工厂函数（如 empty）。
在 NPU 侧使用 NPU 的 allocator 申请设备内存，封装为 NPU 的 storage，最终创建 NPU tensor。

此外，也支持直接创建 NPU storage，相关实现见：

torch_npu/csrc/aten/common/TensorFactories.cpp
torch_npu/csrc/core/NPUStorageImpl.cpp

2. NPU 私有格式的接入方式

NPU 的私有格式信息记录在 NPU 的 storage 中。创建 tensor 或 storage 时，会使用 NPU storage 的创建逻辑，并携带私有格式信息，相关定义见 torch_npu/csrc/core/NPUStorageImpl.h。

3. 操作 NPU 数据的方式

当数据位于 NPU 上后，tensor 中的 device 信息携带 NPU 标识，普通计算算子通过 dispatch 分发到 NPU 算子实现。

aclop 算子：通过 OpCommand 中的接口将 at::Tensor、Scalar 等数据转换为 ACL 相关数据结构（如 aclTensorDesc、aclDataBuffer），最终调用 aclopCompileAndExecute 或 AclopCompileAndExecuteV2 执行操作。相关实现见：
- torch_npu/csrc/framework/OpCommand.cpp
- torch_npu/csrc/framework/OpParamMaker.h
aclnn 算子：通过 ConvertTypes 等模板函数将 at::Tensor、Scalar 等数据转换为 aclTensor，最终调用 aclnn 接口执行操作。相关实现见：
- third_party/op-plugin/op_plugin/utils/op_api_common.h
- third_party/op-plugin/op_plugin/utils/op_api_common.cpp
- third_party/op-plugin/op_plugin/utils/op_api_common_base.h
集合通信算子：使用 HCCL 后端，进入 ProcessGroupHCCL 中的 HCCL 算子实现，将输入输出的数据指针地址传递给 HCCL 算子执行。相关实现见 torch_npu/csrc/distributed/ProcessGroupHCCL.cpp。

三、资源视角：NPU 资源的管理与保障

1. NPU 资源类型

对于 torch_npu 而言，主要涉及以下资源：

device：设备上下文（context）默认使用设备上的默认 context，若涉及 context 切换，必须切回原 context，否则会导致上下文不匹配。
stream：执行流，用于控制任务执行顺序。
event：事件，用于流间同步。
memory：包括设备物理内存、虚拟内存，以及主机侧锁页内存（pin_memory）。

更多信息可参考 CANN 官方文档。

2. NPU 资源的初始化与管理

PyTorch 社区对设备资源仅做简单初始化（如 lazyInitDevice 初始化内存池、P2P 访问等），而 torch_npu 有明确的初始化与去初始化要求：

初始化：参考社区的 _lazy_init，torch_npu 中调用 _npu_init，内部调用 NpuSysCtrl::Initialize，对应社区的 lazyInitDevice。此外，还会调用 getDefaultNPUGenerator 等初始化逻辑。为避免不必要的 context 资源浪费，需要 context 的初始化配置被延迟到流创建或 aclop 调用时进行。
去初始化：借助 Python 的 atexit 机制注册 shutdown 回调，相关实现见：
- torch_npu/__init__.py
- torch_npu/csrc/InitNpuBindings.cpp 中的 THPModule_npu_shutdown
- torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.cpp 中的 Finalize()

需要注意的是，atexit 仅对 spawn 方式启动的子进程生效，fork 方式启动的子进程不会触发。因此，若使用 fork 方式创建子进程并涉及 NPU 资源，需确保同步操作，避免异步执行导致错误。

初始化完成后，可像社区一样申请 stream、event、memory 等资源，相关接口见：

torch_npu/npu/__init__.py
torch_npu/npu/streams.py
torch_npu/npu/utils.py

3. NPU 资源的正确性保障

位置正确性

通过 guard 机制（如 DeviceGuard、StreamGuard）确保资源在正确的设备与流上使用。除了显式代码中的 guard，codegen 生成的算子适配中也内置了 guard，这也是为什么无需显式 set_device，算子仍能在对应设备上正确计算的原因之一。相关实现见：