开篇：XLA是什么？

XLA(Accelerated Linear Algebra)是专用于机器学习的编译器，机器学习的运算中99%都是向量乘以矩阵、矩阵乘以矩阵的计算，XLA是专门用来优化这些计算的。

How to

举个例子，运行在GPU上的model_fn函数会顺序调用multiply、add和reduce_sum这三个op，而且multiply，也就是y * z的计算结果会先从GPU拷贝回host，再拷贝到device作为add的input，同样的，add的计算结果也会以相同的方式传递给下一个op。

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

显然，对于整个函数来说，将中间变量在host和device间来回倒腾是没有意义的。因此，如果把函数看作一个op，那在计算中产生的中间结果就不必返回到host，少了数据传输的时间开销，就可以大幅提升运算效率。

这种将多个op融合成一个op的方法就称为fuse，当前fuse的技术路线有：

• 通过手写或codegen工具来开发fused op，例如在上述例子中就可以开发tf.fused_reduce_sum(x, y, z)。它的优点是代码可控性高，易于性能优化，但缺点是程序缺乏灵活性。像Pytorch这种动态图的框架走的就是这条路线，Nvidia的Apex提供有大量fused kernel，对fused kernel感兴趣的，可以读读LayerNorm核心技术。
• 通过XLA等AI编译器将python函数编译成fused op。这样做的好处是灵活性强，可以fuse任何计算，弊端则是开发难度大，且性能通常会逊色于手写或codegen kernel。

性能

XLA的优化当然不只是fuse，还有对计算图的优化，包括删除无效指令、减少内存占用、替换复杂指令等优化。下图是官方提供的性能报告，经XLA优化过后，Tensorflow BERT MLPerf的训练性能提升了~7倍。除了Tensorflow外，XLA还支持JAX、Julia、PyTorch和Nx等前端。

https://www.tensorflow.org/xla

Just in time（JIT）

jit是指在首次运行时将函数编译成二进制程序，后续再调用该函数时直接运行先前编译好的程序而非python code。@tf.funciton修饰的函数（包括它的子函数）会做jit。除非signature发生了变化，也就是input的shape或dtype和编译时不同，否则get_MSE是不需要重复编译的。

@tf.function
def get_MSE(y_true, y_pred):
  print("compiling ...")
  sq_diff = tf.pow(y_true - y_pred, 2)
  return tf.reduce_mean(sq_diff)

get_MSE(tf.constant(1.0), tf.constant(2.0)) # compile
get_MSE(tf.constant(3.0), tf.constant(4.0)) # It won't recompile
get_MSE(tf.ones([2, 2]), tf.ones([2, 2]) # compile again for new signature

@tf.function将函数内的ops替换成一组（XlaCompile, XlaRun) ops，在运行时前者负责编译，并将编译结果--executable保存到cache，后者负责运行executable。如果cache里已经有编译好的程序就不需要编译了，例如get_MSE(tf.constant(3.0), tf.constant(4.0))。

HLO

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。

https://www.tensorflow.org/xla/architecture

Executable

除了二进制程序，它还包含运行该程序所需要的infos和options。调用executable.run()就可以执行计算图。

END