Tensorflow XLA Service 详解 I

compiler/aot/ 以AOT的方式将tf2xla/接入TF引擎
compiler/jit/以JIT的方式将tf2xla/接入TF引擎, 核心是7个优化器和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()使用实现真正的执行二进制。

编译cubin

调用栈:

-2- Client端Graph编译入口
-3- LocalService端Graph编译入口
-7- Service的Graph编译入口
-8- 根据Client端生成的HloProto表示的Graph转换为Hlo格式
-10- 优化HloModule, backend是XlaOp过滤逻辑用到的, 这里是service/gpu/nvptx_compiler.cc-11- 编译HloModule入口   nvptx_compiler.cc
-12- 构造最终提交到LLVM的llvm::Module对象
-13- 按照PostOrder的顺序依次给HloInstruction分配stream number, 根据是否是GEMM, 决定是否复用operand的stream number. 
-14- 分配Stream number, 决定了HloInstruction最终的处理顺序, 核心工作时确定了thunk_launch_order_以及据此构造的hlo_ordering_. 如果是配置为单Stream, 就是PostOrder, 否则使用BFSLaunchOrder, 会根据HloInstruction之间的依赖关系, 以及一共可用的Stream number数量来给每个HloInstruction分配Stream Number, 原则上, 会将不存在依赖关系的HloInstruction尽量分配到不同的Stream number, 如果有依赖, 那么这个HloInstruction的Stream number会和某个它所依赖的HloInstruction使用相同的Stream number, 所谓的launch order 并不是最终执行的顺序
-15- **显存优化**
-16- visitor 也可以是别的visitor, 之前优化HloInstruction的时候就用到很多
-17- 遍历每一个HloInstruction 构造相应的Thunk, 除了全图的Accept, Instruction也有自己的Accept()用于局部遍历
-19- nvptx_backend_lib.cc  根据HloInstruction 生成ptx
-33- 根据ptx生成cubin, 即把string 的ptx变成uint8-39- 这个有没有优化空间
-42- 这个是性能分析???

执行cubin

调用栈:

-23- 将kernel加入kernel_cache, cache的作用是防止load 的时间占用了执行的时间, 让execute的统计更准确
-34- 获取custream
-39- 加载kernel到GPU执行, 注意, 加载时机和实际执行时机不是一回事, CPU端只需批量加载, GPU负责顺序执行stream上两个kernel

Leave a Reply

This site uses Akismet to reduce spam. Learn how your comment data is processed.

%d bloggers like this: