在深度学习研究领域，PyTorch是不可或缺的工具。PyTorch是一个基于Python的深度学习框架，同时也是一个基于GPU和CPU加速的张量计算库。原生PyTorch并不天然支持NPU，昇腾工程师提供PyTorch Adapter模块来使PyTorch能在NPU设备上适配。本章就来讲讲PyTorch和PyTorch Adapter的原理，为后面实现NPU设备上PyTorch模型迁移与调优打下基础。
PyTorch Adapter
原生的PyTorch分为GPU版本与CPU版本，而我们所需要的NPU版本的PyTorch本质是将CPU版本的PyTorch 适配NPU，利用NPU实现计算加速，对标GPU版本PyTorch。对于PyTorch1.5是侵入式修改，1.8和1.11及后续版本是插件化支持，也就是会有一个torch_npu包，aka. PyTorch Adapter，概念上等同于CUDA版本的PyTorch。那PyTorch Adapter是在哪里做出修改，又在哪里起到什么作用呢？这需要我们对原生PyTorch有一定了解。
原生PyTorch
下图是原生PyTorch的代码结构，主要分为4个部分：
- torch/，包含最熟悉的内容：导入和使用的实际的Python模块，为我们所调用的api。
- torch/csrc/，对应Python接口的C++实现。具体来讲，它实现了在Python和C++领域之间转换的绑定代码，以及PyTorch的一些非常重要的部分，如autograd引擎和JIT编译器。它还包含C++前端代码。
- aten/，“A Tensor Library”的缩写，是一个实现张量操作的C++库。所有调用的算子，CPU算子、CUDA算子的实现都是在这里(.c，.cu)，如果要新写一个算子，也写在这里。aten本身分为两个运算符邻域：“原生”运算符，这是运算符的现代C++实现，以及“遗留”运算符（TH、THC、THNN、THUNN），这是遗留的C实现。
- c10/，包含PyTorch的核心抽象，包括张量和存储数据结构的实际实现。
PyTorch单算子调用流程
-
用户在前端使用算子，可通过nn.Module, nn.Functional, Tensor对象上的函数，torch包函数调用
-
PYBIND11根据注册绑定的映射规则，调用后端C++方法；
-
调用VariableType的计算API;
-
VariableType是Variable的类型定义，即：具有自动求导特性的tensor。 VaraibleType对象提供所有function（算子）粒度的API可供调用。 VaraibleType中的计算API的主要动作：
- 启动Profiling，记录算子调用时间；（有开关）
- 生成Trace记录，用于导出graph；（有开关）
- 调用ATen计算库中的算子计算接口层API；这里会根据input的dtype区分CPU Tensor还是CUDATensor，分发调用CPUType的算子接口API或者CUDAType的算子接口API。
- 根据c返回的output，计算grad_fn，附加在output中。
-
CPU/CUDAType作为算子的接口层，判断本类型\精度是否支持该算子，支持的话则调用计算层TH/THC。
-
计算层TH/THC完成输出内存分配、调用计算库/加速库API、释放输入内存。
从上图可以看出，要适配NPU设备，我们需要在ATen模块里注册NPU的接口实现，调用CANN的能力，这就是PyTorch Adapter做的事情。而PyTorch Adapter的实现，依赖于pytorch的dispatch机制，这部分内容可以参考网上资料：https://zhuanlan.zhihu.com/p/376495783
对于每个operator，dispatcher都会维护一个函数指针表，为每个dispatch key提供对应的实现。这个表中有针对不同后端（CPU、CUDA、XLA）的dispatch条目，也有像autograd和tracing这样的高抽象层级概念的条目。dispatcher的工作是根据输入的tensor和其他一些东西（接下来会提及）计算出一个dispatch key，然后跳转到表所指向的函数。
为了实现分发的能力，提供了三个接口如下图：
只需要将算子注册到dispatcher上，即可复用PyTorch的dispatch机制完成算子分发 与operator registration API交互的三种主要路径： 为op定义schema； 在dispatch key处注册实现； 存在一个fallback方法，你可以用它来为某个key的所有op定义一个handler。
逻辑架构
我们现在有了CANN提供计算能力，有了PyTorch的分发-注册机制，那么只需要在dispatcher上注册npu的key，并对所有算子进行适配，就能调用acl了。
用户python脚本层：定义训练或推理的执行流程； PyTorch的python接口层：分为原生python接口和npu新增python接口； CPython层：定义python到C++实现的映射关系； npu适配层：资源管理，算子适配和下发，特定的功能模块。 资源管理涉及对内存资源，stream资源以及event资源的申请、二次管理、接口封装以及释放控制； 算子适配和下发模块是在原生的算子适配框架下，把PyTorch的ir转为ge的ir，同时申请所需的资源，下发异步task； 特定功能模块是相对独立的模块，后续会随着用户需求或者PyTorch框架演进不断扩展； 底层接口层：npu适配代码直接依赖acl、ge和hccl； 其中资源管理以及单算子模式主要依赖acl接口，图模式新增了对ge接口的调用，hccl接口专用于分布式流程；需要注意的是，与cuda的不同，cuda算子的实现是在pytorch里面的，基于NPU设备的算子是在cann里面实现，在pytorch里留有接口。
参考资料
https://zhuanlan.zhihu.com/p/376495783
http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/
点赞
收藏
回复
分享
举报
在深度学习研究领域，PyTorch是不可或缺的工具。PyTorch是一个基于Python的深度学习框架，同时也是一个基于GPU和CPU加速的张量计算库。原生PyTorch并不天然支持NPU，昇腾工程师提供PyTorch Adapter模块来使PyTorch能在NPU设备上适配。本章就来讲讲PyTorch和PyTorch Adapter的原理，为后面实现NPU设备上PyTorch模型迁移与调优打下基础。
PyTorch Adapter
原生的PyTorch分为GPU版本与CPU版本，而我们所需要的NPU版本的PyTorch本质是将CPU版本的PyTorch 适配NPU，利用NPU实现计算加速，对标GPU版本PyTorch。对于PyTorch1.5是侵入式修改，1.8和1.11及后续版本是插件化支持，也就是会有一个torch_npu包，aka. PyTorch Adapter，概念上等同于CUDA版本的PyTorch。那PyTorch Adapter是在哪里做出修改，又在哪里起到什么作用呢？这需要我们对原生PyTorch有一定了解。
原生PyTorch
下图是原生PyTorch的代码结构，主要分为4个部分：
PyTorch单算子调用流程
用户在前端使用算子，可通过nn.Module, nn.Functional, Tensor对象上的函数，torch包函数调用
PYBIND11根据注册绑定的映射规则，调用后端C++方法；
调用VariableType的计算API;
VariableType是Variable的类型定义，即：具有自动求导特性的tensor。 VaraibleType对象提供所有function（算子）粒度的API可供调用。 VaraibleType中的计算API的主要动作：
CPU/CUDAType作为算子的接口层，判断本类型\精度是否支持该算子，支持的话则调用计算层TH/THC。
计算层TH/THC完成输出内存分配、调用计算库/加速库API、释放输入内存。
从上图可以看出，要适配NPU设备，我们需要在ATen模块里注册NPU的接口实现，调用CANN的能力，这就是PyTorch Adapter做的事情。而PyTorch Adapter的实现，依赖于pytorch的dispatch机制，这部分内容可以参考网上资料：https://zhuanlan.zhihu.com/p/376495783
对于每个operator，dispatcher都会维护一个函数指针表，为每个dispatch key提供对应的实现。这个表中有针对不同后端（CPU、CUDA、XLA）的dispatch条目，也有像autograd和tracing这样的高抽象层级概念的条目。dispatcher的工作是根据输入的tensor和其他一些东西（接下来会提及）计算出一个dispatch key，然后跳转到表所指向的函数。
为了实现分发的能力，提供了三个接口如下图：
只需要将算子注册到dispatcher上，即可复用PyTorch的dispatch机制完成算子分发 与operator registration API交互的三种主要路径： 为op定义schema； 在dispatch key处注册实现； 存在一个fallback方法，你可以用它来为某个key的所有op定义一个handler。
逻辑架构
我们现在有了CANN提供计算能力，有了PyTorch的分发-注册机制，那么只需要在dispatcher上注册npu的key，并对所有算子进行适配，就能调用acl了。
用户python脚本层：定义训练或推理的执行流程； PyTorch的python接口层：分为原生python接口和npu新增python接口； CPython层：定义python到C++实现的映射关系； npu适配层：资源管理，算子适配和下发，特定的功能模块。 资源管理涉及对内存资源，stream资源以及event资源的申请、二次管理、接口封装以及释放控制； 算子适配和下发模块是在原生的算子适配框架下，把PyTorch的ir转为ge的ir，同时申请所需的资源，下发异步task； 特定功能模块是相对独立的模块，后续会随着用户需求或者PyTorch框架演进不断扩展； 底层接口层：npu适配代码直接依赖acl、ge和hccl； 其中资源管理以及单算子模式主要依赖acl接口，图模式新增了对ge接口的调用，hccl接口专用于分布式流程；需要注意的是，与cuda的不同，cuda算子的实现是在pytorch里面的，基于NPU设备的算子是在cann里面实现，在pytorch里留有接口。
参考资料
https://zhuanlan.zhihu.com/p/376495783
http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/