第1个回答 2022-07-30
XLA (Accelerated Linear Algebra)是专用于机器学习的编译器，机器学习的运算中99%都是向量乘以矩阵、矩阵乘以矩阵的计算，XLA是专门用来优化这些计算的。
举个例子，运行在GPU上的 model_fn 函数会顺序调用 multiply 、 add 和 reduce_sum 这三个op，而且 multiply ，也就是 y * z 的计算结果会先从GPU拷贝回host，再拷贝到device作为 add 的input，同样的，add的计算结果也会以相同的方式传递给下一个op。
显然，对于整个函数来说，将中间变量在host和device间来回倒腾是没有意义的。因此，如果把函数看作一个op，那在计算中产生的中间结果就不必返回到host，少了数据传输的时间开销，就可以大幅提升运算效率。
这种将多个op融合成一个op的方法就称为 fuse ，当前fuse的技术路线有：
XLA的优化当然不只是fuse，还有对计算图的优化，包括删除无效指令、减少内存占用、替换复杂指令等优化。下图是官方提供的性能报告，经XLA优化过后，Tensorflow BERT MLPerf的训练性能提升了~7倍。除了Tensorflow外，XLA还支持 JAX 、 Julia 、 PyTorch 和 Nx 等前端。
jit 是指在首次运行时将函数编译成二进制程序，后续再调用该函数时直接运行先前编译好的程序而非python code。 @tf.funciton 修饰的函数（包括它的子函数）会做 jit 。除非signature发生了变化，也就是input的shape或dtype和编译时不同，否则 get_MSE 是不需要重复编译的。
@tf.function 将函数内的ops替换成一组（ XlaCompile , XlaRun ) ops，在运行时前者负责编译，并将编译结果-- executable 保存到cache，后者负责运行executable。如果cache里已经有编译好的程序就不需要编译了，例如 get_MSE(tf.constant(3.0), tf.constant(4.0)) 。
XLA编译器支持的语言（IR）是HLO（High Level Operations），顾名思义这些语言是由一个个op组成，因此，我们在编译前需要先从python code中提取出所有ops，再将它们转换成HLO。
JAX通过tracing的方式，从 @jax.jit 修饰的函数中提取ops，这些ops通过 jaxpr 来表示。然后再通过XLA client提供的API为ops生成相应的HLO。PyTorch/XLA也是采用类似的方法来生成HLO。
Tensorflow的 tf2xla 为每个 Op 创建了一个同名的 XlaOp 用于生成HLO， XlaOp 派生于 Op ，使用相同的注册机制，因此，只要把要编译的子图根据拓扑排序运行一遍就能生成它的HLO。
HLO先经过一系列 pass 优化后再将HLO lowering成ISA，最后将编译好的二进制封装到 executable 。
除了二进制程序，它还包含运行该程序所需要的infos和options。调用 executable.run() 就可以执行计算图。