1 Ascend Extension for PyTorch

Ascend Extension for PyTorch 插件是基于昇腾的深度学习适配框架，使昇腾 NPU 可以支持 PyTorch 框架，为 PyTorch 框架的使用者提供昇腾 AI 处理器的超强算力。

项目源码地址请参见Ascend/Pytorch。

昇腾为基于昇腾处理器和软件的行业应用及服务提供全栈 AI 计算基础设施。您可以通过访问昇腾社区，了解关于昇腾的更多信息。

2 为什么需要此扩展？

PyTorch 本身是个模型训练和推理的 AI 框架，与华为的昇思（mindspore）和百度的飞浆（paddlepaddle）等是同样的产品形态。pytorch 作为业界的顶流，无论是在业界，还是高校或者说是 AI 相关社区，其重要性都是无法撼动的。但是，pytorch 本身只对 cpu 和 gpu 进行了较好的适配，这就导致了新出现的 npu 硬件产品，如果要在 pytorch 上进行无缝对接，肯定得进行一些适配。比如：如何调用新 npu 硬件上配套的高性能算子？

3 Ascend 在 pytorch 上的对接方案

以下架构图来源于昇腾社区。Ascend Extension for PyTorch

该在线对接适配方案的特点包含：

• 最大限度的继承 PyTorch 框架动态图的特性。

• 最大限度的继承原生 PyTorch 的开发方式，可以使用户在将模型移植到昇腾 AI 处理器设备进行训练时，在开发方式和代码重用方面做到最小的改动。

• 最大限度的继承 PyTorch 原生的体系结构，保留框架本身出色的特性，比如自动微分、动态分发、Debug、Profiling、Storage 共享机制以及设备侧的动态内存管理等。

• 扩展性好。在打通流程的通路之上，对于新增的网络类型或结构，只需涉及相关计算类算子的开发和实现。框架类算子，反向图建立和实现机制等结构可保持复用。

• 与原生 PyTorch 的使用方式和风格保持一致。用户在使用在线对接方案时，只需在 Python 侧和 Device 相关操作中，指定 Device 为昇腾 AI 处理器，即可完成用昇腾 AI 处理器在 PyTorch 对网络的开发、训练以及调试，用户无需进一步关注昇腾 AI 处理器具体的底层细节。这样可以确保用户的修改最小化，迁移成本较低。

根据上图，可以获取以下几点内容（全部是个人理解，如果错了还望指正）：

1. 图中标识 1 的部分，这个是业务层，比如训练，推理等；

2. 图中标识 2 的部分，是 torch 的原有 python 接口，方便用户侧调用。

3. 图中标识 3 的部分，ascend 对外提供的 python 接口，比如在分布式下，ascend 是 hccl，而 gpu 是 nccl；在用户脚本侧，就得调用对应的接口；

4. 图中标识 4 的部分，这个是 pybind 层。熟悉 pytorch 的这层作用很明了。python 侧只是用户 python 中脚本呈现的，真正下发到硬件上去执行，都是 c++侧下发。从 python 到 c++，就有一个 pybind 的过程，也就是 cpython。

5. 图中标识 5 的部分，是框架中的通用层，比如算子关系，tensor 实现，memory 申请实现等；

6. 图中标识 6 的部分，这个是 ascend 上特有的，如 aclop 需要编译然后执行（aclop 特有），就是将算子的执行逻辑编译成一个.o 的执行文件，让 c++能够直接调用；

7. 图中标识 7 的部分，算子执行库；也就是 ACL（Ascend c）编译后的算子执行库；

8. 图中标识 8 的部分，hccl（集合通信库）执行库；

9. 图中标识 9 的部分，ascend 的 rumtime。runtime 的作用，也就是给你一个 op 或者说一张图，让其在硬件升运行起来。

下一步，将从源码的角度来看看具体的实现逻辑，相关文章如下：Ascend Extension for PyTorch的源码解析