实践教程｜PyTorch与torch-xla的桥接

今日应用

艺映AI

释放每一帧的创造力和创新，免费AI视频生成平台，专业创建文生视频和AI动态短视频。将文本转换为高质量AI视频，适用于多种场景，如抖音短视频、小说推文、AI短片、AI电影等制作。

今日话题

实践教程｜PyTorch与torch-xla的桥接

 

重点标签 tag1、XLA、PyTorch、torch-xla、编译器

文章摘要

极市导读：文章从XLATensor开始的溯源、注册PyTorch库实现、从PyTorch调用到torch_XLA三个方面来介绍PyTorch与torch-xla的桥接。

前言：XLA (Accelerated Linear Algebra)是一个开源的机器学习编译器，对PyTorch、Tensorflow、JAX等多个深度学习框架都有支持。torch-xla在支持XLA编译的基础上，较大限度地保持了PyTorch的易用性。

PyTorch与torch-xla的桥接：知晓了Trace过程之后，就会好奇一个问题：当用户执行一个PyTorch函数调用的时候，torch-xla怎么将这个函数记录下来的？

从XLATensor开始的溯源：尽管我们现在并不知道怎么调用到torch-xla中的，但我们知道PyTorch Tensor一定要转换成XLATensor，那么我们只需要在关键的转换之处打印出调用堆栈，自然就可以找到调用方。

注册PyTorch库实现：即便我们找到了上面redispatch和codegen的线索，看起来仍然不足以解释PyTorch到torch-xla的桥接，因为PyTorch和torch-xla两个库之间的调用，必须要有符号的映射才可以。

从PyTorch调用到torch_xla：xla调用上面所说的宏进行注册的位置在`RegisterXLA.cpp`这个文件中（codegen的结果），如下：`ORCH_LIBRARY_IMPL(aten, XLA, m) { m.impl(“abs”, TORCH_FN(wrapper__abs)); … }`。

总结：其实torch-xla官方的文档里是有关于代码生成和算子注册这个过程的描述的，只不过一开始我没找到这个文档，走了一点弯路，但是自己探索也会觉得更明了这个过程。

文章来源

原文地址: 点我阅读全文
原文作者: 极市平台