返回介绍

为 XLA 开发一个新后端

发布于 2025-01-22 23:08:20 字数 3058 浏览 0 评论 0 收藏 0

本初级指南针对的是早期用户,他们希望用一种高效的方式轻易地将 TensorFlow 重定向到他们自己的硬件上。本指南不会手把手地讲解,我们假设读者已经了解 LLVMBazel 、以及 TensorFlow。

XLA 提供了一个抽象的接口,让新的架构或加速器可以实现并创建可运行 TensorFlow 计算图的后端。
在新的硬件上重定向 XLA 要比重新实现 TensorFlow 所有的操作更简单、更具扩展性。

大部分实现都属于下列情形之一:

  1. 已有的 CPU 架构尚没有在官方 XLA 中得到支持,存在或不存在已有的 LLVM 后端。
  2. 非 CPU 硬件,已有 LLVM 后端。
  3. 非 CPU 硬件,没有 LLVM 后端。

注意:LLVM 后端既包括官方发布的 LLVM 后端,也包括内部开发的定制 LLVM 后端。

场景 1:已有的 CPU 架构,但官方 XLA 尚不支持

在此场景中,先查看一下已有的 XLA CPU 后端 。通过 LLVM,XLA 让 TensorFlow 更容易重定向到不同的 CPU,因为对于 CPU 来说, XLA 后端之间的主要区别在于 LLVM 生成的代码。Google 在 x64 和 ARM64 架构下测试了 XLA。

如果硬件厂商已经为他们的硬件开发了 LLVM 后端,将此后端与用 XLA 构建的 LLVM 链接起来是比较容易的。
在 JIT 模式中,XLA CPU 后端为主机 CPU 生成代码。对于提前(ahead-of-time)编译,
xla::AotCompilationOptions 可提供一个 LLVM 三元组,来配置目标架构。

如果尚不存在 LLVM 后端,但是存在另一种代码生成器,重用已有 CPU 后端的代码也是可能的。

场景 2:非 CPU 硬件,已有 LLVM 后端

在已有的 xla::CPUCompiler
xla::GPUCompiler 类的基础上,是有可能建模出一个新的
xla::Compiler 实现的。因为它们已经生成了 LLVM IR。
根据硬件的不同,有可能许多 LLVM IR 生成的过程也会不同,但是很多代码是可以和已有的后端共享的。

参考的一个很好的例子是 XLA 的 GPU 后端
这个 GPU 后端针对的非 CPU 的 ISA,因而它的代码生成的很多方面是 GPU 领域专有的。其它类型的硬件,比如,DSP 一类的 Hexagon (它们是有上游的 LLVM 后端的),可以重用 LLVM IR 生成逻辑的部分代码,但是其它部分则是不一样的。

场景 3:非 CPU 硬件,没有 LLVM 后端

如果没有可能利用 LLVM,则最好的选择是为你的硬件及 XLA 实现一个全新的后端。
这是最难的选择,因为你要实现如下的类:

  • StreamExecutor :
    对于许多设备而言,不是 StreamExecutor 中的所有方法都是必要的。更多细节,参见已有 StreamExecutor 的实现。
  • xla::Compiler :
    这个类封装了一个 HLO 计算到一个 xla::Executable 的编译。
  • xla::Executable :
    这个类用于在目标平台上启动一个编译后的计算过程。
  • xla::TransferManager :
    这个类让后端提供针对目标平台的机制,用于从设备内存句柄构造出 XLA 字面量数据。换句话说,它帮助封装了主 CPU 与设备之间的数据传输。

如果您发现本页面存在错误或可以改进,请 点击此处 帮助我们改进。

如果你对这篇内容有疑问,欢迎到本站社区发帖提问 参与讨论,获取更多帮助,或者扫码二维码加入 Web 技术交流群。

扫码二维码加入Web技术交流群

发布评论

需要 登录 才能够评论, 你可以免费 注册 一个本站的账号。
列表为空,暂无数据
    我们使用 Cookies 和其他技术来定制您的体验包括您的登录状态等。通过阅读我们的 隐私政策 了解更多相关信息。 单击 接受 或继续使用网站,即表示您同意使用 Cookies 和您的相关数据。
    原文