Tensorflow XLA HLO — BufferLiveness

XLA Sevice 的 High Level Optimization可以分为设备无关优化和设备相关优化两部分, 设备无关优化是Buffer优化, 以便使用更少的主存或显存(以及其他硬件backend), 设备相关的优化主要是根据硬件特性优化执行流, HLO设备无关优化主要分为3个阶段:

  1. 支持分析指定Buffer的Liveness
  2. 构造逻辑Memory对象到物理Memory对象的映射
  3. 优化物理Memory, 交给特定Backend做specific的进一步优化

本文分析第一个阶段–BufferLiveness分析

所以,BufferLiveness, 就是获取Buffer生命周期关系, 以便决定Buffer复用策略. 源码整体调用栈如下, 该阶段进一步也分3个过程: LogicalBufferAnalysis->TuplePointsToAnalysis->BufferLiveness

BufferAssigner::Run()
  assigner::CreateAssignment()
    liveness = BufferLiveness::Run(module, std::move(hlo_ordering)  //class BufferLiveness
      liveness = new BufferLiveness()
      liveness->Analyze()
        points_to_analysis_ = TuplePointsToAnalysis::Run()
          logical_buffer_analysis = LogicalBufferAnalysis::Run()
            analysis = new LogicalBufferAnalysis()
            analysis.Analyze()
            return analysis
          analysis = new TuplePointsToAnalysis(logical_buffer_analysis)
          analysis.Analyze()
          return analysis
        maybe_live_out_buffers_ = points_to_analysis_->GetPointsToSet(root).CreateFlattenedSet()
      return liveness
    assignment(new BufferAssignment(module, std::move(liveness)

-1- Memory优化入口
-3-16- 获取BufferLiveness, 对比-16-, 获取到的BufferLiveness会用于支撑Buffer优化决策,用OOP的方式优雅的实现pipeline
-8- 获取LogicalBufferAnalysis
-9- 根据HloInstruction的Shape, 构造一组相应的LogicalBuffer, Shape本身是一个树结构,表示一个HloInstruction的输出形状, 一个LogicalBuffer对应一个Shape的一个节点(Subshape), 表示一块逻辑buffer, 用pair<HloInstruction*, ShapeIndex>来索引一个LogicalBuffer。
-10- 作为Liveness分析的原材料-LogicalBuffer已经构造完毕, 存储下来, 并返回analysis, 准备进入下一阶段, 进行TuplePointsToAnalysis.
-11- 用存储有所有LogicalBuffer信息的LogicalBufferAnalysis实例构造TuplePointsToAnalysis实例, Tuple, 用来描述Buffer树状结构的方式, E = {%1, %2, {%3, %4}} 表示这样一个树:深为3, 深1的节点有一个, 树的根, 深为2的节点有3个, %1, %2, 没体现名字的节点暂且叫”X”, 深为3的节点有2个, %3, %4, 这两个节点是上一层中”X”的子节点. PointsTo, “指向”, 在前面的例子中, root的”PointsTo”就是%1, %2和”X”, “X”的”PointsTo”就是%3和%4, 所以TuplePointsToAnalysis就是分析整个计算图中的LogicalBuffer依赖关系并存储在PointsToSet中,后面会详细分析
-12- 执行分析逻辑.
-13- 对计算图中的LogicalBuffer的依赖关系分析完毕, 存储下来, 并返回analysis, 准备进入下一阶段, 进行BufferLiveness
-14- 用TuplePointsToAnalysis实例获取root的alias_buffer, 也就是潜在的需要传出的HloModule的Buffer, 存储在maybe_live_out_buffers_中.
-15- 返回liveness实例, TuplePointsToAnalysis实例会存储LogicalBuffer的依赖关系, 但BufferLiveness并不会存储每一个LogicalBuffer的”liveness”, 而是基于TuplePointsToAnalysis封装了一组判断特定LogicalBuffer的函数.
-16- 将BufferLiveness实例传入构造LogicalBuffer与BufferAllocation映射关系的BufferAssignment实例.

Continue reading