Tensorflow XLA
XLA
XLA(Accelerated Linear Algebra)是TensorFlow计算图的编译器。
官网：
https://www.tensorflow.org/xla?hl=zh-cn
基本架构：
https://www.tensorflow.org/xla/architecture
CSE(Common subexpression elimination)
DCE(Dead code elimination)
TFE(Tensorflow Eager)
应用层
TF目前（v2.4.1）的默认编译选项中已经包含了XLA，但是默认不会启动。
启动方法：
- 手动启动
设置环境变量：
TF_XLA_FLAGS="--tf_xla_enable_xla_devices"
手动指定需要XLA计算的op：
with tf.device("/device:XLA_CPU:0"):
...
这种方法的缺点是：代码需要修改，且XLA不支持有些复杂的op。
- 自动启动
TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit"
代码无需修改。
unit test：
tensorflow/compiler/xla/tests
bazel build //tensorflow/compiler/xla/tests:convolution_test
./bazel-bin/tensorflow/compiler/xla/tests/convolution_test_cpu --gtest_filter="XXXX"
HLO
XLA用HLO(High Level Optimizer)这种中间表示形式，表示正在被优化的计算图。
三个概念，hlo module, computation, instruction。
- hlo module用源码注释的解释，就是一个编译单元，相当于是一个完整可运行的程序。既然是一个程序，就有入口函数，也就是entry_computation，每个module都有且仅有一个entry_computation，相当于main函数，有输入和输出，输入可以是多个参数，但输出只有一个（root instruction的值），如果要返回多个值，需要把多个值构造成一个元组（tuple）返回。
- 一个module可以包含多个computation，除了entry_computation，其他的都是"nested"，也就是被调用。
- HLO instructions就是op了，对应了官网上列出的operation semantics，看注释已经解释的非常清楚了，op融合和向llvm ir转换都是在这个层面进行的。
op的官方定义：
https://tensorflow.google.cn/xla/operation_semantics
HLO也可以有pass：HloModulePass
HloPassPipeline pipeline("pass");
pipeline.AddPass<XXXPass>();
参考：
https://zhuanlan.zhihu.com/p/71980945
tensorflow xla hlo基本概念和pass pipeline
从TF到HLO
一个op从上到下一般有这样几个步骤：
TF op -> XLA op -> HLO op
- TF op -> XLA op:
tensorflow/python/ops/nn_ops.py
conv2d
gen_nn_ops.conv2d
_op_def_library._apply_op_helper("Conv2D")
- XLA op -> HLO op:
tensorflow/compiler/tf2xla/kernels/conv_ops.cc
REGISTER_XLA_OP(Name("Conv2D"), Conv2DOp);
ConvOp::Compile
MakeXlaForwardConvOp
HandleConcatenate
底层实现
XLA支持两种接入模式：
-
JIT。compiler/jit/
-
AOT。compiler/aot/
将整个计算图用_XlaCompile & _XlaRun替换。
compiler/tf2xla/ -> compiler/xla/client/ -> compiler/xla/service/
最终的计算由service负责实现。
compiler/aot/
以AOT的方式将tf2xla/接入TF引擎。
compiler/jit/
以JIT的方式将tf2xla/接入TF引擎， 核心是9个优化器和3个tfop，其中XlaCompileOp调用tf2xla的“编译”入口完成功能封装，XlaRunOp调用xla/client完成“运行”功能。
compiler/tf2xla/
对上提供xla_compiler.cc:XlaCompiler::CompileFunction()供jit:compile_fn()使用将cluster转化为XlaComputation。核心是利用xla/client提供的接口，实现对XlaOpKernel的“Symbolic Execution”功能。每个XlaOpKernel子类均做的以下工作: 从XlaOpKernelContext中取出XlaExpression或XlaOp，调用xla/client/xla_buidler.h提供的方法完成计算，将计算结果的XlaOp存入XlaKernelContext。
compiler/xla/client/
对上提供xla_builder.cc:Builder等供CompileFunction()使用，将Graph由Op表达转化为HloModuleProto:HloComputationProto:HloInstructionProto表达并保存在XlaComputation中。
对上提供local_client.cc:LocalClient::Compile()，作为编译入口供jit：BuildExecutable()使用，将已经得到的XlaComputation交给service并进一步编译为二进制。
对上提供local_client.cc:LocalExecutable::Run()，作为运行入口供jit/kernels/xla_ops.cc:XlaRunOp使用，通过Key找到相应的二进制交给service层处理。
compiler/xla/service/
对上提供local_service.cc:LocalService::BuildExecutable()供LocalClient::Compile()使用实现真正的编译，承接XlaComputation封装的HloProto，将其转化为HloModule:HloComputation:HloInstruction表达，对其进行优化之后，使用LLVM后端将其编译为相应Executable后端的二进制代码。
对上提供executable.cc:Executable::ExecuteOnStream()供LocalExecutable::Run()使用实现真正的执行二进制。
backward
tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md
CanonicalizeBackwardFilterConvolution
Conv2DBackpropFilter
GetKnownXLAWhitelistOp
XlaOpRegistry::GetAllRegisteredOps
REGISTER_XLA_OP
HloInstruction::Visit
class ConvBackpropFilterOp : public XlaOpKernel
MakeXlaBackpropFilterConvOp
ConvGeneralDilated
https://discuss.tvm.apache.org/t/rfc-mlir-frontend/6473
backend
官方backend：
tensorflow/compiler/xla/service
第三方的XLA backend接入：
tensorflow/compiler/plugin
第三方的XLA backend中，比较出名的是graphcore。
它的TF实现：
https://github.com/graphcore/tensorflow/tensorflow/compiler/plugin/poplar
XLA的主要目的是方便硬件厂商更好的适配tensorflow。因此，作为XLA基础的HLO，其op数非常少，仅有不到100个。用户只要实现了这些op，就可以接入tf了——其他不支持的tf op，都被分解为简单的HLO op。
HloTransposeInstruction
HandleTranspose
HLO op的弊端是颗粒度太细，导致执行效率不高。因此，XLA还提供了高级op的注册功能，主要是用xla::CustomCall来实现。
MaxPool2DGradOp
REGISTER_XLA_OP(Name("MaxPoolGrad").Device(DEVICE_IPU_XLA_JIT),
MaxPool2DGradOp);
xla::CustomCall(&b, PoplarOp_Name(PoplarOp::MaxPoolGrad), args,
input_shape, attribute_map_.Serialise());
model test:
tensorflow/compiler/plugin/poplar/docs/example_tf2_model_fit.py
混合backend
XLA支持混合多backend的运行，可用tf.debugging.set_log_device_placement(True)查看相关的设备指派信息。
设备指派主要由Placer模块负责：
https://www.cnblogs.com/deep-learning-stacks/p/9823486.html
TensorFlow中的Placement启发式算法模块——Placer
backend优先级
REGISTER_LOCAL_DEVICE_FACTORY(DEVICE_XLA_XXX_NPU, XlaXXXNpuDeviceFactory, 500);
CPU的优先级是50，添加的backend的优先级只要大于50，就可以得到调度权。
unit test
写好的backend需要测试，同时Unit Test也是编写一个backend的入门级入口。
bazel build -c dbg tensorflow/compiler/xla/tests:convolution_variants_test
./bazel-bin/tensorflow/compiler/xla/tests/convolution_variants_test_cpu --gtest_filter=*BackwardInputLowPaddingLessThanHighPadding*
参考
https://mp.weixin.qq.com/s/RO3FrPxhK2GEoDCGE9DXrw
利用XLA将GPU性能推向极限
https://mp.weixin.qq.com/s/MPI9KERDS-Al4DTBDRV04w
TensorFlow XLA工作原理简介
https://sketch2sky.com/
一个XLA方面的blog
https://tensorflow.juejin.im/performance/xla/jit.html
使用即时编译
https://blog.slinuxer.com/2019/06/tensorflow-xla
TensorFlow XLA初步接触
https://github.com/horance-liu/tensorflow-internals
电子书《TensorFlow Internals》