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中每条指令。

参考：

https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compile r/xla/mlir_hlo
https://mlir.llvm.org/docs/LangRef/#high-level-structure

XLA IR：HLO、LHLO、MHLO和LMHLO相关推荐

Tensorflow XLA
XLA XLA(Accelerated Linear Algebra)是TensorFlow计算图的编译器. 官网: https://www.tensorflow.org/xla?hl=zh-cn 基 ...
TensorFlow XLA优化原理与示例
TensorFlow XLA优化原理与示例 XLA概述 XLA(加速线性代数)是用于优化TensorFlow计算的线性代数的域特定编译器.结果是在服务器和移动平台上的速度,内存使用率和可移植性得到了改 ...
TensorFlow XLA 初探
在此分享一下个人在学习XLA时的一点理解和心得~ XLA(加速线性代数)是一种针对特定领域的线性代数编译器,能够加快 TensorFlow 模型的运行速度,而且可能完全不需要更改源代码.它可以提高运行 ...
阿里 BladeDISC 深度学习编译器正式开源
简介:随着深度学习的不断发展,AI模型结构在快速演化,底层计算硬件技术更是层出不穷,对于广大开发者来说不仅要考虑如何在复杂多变的场景下有效的将算力发挥出来,还要应对计算框架的持续迭代.深度编译器就成了 ...
AICompiler动态shape编译框架
AICompiler动态shape编译框架移动互联网的兴起,不仅产生了海量数据,也对人机交互有了新的定义.企业如何动态处理不同规格图片数据,如何更灵活处理不同长度的对话语料等等,提升企业运营效率,争 ...
使用nGraph的Intel®Xeon®上的高性能TensorFlow
使用nGraph的Intel®Xeon®上的高性能TensorFlow High-performance TensorFlow* on Intel® Xeon® Using nGraph 最近宣布了n ...
AICompiler动态shape编译框架案例和效果数据
简介:欢迎走进走进阿里云机器学习PAI AICompiler编译器系列.近期,阿里云机器学习PAI团队全新上线一套Dynamic Shape Compiler框架,不仅作为AICompiler技术栈中 ...
AI编译器与传统编译器的联系与区别
AI编译器与传统编译器的区别与联系总结整理自知乎问题针对神经网络的编译器和传统编译器的区别和联系是什么?. 文中提到的答主的知乎主页:@金雪锋.@杨军.@蓝色.@SunnyCase.@贝壳与知了. ...
first meet ot MLIR
first meet ot MLIR definition what question does it solve? XLA HLO tensorflow xla hlo 基本概念: 初见MLIR d ...

XLA IR：HLO、LHLO、MHLO和LMHLO