Jax工作原理

一、Jax

    可以将Jax分为由纯Python组成的jax和主要由C++组成的jaxlib，jaxlib包含XLA、XLA使用的LLVM组建片段、MLIR基础设施（StableHLO Python绑定）和专门用于快速JIT和PyTree操作的JAX特定C++库。

1.1、Jax

    最左侧是开发者自己编写的Python代码，JAX会追踪并变换成JAX IR中间 表示，并按照Python代码，通过jax.jit将其编译成HLO(High Level Optimized)代码，代表高级的优化代码，提供给XLA进行读取。XLA在获取编译的HLO代码后，会分配到对应的CPU、GPU、TPU或者ASIC。

1.2、jaxlib

    jaxlib被分为两个主要的存储库，即jaxlib主Jax存储库的子目录和XLA源代码。XLA内部的JAX特定部分主要位于xla/python子目录（JAX的C++组件位于XLA树内部的原因部分是历史原因，部分是技术原因，技术原因在于XLA C++ API不稳定，通过将XLA:Python绑定保留在XLA树中，可以将他们的C++实现与XLA的C++ API进行原子更新。）。

二、XLA

    XLA(Accelerated Linear Algebra, 加速现行代数)，是Google推出的高性能机器学习领域编译器，它可以在不更改源代码的条件下加速Tensorflow模型。机器学习的运算中99%都是向量乘以矩阵、矩阵乘以矩阵的计算，XLA是专门用来优化这些计算。

2.1、存在的问题

    一个函数调用中有多个计算操作，每个计算操作都需要将结构从GPU拷贝到CPU，这种来回的数据拷贝是没有意义的。因此，如果把函数看做一个操作，那么在计算中产生的中间结果不必返回到host，少了数据传输的时间开销，就可以大幅提升运算效率。
    这种将多个操作融合成一个操作的方法称为fuse。

2.2、技术线路

    当前fuse的技术线路有：

• 通过手写或codegen工具来开发fused操作。
• 通过XLA等AI编译器将python函数编译成fused操作。

2.3、AI编译器

    传统编译器分为三部分：前端(FrontEnd)、优化器(Optimizer)和后端(BackEnd)。

• 前端：负责词法分析和语法分析，将源代码转化为抽象语法树
• 优化器：在前端的基础上，对得到的中间代码进行优化，是代码更加高效
• 后端： 将已经优化的中间代码转化为针对格子平台的机器代码

2.4、XLA原理

    现代编译器核心离不开IR(Intermediate Representation，中间表示)，编译过程是将高级语言不断变化最终编程机器可以运行地机器指令，整个过程是一个复杂的过程。

• XLA Graph：节点和边缘组成，其中节点是现行代数操作，边缘是数据流，比如tensors
• XLA Graph从线性代数降低到编译的形式，到达LLVM IR，这种中间表示形态
• 通过Code Generator生成实际的机器平台代码

2.5、XLA编译

    XLA基于LLVM框架开发，前端的输入是Graph，前端没有将Graph直接转化为LLVM IR。首先XLA的功能主要体现在两个方面：

• 即使编译(Just In Time)：边运行边编译
• 超前编译(Aheda Of Time)：运行前编译

    这两种编译方法的主要区别在于是否在运行时进行编译，对于AI训练模型中，AOT模式下更具有性能优势。

    此外，XLA还有着大部分深度学习编译器都有的功能：摆脱计算库的限制，自动生成算子代码并支持在多硬件上的良好移植性。

2.6、XLA工作原理

    XLA作用于计算图如下：

    这里假设XLA仅支持matmul和add，XLA通过图优化方法，在计算图中找到合适被JIT编译的区域：

XLA把这个区域定义为Cluster，作为一个独立的JIT编译单元，计算图中通过Node Attribute标示：

然后另一个的图优化方法，把Cluster转化成Function子图，在原图上用一个Caller节点表示这个Function在原图的位置:

最后调用图优化方法，把Function节点转化成特殊的XLA节点。

参考

AI编译器XLA调研
XLA优化原理简介
机器学习洞察 | 一文带你“讲透” JAX