tensorflow xla LICM

先看什么是LICM，即训练不变代码外提，这是编译优化对循环优化的常用手段。

参考课件：

[http://www.cs.toronto.edu/~pekhimenko/courses/cscd70-w18/docs/Lecture%205%20[LICM%20and%20Strength%20Reduction]%2002.08.2018.pdf](http://www.cs.toronto.edu/~pekhimenko/courses/cscd70-w18/docs/Lecture 5 [LICM and Strength Reduction] 02.08.2018.pdf)

基本上分2个阶段：

找到循环
识别不变式，并把代码外提

对xla来说，只关注第2点，课件中对Loop invariant的定义

operands are defined outside loop or invariant themselves
not all loop invariant instructions can be moved to preheader

也就是说，判断循环不变式主要看operands，如果本身是常量或定义在循环外那就可以认为是invariant，同时需要主要，不是所有的loop invariant都可以提出去，其实也就是最终是否能做code motion还是取决于是否起到正面的效果并且不影响程序正确执行。

tensorflow xla里的代码在这里

[TF:XLA] Add while loop invariant code motion · tensorflow/tensorflow@51895fe

首先这个功能是整个优化pipeline中的一个pass，所以继承了HloModulePass并实现Run方法。

这个Run方法是所有pass必须实现的接口：

StatusOr Run(HloModule* module) override;

返回值表示是否对module做了修改。

针对WhileLoopInvariantCodeMotion的Run具体又做了2件事：找到while指令，提取while指令中的循环不变式。最终调用的是TryHoistingInvariantInstructionsFromWhileBody：

这个函数的大概实现如下：

当前这个pass只处理参数是元组的情况，对于只有真·一个参数的情况最好用单独的pass处理（我理解就是实现更简单）
如果循环本身只执行一次，跳出当前pass
找到所有循环不变式，并放入一个待提升表
对所有指令进行排序，保证操作数一定出现在操作符的前面
如果循环包含kDomain指令，跳出当前pass
对while循环体内的指令逐个判断：有副作用，是参数，控制结构都不做优化
当一个指令的所有操作数都是不变式或常量的时候，这个指令也是不变式
对于一些特定的指令，单独外提没有意义，先加入待提升表
收集可以提升的指令
指令提升

其中黑体的3处具体逻辑如下：

找到所有循环不变式

/*static*/ std::vector WhileUtil::GetInvariantGTEsForWhileBody(
    const HloComputation& while_body) {
  std::vector result;
  const HloInstruction::InstructionVector root_operands =
      while_body.root_instruction()->operands();
  for (int i = 0; i < root_operands.size(); i++) {
    HloInstruction* instr = root_operands[i];
    if (instr->opcode() == HloOpcode::kGetTupleElement &&
        instr->tuple_index() == i &&
        instr->operand(0) == while_body.parameter_instruction(0)) {
      result.push_back(instr);
    }
  }
  return result;
}

判断root_instruction的operands是否不变量，需要满足3个条件

特定指令不单独外提

bool WhileLoopInvariantCodeMotion::NotWorthHoistingIndividually(
    const HloInstruction& instruction) {
  switch (instruction.opcode()) {
    default:
      return false;
    // 注意这里有个kConstant
    case HloOpcode::kConstant:
      return !hoist_constants_;

    case HloOpcode::kBitcast:
    case HloOpcode::kBroadcast:
    case HloOpcode::kIota:
    case HloOpcode::kReshape:
    case HloOpcode::kReverse:
    case HloOpcode::kSlice:
    case HloOpcode::kTranspose:
    case HloOpcode::kTuple:
      return true;
  }
}

指令外提

指令外提在xla中的实际做法是把这些不变式附加到while指令的参数后，作为多余的参数处理。大致分为3步：1. 生成新的while指令（删除原来的while指令）；2. 把不变式加到新while指令的参数末尾；3. 删除不变式指令

看源码的时候有些通用函数的作用如下：

FindOrDie：容器内必存在的时候使用，防止不存在的时候程序出错，因为如果不存在执行这个函数程序直接崩溃
ReplaceInstruction: HloComputation的一个方法，删除旧的指令，添加新的指令
RemoveInstruction: HloComputation的方法，删除指令
AddInstruction: HloComputation的方法，添加指令
ReplaceAllUsesWith: HloInstruction的方法，把调用指令的use换成新的指令，
                      也就是更新使用了原指令的指令的内部引用
parameter_instruction(0): HloComputation的方法，获取computation的输入参数（指令），很多时候是个元组
CreateGetTupleElement: HloInstruction的方法，创建GetTupleElement指令，GetTupleElement和Tuple指令是一对互操作
AppendSuffix：元组后添加元素
ExtractPrefix：截取元组的前n个元素

所谓的优化，就是对这些instruction和computation按照某种优化算法进行操作变换，读源码的时候核心需要注意的也就是3点：

优化算法
指令
指令支持的操作

从头来看，首先就是指令：

https://www.tensorflow.org/xla/operation_semantics#while

其次是指令的操作：分别定义在hlo_instruction.h和hlo_computation.h中，最后是优化算法，这块可以从LLVM的pass中借鉴和学习，同时结合具体场景做等价的实现。

original refer

Pete Yuan

Comments