HLO

HLO（High Level Optimizer），是XLA IR。XLA支持一组正交的基础operators（原子算子），其他的operators都可以由这组基础算子组合而成。HLO IR是分层的嵌套结构，由以下三个层次组成：

HloModule：HLO IR最高层的表示，可以理解成整个程序。一个HloModule可以包含很多个HloComputation。
HloComputation：HLO IR中间层的表示，相当于程序中的一个函数。一个HloModule只能有一个entry_conputation，其他的computation是被entry_computation调用的。我们可以把entry_computation类比作main函数。每个HloComputation可以包含多个HloInstruction。
HloInstruction：HLO IR最底层的表示，相当于程序中的一条指令。每个HloComputation只能有一个root_instruction。root_instruction的output就是该computation的output。computation的input用parameter表示。HloInstruction也可以调用HloComputation，HloInstruction中有一个called_computations_来表示该instruction调用的computation。不同HloInstruction参数和属性可能存在区别，HloInstruction只是一个base class，特殊的Instruction可以由HloInstruction派生得到，例如HloSliceInstruction。

在这里插入图片描述
HloInstruction通过call_computations_调用另一个comoutation，实现IR的嵌套。例如XLA的融合优化，会将被融合的指令替换成一条HloFusionInstruction，被融合被组装到一个HloComputation，并被HloFusionInstruction调用。

 1 HloModule m, is_scheduled=true, entry_computation_layout={(f32[8]{0},f32[8]{0})->f32[8]{0}}
 2                                                                                           
 3 fused_computation {                                                                       
 4   p0 = f32[8]{0} parameter(0)                                                             
 5   p1 = f32[8]{0} parameter(1)                                                             
 6   multiply0 = f32[8]{0} multiply(p0, p1)                                                  
 7   ROOT add0 = f32[8]{0} add(multiply0, p1)                                                
 8 }                                                                                         
 9                                                                                           
10 ENTRY PredFloatMOF {                                                                      
11   p0.1 = f32[8]{0} parameter(0)                                                           
12   p1.1 = f32[8]{0} parameter(1)                                                           
13   ROOT fusion = f32[8]{0} fusion(p0.1, p1.1), kind=kLoop, calls=fused_computation                                                                                                          
14 }

XLA HLO的op操作的是immutable、静态shape和显示广播的tensor。
HLO的入口：

TF、PT和JAX转换成HLO。
经过mlir_bridge，先转成mlir（tf_executor dialect and mhlo），在mlir上做一些优化，再转为HLO。此路线在XLA中还是experimental，还在开发中。

出口：

LHLO、MHLO或LMHLO。XLA的pipeline是将HLO转成MHLO和LMHLO，然后用于codegen。

LHLO

LHLO：“late”-HLO，经过buffer assignment后的HLO。HLO和LHLO的区别在于HLO注重的是tensor的表达，不考虑到内存的分配，比如tensor<8x32x16xfp32>，仅仅表示为tensor，没有具体的内存信息，没有side-effect的。LHLO会为tensor开辟内存空间，也即经过buffer assign，buffer assign相当于传统编译器中的内存分配，LHLO的op具有side-effect。
入口：

出口：

用于codegen。

MHLO

MHLO：“meta”-HLO dialect，是HLO风格的MLIR dialect，并且在IR上扩展支持了dynamic shape。XLA HLO 的shape是静态不可变的，不同shape需要重新编译；MHLO支持动态shape，IR本身有能力表达shape计算和动态shape信息的传递。

MLIR-HLO使得HLO可以从XLA中独立出来，可以结合MLIR构建端到端的编译器。

入口：

TF Graph
XLA HLO

出口：

Linalg IREE、LMHLO，直接用于codegen。使用mlir构建编译器可以选择使用hlo为输入，转换为mlir的MHLO，做相关的分析、变换和优化，再转成LMHLO或者Linalg等做codegen。
XLA HLO

LMHLO

LMHLO：“late”-“meta”-HLO dialect。是LHLO风格的MLIR dialect。

入口：

MHLO或者XLA HLO经过shedule 和 buffer assign 后再转换的到。

出口：

用于codegen到llvm ir。

XLA IR 总结

(M)HLO的op是tensor类型，是不可变的 (immutable)、并且不具有 side effect，tensor的数据流分析（例如ssa def-use chain ）和转换会比较容易。而L(M)HLO是经过buffer assign的，buffer 是可变的 (mutable)和有别名的 (alias)，buffer上分析和转换需要比较复杂的依赖分析 (dependency analysis) 和别名分析 (alias analysis)。XLA会在HLO上做优化相对比较容易，包括传统的图优化（代数化简、死代码删除等）、融合相关优化等等。在完成HLO层的优化后转为LMHLO，利用mlir LMHO dialect做codegen生成llvm ir。
XLA pipeline中是直接将XLA HLO转换为LMHLO，然后在LMHLO上codegen。
按照xla的技术路线，lmhlo将合并到mhlo中，将不再区分mhlo和lmhlo，mhlo可以同时操作tensor和memref。

HLO->LMHLO

主要的处理过程如下：
1、buffer assignment。

hlo instruction schedule，确定buffer liveness(减少内存使用)等。
buffer assignment。

2、HLO转换为LMHLO。

创建mlir module、funcOp及其参数。
以SequentialHloOrdering顺序（可以先简单理解成指令执行顺序，这里会考虑到节省内存）处理entry computation中的每条指令，不同的指令类型进入不同的处理函数。例如：为hlo的instruction创建lmhlo的op；有call_computation时，以root节点后序遍历（post order）顺序处理computation中每条指令。

参考：