Tensorflow XLA Service 详解 II

本文主要介绍在XLA service阶段针对HloInstruction做的一些显存优化, 对于训练框架来说, 显存优化的工作至关重要, 主要是由于现阶段GPU+CUDA远没有CPU+Linux组合强大, 后者有完善的建立在虚拟内存基础上的内存管理机制, 内存的高效使用由linux kernel来负责, 即便物理内存不足, 还可以使用swap, 内存压缩等技术确保内存的高效供应, 而在GPU+CUDA里, 这方面的工作很大程度让渡给了程序员自己来搞定, GPU程序接触到的就是物理显存, 如果程序的显存申请超过显存容量, 整个程序就会直接coredump, 此外, 显存本身就集成在GPU板卡上, 无法像内存一样扩展, 而GPU本身造价昂贵, 最后, 在深度学习训练中, 大力出奇迹的现状下, 显存的消耗明显超过的摩尔定律, 这也加剧了显存供求关系的矛盾, 正式由于训练框架做了大量的优化, 才能让模型跑起来.

XLA Service的显存优化设计思想与tensorflow整体一样遵循”静态图”的设计: 先整体优化, 再落地实施. 其中, xla/service/buffer_assignment.cc 是整个显存优化的核心, 在1.14版本中, xla/service/支持两种后端: cpu和gpu, 纷纷针对两种backend有进一步的优化算法, 本文主要针对GPU的优化逻辑进行分析

核心文件

内存优化公共:
xla/service/buffer_assignment 内存优化核心文件
xla/service/buffer_liveness.cc  内存片生命周期分析
GPU相关:
xla/service/gpu/buffer_allocations.cc BufferAllocation的组合
xla/service/gpu/gpu_hlo_scheudle.cc   Hlo的处理顺序, 与显存的优化策略息息相关, 简单地说, 按照BFS并行执行的HloInstruction消耗的显存肯定大于所有的HloInstruction都顺序执行. 

核心抽象

HloSchedule

XxxSchedule是TF的代码风格, 类似的有更底层用于Thunk调度的ThunkSchedule, 以及Service提供的HloSchedule. HloSchedule中最重要的就是封装了HloOrdering.

HloSchedule

XLAService内存优化的本质是处理LoigicalBuffer和BufferAllocation之间的关系, 原则是使用尽可能少的BufferAllocation去承载尽可能多的LogicalBuffer, 而如何使用的更少, 就涉及到了对Hlo图的分析, 就涉及到了Ordering的问题, 使用不同策略生成Ordering, 直接影响两个LogicalBuffer之间的约束关系, 最简单的, 在图遍历中, 使用DFS和BFS的2种方式遍历会导致图上节点的内存依赖关系大有不同.

HloOrdering是描述HloInstruction加载序列的基类, 派生类有PredecessorHloOrdering, DependencyHloOrdering 和 SequentialHloOrdering, 其中, DependencyHloOrdering基于依赖关系, 所以可以并行, 性能更高, 但耗更多的内存, 而SequentialHloOrdering完全串行, 性能相对低, 但可以节约更多内存, 而 PredecessorHloOrdering 是个虚类, 需要子类进一步填充predecessors_, 这也是GPU后端使用的方式.不同的Ordering会影响内存的依赖关系, 进一步影响Launch到GPU后Kernel的执行序列. 

GpuHloSchedule

Build()      构造GpuHloSchedule单例, 同时会根据Stream的数量采用BFS或Sequence的LaunchOrder,, 这个LaunchOrder会用于构造GpuHloOrdering, 以及thunk_launch_order_
ThunkLaunchOrder() 拿到上述thunk_launch_order_
ConsumeHloOrdering() 拿到上述GpuHloOrdering

GpuHloOrdering

在使用Multi-stream的情况下, 一方面stream no的分配算法会遍历图给每个HloInstruction分配不同的stream no, 另一方面, GpuHloOrdering也会据此选择BFS算法来生成Ordering存入thunk_launch_order以及HloOrdering.predecessors_ , 后续的显存优化会据此优化显存, 确保Multi-Stream在实际执行的过程中不会产生显存问题.  

BufferAssigner

BufferAssigner 用于构造BufferAssignment对象, 可以理解为BufferAssignment的wrapper

BufferAssignment

内存管理入口类. 所有关于内存优化的内容均在其Run()中实现, 在nvptx_compiler.cc:RunBackend()调用

BufferLiveness

描述一块内存的生命周期, “computes liveness of the output buffers of HLOs and their interference”, 为显存优化的决策提供支持. BufferLiveness 的构造过程经历3层分析, 从内而外, 即从前到后:
LogicalBufferAnalysis: 构造LogicalBuffer, 每个Instruction都可以有多个LogicalBuffer, 通过Accept机制遍历所有的Instruction
TuplePointsToAnalysis: 构造PointToSet, 需要之前的LogccalBuffer,  最终也是通过Accept机制遍历所有的Instruction二者其实的都是Hlo的DFS遍历handler

-2- 规划内存的总入口, 注意入参
-15- liveout的部分只发生在root instruction???

LogicalBuffer

“虚拟内存”, LogicalBuffer 对内存需求的抽象, 内存优化的实质就是调整LogicalBuffer和BufferAllocation的关系, 将来会被HloValue替换掉. 同时, 整个内存优化过程都使用color的概念, 相同的color表示可以融合, 这点在BufferAllocation和BufferSet中都有用到, 显然, 处于一个BufferAllocation之下的LogicalBuffer都有相同的color, 要不也不会在一起, 相应的BufferAllocation实例也使用想用的color进行标记
什么样的两个LogicalBuffer会被分配给同一个BufferAllocation?
1. kWhile, kCall, kConditional 的LogicalBuffer
2. cannot_merge_buffer_sets(): !(两个buffer的id不同(不是一个buffer) && 会发生干涉)

-1- LogicalBuffer只存Shape(确切的是SubShape) id 和inst指针
-3- 所有的LogicalBuffers对象都会丢到, LogicalBufferAnalysis::logical_buffers中
-7- 所以, 一个LogicalBuffer只对应一个Subshape的内存, 不是整个Instruction

BufferAllocation

BufferAllocation  物理内存的抽象, 一个 BufferAllocation 实例最终对应一块连续物理内存, 通常, 一块物理内存会被多个不干涉的LogicalBuffer公用.一个BufferAllocation实例可以分为:

BufferAllocation::Slice

BufferAllocation里的一片, 关联在该BufferAllocation的每一个LogicalBuffer都会有对应的Slice

GetSlice 根据 buffer,  重新构造一个 Slice 对象

ShapeIndex

一个 ShapeIndex 对应的是一个 LogicalBufferList, 真正使用的时候是用的 ShapeIndex 来索引到buffer

-4- index是ShapeIndex, 描述一个Shape的索引, 不是Intuction. buffers 是 BufferList,  此处可以看到, 直接用index来索引得到的BufferList, 

-2- 使用 ShapeIndex 索引 直接->Allocation::Slice->Thunk, 

BufferAllocations

GPU后端对一组BufferAllocation的抽象

-1- 物理内存对象, se::DeviceMemoryBase是StreamExecutor对一块线性内存资源的抽象,  提供了最基础offset+size的计算方法

-2- 每个 buffer 都有一个 BufferAllocation
-8- 真正分配新的buffer

优化过程

-2- 这个backend其实不是llvm的backend 而是HLO的backend?
-6- //buffer_assignment.cc , 都是对元数据的管理, 没有真的分配, 包含了所有的内存管理代码, 显存优化的入口, 这里ConsumerHloOrdering()返回的是hlo_ordering
-8- 就是下文的buffer_size
-11- 对于一个tuple描述的shape, size的计算
-15- 遍历 module//buffer_liveness.cc  –> DFS analysis ???, 构造一个liveness, 需要针对HloModule的TuplePointsToAnalysis::Run()->LogicalBufferAnalysis::Run()等一系列分析, 
-19- 遍历 HloModule, 为每个HloInstruction构造相应的Logicalbuffer, 存入LogicalBufferAnalysis, 并返回其指针, class LogicalBufferAnalysis : public DfsHloVisitorWithDefault, 说明也是用来遍历的handler
-21- 遍历节点, 执行各种handler, 建立logicalbuffer和instruction及其subshape的关系, 并存储在logicalbuffer中
-22- 构造LogicalBuffer对象
-26- 所有非fusion的buffer构造完毕
-27- 这个循环收集所有的fusion_inst
-30- gather 该computation下的所有的fusion instruction
-31- **为什么要放到的最后呢???**
-32- 给fusion的inst分配logical buffer
-44- 构造BufferAssignment
-52- 当前只有这3种op可以colocated
-53- 将这个instruction对应的LogicalBuffer放到BufferSet
-54- 将BufferSet放到BufferSets
-60- colocate意味着最终分配的是同一块物理内存
-62- 使用邻接矩阵的方式表示一张图, 用int64标识一个节点. 遍历每一个node, 首先遍历其每一个neighbor, 这个neighbor不是相邻节点, 而是通过 cannot_merge_buffer_set 设置的, 比如id()不同&&… 他们用过的color就标记为不可用, 取一个neighbor都没用过的color, 最小未用, 给这个node, 最终的效果是, 每一个node和他的neighbors节点的颜色都不同, 但可能和其他节点相同, 相邻的不能共用buffer, 猜测应该是不能共用输入及输出, 前面的输出还是可以作为后面的输入的. 最终, 就是在这里解决内存冲突: 不相邻的标记为相同颜色, 说明后续可能可以合并.  有意思的是, 为什么相邻的不能使用一个buffer呢??? 和我之前的理解正好相反???. MergeColocatedBufferSets 遍历所有的BufferSet, 将能合并的set合并了,  合并之前, 每个BufferSet都和Instruction/node关联, 合并之后, BufferSets中的一个Set里的buffer就不是和instruction强相关了, 可能是几个合在一起的, 不过, 不存在把一个BufferSet拆成几个的情况, 两个bufferset, 要么能合并, 要么不能合并, 不会把一个BufferSet内的LogicalBuffer拆分. 
-70- interference_map表示vector的vector,  表征的是colocated_buffer_sets中两个set的关系,  而从这里看, 每个set里存储的一组logicalbuffer其实都是一个node使用, 这一点其实也可以从colocated_buffer_sets的构造过程可以看出: 3种op每种op都会先添加一组上logical_buffer到colocated_buffer_set, 再把set放到sets中
-72- 将 interference_map中表征的有冲突的node用不同的color表示
-77- num_sets是最大的color值, 即能将nodes划分为互斥最小集合的数量
-80- 将不冲突的, 即color相同的node都放一起
-87- AssignColocatedBufferSets, 代码里的AssignXXX都是构造/分配BufferAllocation的意思, 比如这个 AssignColocatedBufferSets, 就是根据BufferSets分配BufferAllocation, 遍历colocated_buffer_sets 里的每一个Set里的每一个LogicalBuffer, 将一个Set的LogicalBuffer都放在一个BufferAllocation中, 一个Allocation里的LogicalBuffer不一定一样大, 将二者关联的时候会传入buffer_size的, 此时, offset都还是0 
-92- Module->computation->instruction的嵌套关系
-94- 根据assigner构造时传入的函数对象计算,  这里是遍历所有的buffersize, 将相同的合并到一个allocation中
-95- 为 buffer 分配allocation
-96- **每个allocation内的buffer或buffer_size是固定的?**
-99- 将这个buffer 赋给已有BufferAllocation对象, 这里 offset是0, size就是buffer_size, 是第一次存储这两个值
-106- 第一个顺序执行的点
-107- 至此, 就完成了初始的内存分配代码, 接下来 AssignBufferForComputation(), 就是将 Computation都关联到Allocation, 首先是 global_computations, 遍历所有的Instruction->points_to_analysis->GetBuffersDefinedByInstruction(), 将所有的Buffer按照从大到小排序, 之所以从大到小排序, 应该和 BufferAllocation的Size 要比里面的所有 Buffer的size 都大的原因, 这样就能充分利用 allocation, 毕竟越少越好, 前面的用大块 buffer 搞了大的, 后面就可以乘凉, NewAllocation 本身就会将自己链入Assignment, 总之, 这里会通过一系列判断, 将大部分buffer存入尽可能少的Allocation中. global_computations 进行Assign的过程中, 有一些buffer被加入了sequential的范围,  这部分buffer在AssignBuffersWithSequtialOrdering, 同样, 这里Assign也是找Allocation的意思, 这里会用到堆模拟器, 会计算很多offset, size之类的, 将buffer加入到Heap中, NewEmptyAllocation(result.heap_size, 就能看出, 此外, 由于根据hlo_ordering, stream_no> 1时不会所有的都是顺序的, 所以, 会存在根据hlo_ordering最终有多个chunk的情况, 更具体的, 一个computation对应多个LogicalBuffer, 并进行一个simulate, 得到一组chunk, 和一个allocation. 
-112- 有关于顺序执行的代码
-116- 将传入的computation下的每一个instruction的每一个LogicalBuffer都入队
-118- 保存传入的computation下每一个instruction及其位置
-130- 根据注释, :// Output tuple table may be allocated at run-time, so make sure we don’t // overwrite them.
-140- 处理在sorted_buffer中, 但不在 colocated_buffers 中的, 前文的 colocated_buffers 只有3种HloOpCode可以用
-155- 获取每一个已有的Allocation
-161- 说明该buffer已经找到了 allocation
-162- 说明之前没有找到能用的
-163- 找到最接近的大于的key
-172- 从小到大(即从后向前) 找合适的allocation存储buffer
-186- 不是liveout的buffer, 即temp的buffer, 先存着, 没有相应的Allocation, 这里存的是comutation和buffer们的映射关系
-189- 实在不行了, 分配个新的Allocation
-193- 计算一个buffer的 live out 或 interfere 等
-198- **为什么sequence会节约显存???
-199- buffer 顺序执行, 重点关注
-202- 这里用的LazyBestFitHeap, ChoostBestHeapAlgorithm
-208- 默认构造一个buffer只有0, 通过default_colorer(), 还有其他逻辑导致可能没有0吗?-213- out/trace.log:2019-08-02 10:44:04.437956: I tensorflow/compiler/xla/service/buffer_assignment.cc: 1262] Running per-computation heap simulation
-215- 这里的是LogicalBuffer->BufferValue里的color, 不是前文用来合并BufferSet的color, 是通过buffer_liveness.h 中的 DefaultColorer赋值的. 这部分buffers都有color, 但没有被colocate合并过的, 也没有被合并单奥已有allocation中的temp buffer
-216-  single_colored_set里的buffer是可以合并的.
-237- 分析时候可以共享内存
-242- 根据堆模拟器, 重新对齐buffer的offset和size???
-243- 用整个优化过的堆大小来定义新的allocation, 一个allocation里是可以存放很多不同offset和size的buffer的. allocation有自己的size, 全文来看, 同样的color分配一个allocation. 那直接取最大的不就好了, 为啥还有堆模拟器之类的???–>模拟器怎么运作的??? 所以最终, Assignment里存储的tmp buffer是最小数量的Allocation和他的size, 用color区分 + 之前3种已经处理好的Op,  + 其他
-245- 会存入allocation
-256- 把多个allocation里的buffer存到同一个combined_allocation里
-257- 划分的原则
-270- XLA HLO Visitor机制
-271- 使用上述代码构造的buffer来构造Thunk

Leave a Reply

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

%d bloggers like this: