XLA - 加速线性代数

XLA : Accelerated Linear Algebra

XLA（加速线性代数）是一款用于GPU、CPU和ML加速器的开源机器学习（ML）编译器。

XLA 编译器采用流行的 ML框架（如PyTorch、TensorFlow和JAX）中的模型，并对其进行优化，以在包括 GPU、CPU和ML加速器在内的不同硬件平台上实现高性能执行。

如今，XLA 支持多种机器学习框架前端（包括 PyTorch、TensorFlow 和 JAX），并且是 OpenXLA 项目的一部分。

OpenXLA 项目是一个用于机器学习的开源编译器技术生态系统，由领先的机器学习硬件和软件组织协作开发。

在创建 OpenXLA 项目之前，XLA 是在 TensorFlow 项目中开发的，但基本目标保持不变：

提高执行速度。 编译子图以缩短短期有效的操作的执行时间并消除运行时的开销，融合流水线操作以减少内存开销，并专门化已知张量形状以实现更积极的常量传播。
提高内存用量。 分析和调度内存用量，消除许多中间存储缓冲区。
减少对自定义操作的依赖。 通过提升自动融合的低级别操作的性能，使其匹配最初手动融合的自定义操作的性能，从而不再需要许多自定义操作。
提高可移植性 。为新型硬件编写新后端相对简单，以便大部分机器学习模型无需修改即可在该硬件上运行。这与为新硬件专门处理单个单体式运算的方法相反，后者需要重写模型才能利用这些运算。

在这里插入图片描述

XLA 编译器从 StableHLO 中定义的机器学习框架获取模型图，并将它们编译为适用于各种架构的机器指令。

StableHLO 定义了一个版本化操作集（HLO 表示高级操作），可在机器学习框架和编译器之间提供可移植层。

通常，将模型图转换为针对目标优化的可执行文件的编译过程包括以下步骤：

XLA 对独立于目标的 StableHLO 图执行多次内置优化和分析传递，例如用于为计算分配运行时内存的 CSE、与目标无关的操作融合和缓冲区分析。在此优化阶段，XLA 还会将 StableHLO 方言转换为内部 HLO 方言。
XLA 将 HLO 计算发送到后端，以进一步优化 HLO 级别的优化，这次同时考虑到特定于目标的信息和需求。例如，GPU 后端可以执行对 GPU 编程模型特别有益的操作融合，并确定如何将计算划分为数据流。在此阶段，后端还可以将某些操作或其组合与优化的库调用进行模式匹配。
然后，后端执行特定于目标的代码生成。XLA 中包含的 CPU 和 GPU 后端使用 LLVM 进行低层级 IR、优化和代码生成。这些后端发出以高效方式表示 HLO 计算所需的 LLVM IR，然后调用 LLVM 以从此 LLVM IR 中发出原生代码。

在此过程中，XA 编译器是模块化的，因为它可以轻松插入替代后端，以以某些新颖的硬件架构为目标。GPU 后端目前通过 LLVM NVPTX 后端支持 NVIDIA GPU。CPU 后端支持多个 CPU ISA。