10. MUSA AI 框架算子适配及 CV 图像处理开发示例
10.1. muPyTorch-muDNN 对接扩展
muPyTorch-muDNN对接扩展,可以让PyTorch框架使用muDNN算子库进行扩展,从而可以让用户在不对模型进行较大改动的情况下,使用musa-plugin在MTGPU上运行并计算AI模型。本节主要关注:PyTorch的dispatch机制、算子调用以及开发者如何进行muPyTorch与muDNN算子库的对接。对接完成后,PyTorch中的算子就可以通过muDNN算子库,运行在MTGPU设备上,进行算子计算和模型训练及推理。
10.1.1. 算子派发机制
PyTorch中有很多算子,同一种算子又分为很多类型。以add算子为例,按device类型可以分为CPU、GPU、NPU、TPU、FPGA等设备;按layout类型可以分为普通张量和稀疏张量,不同的张量有不同的布局;按dtype类型可以分为fp64、fp32、fp16、int8,甚至bf16、hf32等不同类型。那么当python端调用一个torch.add算子时,最后实际执行的算子该怎么选择呢?此时,就要提到PyTorch的dispatch机制了,接下来便对dispatch机制进行具体介绍。