Tensorflow XLA

XLA

XLA(Accelerated Linear Algebra)是TensorFlow计算图的编译器。

官网:

https://www.tensorflow.org/xla?hl=zh-cn

基本架构:

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

Tensorflow XLA_第1张图片

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》

你可能感兴趣的:(tensorflow,深度学习,人工智能)