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》
Tensorflow XLA相关推荐
- TensorFlow XLA优化与Memory
TensorFlow XLA优化与Memory XLA概述 XLA(加速线性代数)是用于优化TensorFlow计算的线性代数的域特定编译器.结果是在服务器与移动平台上的速度,内存使用率与可移植性得到 ...
- TensorFlow XLA加速编译器
TensorFlow XLA加速编译器 加速线性代数器(Accelerated linear algebra,XLA)是线性代数领域的专用编译器.根据 https://www.tensorflow.o ...
- TensorFlow XLA优化原理与示例
TensorFlow XLA优化原理与示例 XLA概述 XLA(加速线性代数)是用于优化TensorFlow计算的线性代数的域特定编译器.结果是在服务器和移动平台上的速度,内存使用率和可移植性得到了改 ...
- TensorFlow XLA 初探
在此分享一下个人在学习XLA时的一点理解和心得~ XLA(加速线性代数)是一种针对特定领域的线性代数编译器,能够加快 TensorFlow 模型的运行速度,而且可能完全不需要更改源代码.它可以提高运行 ...
- 使用TensorFlow XLA辅助实现BERT预训练加速
XLA 简介 XLA 是 TensorFlow 图表的编译器,只需更改极少的源代码,便可加速您的 TensorFlow ML 模型.这篇文章将介绍 XLA,并说明如何在您自己的代码中试用 XLA. 在 ...
- tensorflow XLA 笔记
笔记 通过session_config 开启XLA 特定op关闭 XLA 通过session_config 开启XLA from tensorflow.core.protobuf import con ...
- 深入浅出 python epub_《机器学习从认知到实践(第2辑)(套装共3册,Python+TensorFlow)》epub+mobi+azw3...
套装书目 · · · · · · <Python机器学习实践指南> <TensorFlow机器学习项目实战> <TensorFlow技术解析与实战> 分册简介 ...
- TensorFlow学习笔记——《TensorFlow技术解析与实战》
著名历史学家斯塔夫里阿诺斯在<全球通史>中,曾以15世纪的航海在"物理上"连通"各大洲"作为标志将人类历史划分为两个阶段.在我正在写作的<互联 ...
- Tensorflow教程
前言 1.TensorFlow 是由 Google Brain 团队为深度神经网络(DNN)开发的功能强大的开源软件库: 2.TensorFlow 允许将深度神经网络的计算部署到任意数量的 CPU 或 ...
最新文章
- PyTorch学习之六个学习率调整策略
- 10.无重复字符的最长子串---使用滑动窗口方法和哈希表来解决
- java 内存测试_请你说一下java jvm的内存机制
- Windows Server 2008 R2使用LDP恢复已删除的用户实战
- WinCE开机Logo的实现(USB下载图片到nandflash)
- linux 下设置定时任务
- Vue如何循环提取对象数组中的值
- 过拟合解决方法之L2正则化和Dropout
- PTA 7-5 jmu-java-m05-自定义Judgeable接口 (10 分) JAVA
- xcel Home 数据透视表初级班(10118班)小结第二课时
- MongoDB文档增删改查
- 蛮好用的Gungho重点工作督查督办跟踪管理系统
- springboot集成kettle, 访问不到ktr文件
- Windows 系统维护
- 计算机组成与系统结构(课程设计)
- 关于计算机编程的收获的作文,编程的乐趣作文800字
- CCF-GAIR 2020 源创专场 | 免费领取学术票
- 目前最新版AndroidStudio详细安装教程,纯萌新安装
- python小论文范文3000字_完整的论文范文3000字
- 日志管理系统排名_目前较好的日志管理系统有哪些?