Tensorflow XLA Service Buffer优化详解

下图是Tensorflow 架构图,以及XLA在Tensorflow中的位置

通过层层转换(参考Tensorflow XLA Service 详解 I),Graph在进入XLA Service前已经被表达为HloModule的形式,而作为图编译器的核心,XLA Service就负责将HloModule表达的计算图编译为可以直接在不同硬件平台(Backend)执行的程序,而编译的核心,就是优化代码,包括设备无关的优化和设备相关的优化:

  1. 优化HloModule所表示的计算图,并将其转化为LLVM HLO
  2. 基于LLVM,生成硬件相关的二进制

作为通用的编译器框架,LLVM 会对LLVM HLO做大量的优化,生成高效的Target binary, 所以作为XLA的开发者,主要关注阶段1的优化: 设备无关的图优化算法。Compiler是XLA适配硬件的接口类,每个适配XLA的硬件都必须实现其中的方法。尤其是RunBackend(),参考Tensorflow XLA Service 详解 I 一文,该接口是进行图优化和编译的入口,统领整个优化和编译过程。

同时,XLA Service还实现了一组通用的优化方法(各种Schedule策略,各种Memory优化算法)供各个硬件平台的编译器使用,当然,主要是供给RunBackend()调用。以XLA的GPU平台的编译优化流程为例:

NVPTXCompiler::RunBackend()
  hlo_schedule = GpuHloSchedule::Build(*module, *stream_assignment, pointer_size_)
  BufferAssigner::Run(hlo_schedule->ConsumeHloOrdering()...)
  entry_computation->Accept(&ir_emitter)
  CompileToPtx() 
  CompilePtxOrGetCachedResult()

-1- 从XLA Service通用层中选择适合GPU的Schedule策略
-3- 基于Schedule策略,进行设备无关的Buffer优化,主要关注尽可能的减少Buffer的大小。注意,这里是设备无关的优化,是无法利用硬件Memory特性的。
-4- 将HloModule转化为LLVM IR
-5,6- 利用LLVM框架,将LLVM IR编译为二进制代码。

本文主要关注-3-,是XLA优化的核心。

对BufferAssigner::Run()进一步分解。

NVPTXCompiler::RunBackend()
  hlo_schedule = GpuHloSchedule::Build(*module, *stream_assignment, pointer_size_)
  //this analysis figures out which temp buffers are required to run the computation
  BufferAssigner::Run(hlo_schedule->ConsumeHloOrdering()...)
    assigner.CreateAssignment(HloModule, hlo_ordering, buffer_size)
      liveness = BufferLiveness::Run()
      assignment = new BufferAssignment(module, liveness, ...)
      set<LogicalBuffer*> colocated_buffers
      set<BufferAllocation::Index> colocated_allocations
      vector<ColocatedBufferSet> colocated_buffer_sets
      BuildColocatedBufferSets(&colocated_buffer_sets)
      colorer_(assignment->liveness())
      AssignColocatedBufferSets(colocated_buffer_sets, assignment, &colocated_buffers, &colocated_allocations);
      GatherComputationsByAllocationType(module, &thread_local_computations, &global_computations)
      for computation : global_computations:
        AssignBufferForComputation(computation, false, buffers_to_assign_sequentially)
      AssignBuffersWithSequentialOrdering(buffers_to_assign_sequentially, ,assignment)
      for computation : thread_local_computations:
        AssignBuffersForComputation()
      for buffer : assignment->liveness().maybe_live_out_buffers():
        if assignment->HasAllocation(buffer):
          assignment->GetMutableAssignedAllocation(buffer).set_mayby_live_out(true)
      assignment->CombineTempAllocations()
      return std::move(assignment)
  entry_computation->Accept(&ir_emitter)
  CompileToPtx() 
  CompilePtxOrGetCachedResult()

-6- 进行BufferLiveness分析,分析整个HloModule的LogicalBuffer的干涉关系,为后续优化提供依据
-11- BuildColocatedBufferSets, 依据Bufferliveness的分析,将所有的LogicalBuffer分为几个Bufferset,并进行初步的Set融合,每个Bufferset内
参照注释, colocated buffer sets, 每个set都是一组可以共享BufferAllocation的LogicalBuffer, 共享Allocation,意味着共享同一块物理内存(GPU的显存)
-12- colorer_ 缺省被赋值为BufferLiveness::DefaultColorer(), 所有的LogicalBuffer实例的color都会被设置为0
-13- AssignColocatedBufferSets, 为Bufferset分配BufferAllocation, 每一个LogicalBufferSet 与其关联, 这里用到了buffer_size_, 这个函数是判断一个LogicalBuffer大小, LogicalBuffer的大小要和相应的Allocation一样, 具体可以参考tf2xla/while_op.cc tf2xla/if_op.cc xla/client/builder.cc kConditional代码,可以看到明显的要求各个body的Shape要一致。通过TEST用例也能确认
-14- GatherComputationsByAllocationType,根据内含的LogicalBuffer的属性,将Allocation分为global和thread local两类,这部分是理解显存优化的关键,后文详细
-16- AssignBufferForComputation,关联Allocation和XlaComputation,此调用点只针对global,temp buffer被收集到buffers_to_assign_sequentially, 延后处理,

BufferLiveness::Run()

Continue reading