PyTorch Adapter

概况

PyTorch Adapter：用于昇腾适配PyTorch框架，为使用PyTorch框架的开发者提供昇腾AI处理器的超强算力。
课程地址：https://www.hiascend.com/zh/developer/courses/detail/2203141624362937328
PyTorch Adapter作为昇腾CANN与原生PyTorch的桥梁，使得PyTorch能够调用昇腾的算力。（可以理解为一个适配器或插件）

PyTorch Adapter作为对PyTorch(CPU)的补充，可在概念上等价于PyTorch(GPU)。

• CUDNN是基于CUDA的深度学习GPU加速库
• 昇腾Driver+Firmware对标NVIDIA的Driver
• Ascend910/Ascend310为昇腾处理芯片
• V100/A100为NVIDIA用于深度学习机器学习的AI加速器

实现介绍

了解昇腾PT Adapter之前可以先看一下原生pytorch的设计。

PyTorch内部机制：
http://blog.ezyang.com/2019/05/pytorch-internals/
https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md#codebase-structure

PyTorch的目录结构：
torch/
torch/csrc/
aten/src/ATen/
c10/

• torch 真正的Pytorch库，所有的不在csrc目录里的都是python模块，遵循Pytorch前端模块结构
• csrc Pytorch的C++库，这个目录的库混合了python绑定代码，与C++的重型后端，参考setup.py来理解python的代码组织逻辑
• aten Pytorch的C++ tensor库。native算子的最新实现，如果你想写一个新的算子，应该在这里写，大多数CPU算子被挪到了顶层目录，除了那些需要特殊编译的算子
• C10 必须的核心组件，在Server端与移动端都可以工作

PyTorch在CPU-CUDA上利用dispatch机制实现异构计算。该机制实现如：

• 在设备输入tensor的情况下，实现CPU和CUDA算子实现的切换
• 在自动求导是否必要的情况下，实现自动求导和后端算子实现的切换
• 必要时应用自动转换以实现自动混合精度

pytorch中的dispatcher：
https://zhuanlan.zhihu.com/p/376495783

对于每个operator（算子），dispatcher都会维护一个函数指针表，为每个dispatch key提供对应的实现。这个表中有针对不同后端（CPU、CUDA、XLA）的dispatch条目，也有像autograd和tracing这样的高抽象层级概念的条目。dispatcher的工作是根据输入的tensor和其他一些东西（接下来会提及）计算出一个dispatch key，然后跳转到表所指向的函数。
因此PyTorch Adapter只需要将算子注册到dispatcher上，即可复用PyTorch的dispatch机制完成算子分发。
与operator registration API（算子注册的API）交互的三种主要路径：

• 为op定义schema
• 在dispatch key处注册实现
• 存在一个fallback方法，你可以用它来为某个key的所有op定义一个handler

CANN提供计算能力，通过PyTorch的分发-注册机制，只需要在dispatcher上注册npu的key，并对所有算子进行适配，就能调用acl了。