LLVM笔记(5) - SMS - 代码天地

1. SMS介绍
SMS(Swing Modulo Scheduling, 摇摆模调度)是一个基于循环的与架构无关的SWP(software pipelining)指令调度框架, 其目的是通过将当前迭代的指令与上一迭代同时发射来提升并行度. 考虑以下伪代码:

1 BB1:
2   A0 = A1  (1)
3   A2 = A1  (2)
4   A3 = A2  (3)
5   b BB1 if cond

不考虑跳转指令, 该循环需要独立发射三条指令(指令之间相互依赖). 但如果将其改变为以下方式:

1   A0 = A1
2 BB1:
3   A2 = A1
4   A3 = A2; A0 = A1;
5   b BB1 if cond
6   A2 = A1
7   A3 = A2

指令1与指令3间无直接依赖, 因此可以将本次迭代的指令3与下次迭代的指令1放在同一时刻发射, 提高并行度.

关于SMS原理的更多了解, 参见include/llvm/CodeGen/MachinePipeliner.h中涉及的三篇论文.
1. "Swing Modulo Scheduling: A Lifetime-Sensitive Approach"
2."Lifetime-Sensitive Modulo Scheduling in a Production Environment"
3. "An Implementation of Swing Modulo Scheduling With Extensions for Superblocks"

2. SMS实现分析
SMS算法由三部分组成. 首先, 计算出最小初始间隔(MII, minimal initiation interval). 第二步建立dependence graph, 计算每条指令的相关信息(ASAP ALAP MOV Height Depth ...). 最后使用论文中提及的方式为节点排序.

3. 常见概念
MII(minimal initiation interval)指完成循环所需的最小间隔, 它等于ResMII与RecMII两者的较大值, MII是理想的调度结果. ResMII(Resource MII)是根据一次循环所需的功能单元(FU, function unit)去除以机器所有功能单元的结果, 即如果一个循环包含4条指令而芯片只有两个运算单元则(不考虑指令间依赖)最少需要2 cycle才能完成一次循环. 更进一步, 如果循环中包含多类FU的使用(且每类FU之间相互无法替换), 则ResMII是分别对每一类FU计算ResMII后的最大值. RecMII(Recurrence MII)是循环中环路完成一次迭代所需的最小间隔, 如果循环存在多条数据链路则分别计算后取其最大值.
Prolog / Epilog分别指实现SWP后被提前到循环之前/放到循环之后执行的代码块. Kernel指实现SWP后的循环代码块.
Dependence Graph指指令间存在的依赖关系, 具体分为对寄存器的依赖与对内存的依赖. 对寄存器的依赖又分三种: 先写后读(RAW, read after write, 数据依赖, 又称真依赖)在LLVM中用SDep::Data表示, 先读后写(WAR, write after read, 反依赖)在LLVM中用SDep::Anti表示, 先写后写(WAW, write after write, 输出依赖)在LLVM中用SDep::Output表示. 对同一地址的读写也会产生依赖, 在LLVM中统一使用SDep::Order表示(不论是laod-store还是store-load还是store-store), 另外特殊指令带有side-effect属性时也使用顺序依赖表示. 顺序依赖又分为Barrier, AliasMem, Artificial, Weak等多个子类型, 其具体定义见include/llvm/CodeGen/ScheduleDAG.h.

4. 代码分析
MachinePipeliner类的入口为MachinePipeliner::swingModuloScheduler(), 函数首先检查循环是否仅包含一个BB(即没有change of flow). SwingSchedulerDAG是ScheduleDAGInstrs的子类, 父类的startBlock()/enterRegion()/exitRegion()/finishBlock()都是virtual的, 用于特定调度器的initial/clean up. SwingSchedulerDAG::schedule()(defined in lib/CodeGen/MachinePipeliner.cpp)实现swing modulo调度算法.
4.1. 建立dependence graph
为建立依赖图首先需要获取alias分析结果(否则无法分析order dependence), 再调用公共框架接口ScheduleDAGInstrs::buildSchedGraph()(defined in lib/CodeGen/ScheduleDAGInstrs.cpp)建立依赖图(关于ScheduleDAGInstrs类后文分析). 注意公共框架返回的依赖图是基于顺序执行的代码块的, 由于SWP的特殊性我们还要加上跨迭代的依赖(loop carried dependence, 即本次迭代的指令与下次迭代间的依赖).
由于在SSA模式下跨迭代的寄存器使用必然存在phi node, 因此我们只需为phi node添加依赖. SwingSchedulerDAG::updatePhiDependences()(defined in lib/CodeGen/MachinePipeliner.cpp)会为phi node添加data dep如果存在指令引用该phi, 并为该指令添加anti dep(在本次迭代的引用指令与下次迭代的phi指令间存在先读后写的关系).
对于跨迭代的内存访问间依赖交给SwingSchedulerDAG::addLoopCarriedDependences()(defined in lib/CodeGen/MachinePipeliner.cpp)处理. 该接口会遍历循环内所有load/store指令, 对每条load指令如果存在由alias关系的store则检查是否能通过load指令reach到store指令(一般是顺序load后store的情况, 不需要额外dep), 如果不能reach则为该sotre添加barrier dep. 再调用ScheduleDAGTopologicalSort::InitDAGTopologicalSorting()(defined in lib/CodeGen/ScheduleDAG.cpp)为节点排序方便之后处理, 至此依赖图基本建立完毕.
当前框架下还有两个针对依赖图的优化, SwingSchedulerDAG::changeDependences()会尝试转换指令使用前次迭代的值来降低RecMII(主要针对pre-inc/post-inc的load/store, 缩短dep chain). SwingSchedulerDAG::postprocessDAG()用来调用一些基于特定目的实现的Mutation的hook, 当前框架实现一个名为CopyToPhiMutation, 这个没看懂干嘛的, 为什么延后COPY的调度可以获得更好的性能?

为更具象的描述代码, 这里以一个testcase为例:

 1 [23:34:53] hansy@hansy:~$ cat test.c
 2 void test(int * __restrict in, int * __restrict out, int cnt)
 3 {
 4     int i;
 5     #pragma nounroll
 6     for (i = 0; i < cnt; i++) {
 7         *in = *out;
 8         in++;
 9         out++;
10     }
11 }
12 [23:34:59] hansy@hansy:~$ ~/llvm/llvm_build/bin/clang test.c -w -O2 -S -mllvm -debug-only=pipeliner --target=hexagon -mllvm -print-after-all 2>1.ll
13 [23:35:03] hansy@hansy:~$

其中__restrict是c99引入的关键字, 在这里用以减少内存访问依赖, #pragma nounroll是llvm的pragma, 保证循环不被展开优化, 用在这里减少循环内指令数(变相减少打印), 也可以去掉以上语法糖来观察SMS的变化. 篇幅有限, 这里截取部分打印.

loop block before SMS

448B    bb.4.for.body (address-taken):
    ; predecessors: %bb.4, %bb.6
      successors: %bb.5(0x04000000), %bb.4(0x7c000000); %bb.5(3.12%), %bb.4(96.88%)

464B      %2:intregs = PHI %7:intregs, %bb.6, %5:intregs, %bb.4
480B      %3:intregs = PHI %8:intregs, %bb.6, %6:intregs, %bb.4
496B      %15:intregs, %6:intregs = L2_loadri_pi %3:intregs(tied-def 1), 4 :: (load 4 from %ir.out.addr.07, !tbaa !2)
512B      %5:intregs = S2_storeri_pi %2:intregs(tied-def 0), 4, %15:intregs :: (store 4 into %ir.in.addr.08, !tbaa !2)
528B      ENDLOOP0 %bb.4, implicit-def $pc, implicit-def $lc0, implicit $sa0, implicit $lc0
544B      J2_jump %bb.5, implicit-def dead $pc

dependence graph

SU(0):   %2:intregs = PHI %7:intregs, %bb.6, %5:intregs, %bb.4
  # preds left       : 0
  # succs left       : 2
  # rdefs left       : 0
  Latency            : 0
  Depth              : 0
  Height             : 1
  Successors:
    SU(3): Data Latency=0 Reg=%2
    SU(3): Anti Latency=1
SU(1):   %3:intregs = PHI %8:intregs, %bb.6, %6:intregs, %bb.4
  # preds left       : 0
  # succs left       : 2
  # rdefs left       : 0
  Latency            : 0
  Depth              : 0
  Height             : 1
  Successors:
    SU(2): Data Latency=0 Reg=%3
    SU(2): Anti Latency=1
SU(2):   %15:intregs, %6:intregs = L2_loadri_pi %3:intregs(tied-def 1), 4 :: (load 4 from %ir.out.addr.07, !tbaa !2)
  # preds left       : 2
  # succs left       : 1
  # rdefs left       : 0
  Latency            : 1
  Depth              : 1
  Height             : 0
  Predecessors:
    SU(1): Data Latency=0 Reg=%3
    SU(1): Anti Latency=1
  Successors:
    SU(3): Data Latency=0 Reg=%15
SU(3):   %5:intregs = S2_storeri_pi %2:intregs(tied-def 0), 4, %15:intregs :: (store 4 into %ir.in.addr.08, !tbaa !2)
  # preds left       : 3
  # succs left       : 0
  # rdefs left       : 0
  Latency            : 1
  Depth              : 1
  Height             : 0
  Predecessors:
    SU(2): Data Latency=0 Reg=%15
    SU(0): Data Latency=0 Reg=%2
    SU(0): Anti Latency=1
ExitSU:   ENDLOOP0 %bb.4, implicit-def $pc, implicit-def $lc0, implicit $sa0, implicit $lc0
  # preds left       : 0
  # succs left       : 0
  # rdefs left       : 0
  Latency            : 0
  Depth              : 0
  Height             : 0

NodeFunction & NodeSets

    Node 0:
       ASAP = 0
       ALAP = 0
       MOV  = 0
       D    = 0
       H    = 1
       ZLD  = 0
       ZLH  = 1
    Node 1:
       ASAP = 0
       ALAP = 0
       MOV  = 0
       D    = 0
       H    = 1
       ZLD  = 0
       ZLH  = 2
    Node 2:
       ASAP = 0
       ALAP = 0
       MOV  = 0
       D    = 1
       H    = 0
       ZLD  = 1
       ZLH  = 1
    Node 3:
       ASAP = 0
       ALAP = 0
       MOV  = 0
       D    = 1
       H    = 0
       ZLD  = 2
       ZLH  = 0

  Rec NodeSet Num nodes 2 rec 1 mov 0 depth 1 col 0
   SU(0) %2:intregs = PHI %7:intregs, %bb.6, %5:intregs, %bb.4
   SU(3) %5:intregs = S2_storeri_pi %2:intregs(tied-def 0), 4, %15:intregs :: (store 4 into %ir.in.addr.08, !tbaa !2)

  Rec NodeSet Num nodes 2 rec 1 mov 0 depth 1 col 0
   SU(1) %3:intregs = PHI %8:intregs, %bb.6, %6:intregs, %bb.4
   SU(2) %15:intregs, %6:intregs = L2_loadri_pi %3:intregs(tied-def 1), 4 :: (load 4 from %ir.out.addr.07, !tbaa !2)

  NodeSet Num nodes 2 rec 1 mov 0 depth 1 col 0
   SU(0) %2:intregs = PHI %7:intregs, %bb.6, %5:intregs, %bb.4
   SU(3) %5:intregs = S2_storeri_pi %2:intregs(tied-def 0), 4, %15:intregs :: (store 4 into %ir.in.addr.08, !tbaa !2)

  NodeSet Num nodes 2 rec 1 mov 0 depth 1 col 0
   SU(1) %3:intregs = PHI %8:intregs, %bb.6, %6:intregs, %bb.4
   SU(2) %15:intregs, %6:intregs = L2_loadri_pi %3:intregs(tied-def 1), 4 :: (load 4 from %ir.out.addr.07, !tbaa !2)

NodeSet size 2
  Bottom up (default) 3 0
   Switching order to top down
Done with Nodeset
NodeSet size 2
  Bottom up (preds) 2 1
   Switching order to top down
Done with Nodeset

Schedule

Try to schedule with 1
Inst (3)   %5:intregs = S2_storeri_pi %2:intregs(tied-def 0), 4, %15:intregs :: (store 4 into %ir.in.addr.08, !tbaa !2)

    es: -2147483648 ls: 2147483647 me: 2147483647 ms: -2147483648
    insert at cycle 0   %5:intregs = S2_storeri_pi %2:intregs(tied-def 0), 4, %15:intregs :: (store 4 into %ir.in.addr.08, !tbaa !2)
Inst (0)   %2:intregs = PHI %7:intregs, %bb.6, %5:intregs, %bb.4

    es: 0 ls: 0 me: 2147483647 ms: -2147483648
    insert at cycle 0   %2:intregs = PHI %7:intregs, %bb.6, %5:intregs, %bb.4
Inst (2)   %15:intregs, %6:intregs = L2_loadri_pi %3:intregs(tied-def 1), 4 :: (load 4 from %ir.out.addr.07, !tbaa !2)

    es: -2147483648 ls: 0 me: 2147483647 ms: -2147483648
    insert at cycle 0   %15:intregs, %6:intregs = L2_loadri_pi %3:intregs(tied-def 1), 4 :: (load 4 from %ir.out.addr.07, !tbaa !2)
Inst (1)   %3:intregs = PHI %8:intregs, %bb.6, %6:intregs, %bb.4

    es: 0 ls: 0 me: 2147483647 ms: -2147483648
    insert at cycle 0   %3:intregs = PHI %8:intregs, %bb.6, %6:intregs, %bb.4
Schedule Found? 1 (II=2)
cycle 0 (0) (0) %2:intregs = PHI %7:intregs, %bb.6, %5:intregs, %bb.4

cycle 0 (0) (1) %3:intregs = PHI %8:intregs, %bb.6, %6:intregs, %bb.4

cycle 0 (0) (2) %15:intregs, %6:intregs = L2_loadri_pi %3:intregs(tied-def 1), 4 :: (load 4 from %ir.out.addr.07, !tbaa !2)

cycle 0 (0) (3) %5:intregs = S2_storeri_pi %2:intregs(tied-def 0), 4, %15:intregs :: (store 4 into %ir.in.addr.08, !tbaa !2)

先来看下dependence graph, 循环内总共四条Machine IR, 其中两条为PHI, 另外两条分别为load与store. 先看第一条PHI, 因为它的def-reg %2被store指令使用(先写后读)所以存在真依赖, 同理它的use-reg %5被store指令定义所以又存在反依赖. 另一条PHI与load指令间依赖关系类似. 最后load定义的%15被store使用, 所以两者间还有一个真依赖. 注意虽然PHI与store之间都存在反依赖, 但其含义不同: PHI指令的反依赖表示PHI不能调度到本迭代的store后, store的反依赖表示不能调度到下次迭代的PHI之后.

4.2. 计算ResMII/RecMII
ResMII比较容易计算(最简单的办法是指令数除以功能单元数, 当然代码中使用更复杂的DFA做更精确的评估), SwingSchedulerDAG::calculateResMII()基于DFA实现了ResMII的计算.
RecMII的计算则复杂许多, 首先我们要找到所有的关键路径(critical path). SwingSchedulerDAG::findCircuits()会首先查找循环中所有的环路并将其记录到NodeSets中. 思路是首先交换反依赖(否则跨迭代环路不结束), 然后构建一个依赖矩阵, 依赖矩阵的作用是记录所有backedge, 判断回边的依据: 到phi的反依赖(必然是回边), 顺序的输出依赖链上第一个与最后一个(最后一个不能调度到下个迭代的第一个之前), 以及读写之间的顺序依赖(为什么写写之间的顺序依赖不算回边? 我的猜测是如果循环中只有写写那么相对顺序就不重要了, 只要保证最后一次写正确即可?). 最后通过索引改矩阵查找从每个节点出发回到该节点为止的一条链路.
可以看到例子中一共找到两条环路, 第一条phi到store指令的基址, store指令基址再到phi, 第二条类似phi到load指令再到phi. 找到所有环路后即可计算RecMII, 对每条环路计算所需的latency总和取最大值即可.

4.3. 计算NodeFunction与NodeOrder
为了方便之后的调度还需要计算节点的一些特殊属性, SwingSchedulerDAG::computeNodeFunctions()代码注释已经很好的说明这些属性的作用. 如字面意思ASAP指该节点最早调度时机, ALAP值该节点最晚调度时机, MOV指该节点可调度空间(等于ALAP - ASAP), D指节点深度, H指节点高度. 注意由于有zero latency指令的存在(一般是各类phi copy reg_sequence等伪指令), ASAP/ALAP与D/H存在区别.
计算NodeOrder(TODO).

4.4. 调度
SwingSchedulerDAG::schedulePipeline()按NodeOrder顺序排序. 节点所插入的cycle是由NodeFunction与已排序的节点间的依赖决定的. SMSchedule::computeStart()用来计算当前节点最早与最晚插入时机, 如果计算的earlystart > latestart则调度失败, 一般情况是NodeOrder顺序非最优顺序. 否则调用SMSchedule::insert()尝试插入节点(注意伪指令一般视作zerocost指令无需排序), 由于NodeOrder过程中只考虑了顺序的依赖, 因此这里还要考虑跨迭代的依赖(即不光考虑当前cycle是否能插入该节点, 还要考虑(cycle + II * stage)中的节点是否与该节点冲突, 否则最后无法将later stage中指令折叠到当前cycle). 所有节点排序完成后就得到顺序依赖的指令调度, 此时调用SMSchedule::finalizeSchedule()将later stage指令往前折叠. 该接口分为三步, 第一将later stage指令拷贝到first stage, 第二步建立每个stage的寄存器映射, 最后为每个cycle内的指令重新排序(包括zerocost指令). 第三步通过SMSchedule::orderDependence()实现.

4.5. 生成pipelined loop
finalizeSchedule()产生了并行的循环代码, 但这还不是最终的kernel, 期间还需要完成三件事: 重新计算寄存器(对应不同迭代), 重新生成phi节点, 生成prolog与epilog. 这些任务由SwingSchedulerDAG::generatePipelinedLoop()完成. 主要接口updateInstruction()根据VRMap与stage为寄存器重命名, generateExistingPhis()替换之前的phi节点, generatePhis生成新的phi节点. 这块是诟病最多的代码, 像reduceLoopCount()这些为Hexagon实现的hack, 最近社区还有讨论重构这块代码(http://lists.llvm.org/pipermail/llvm-dev/2019-July/134002.html).

5. 优化&问题定位
SMS是比较复杂的后端优化模块, 问题定位主要抓以上几个流程点: 首先看依赖图是否有问题, MII的计算是否正确, NodeFunction与NodeOrder计算, 排序结果是否有问题, prolog/epilog是否有误. 每个流程的问题不能影响下一流程. e.g. SMS流程后出现use before def问题, 如果觉得排序有误就不用看寄存器映射是否有问题, 先分析排序是否正确, 因为finalizeSchedule()后指令就不能重排序, 而VRMap是依据当前指令排序生成的, 排序不正确后面修改phi也不能解决问题.

LLVM笔记(5) - SMS

猜你喜欢