体系结构论文（四十九）：Partitioned Scheduling and Parallelism Assignment for Real-Time DNN Inference Ta【DAC‘24】

Partitioned Scheduling and Parallelism Assignment for Real-Time DNN Inference Tasks on Multi-TPU

这篇文章讨论的是如何在多TPU系统上进行DNN推理任务的分区调度和并行分配。简单来说，文章的目的是提高在多个TPU上同时运行DNN模型的效率。

一、Intro

• 流水线并行处理： 为了提升DNN推理的性能，模型的不同层被分为多个阶段，这些阶段可以在不同的TPU上并行处理。每个TPU负责处理模型的一部分，这样可以显著提高推理吞吐量，因为不同部分可以同时处理。

图1: 在四个TPU上执行十帧的流水线操作

• 该图展示了四个TPU（TPU1到TPU4）处理十个帧的流水线化执行过程。

• 处理流程：每个TPU处理某个阶段的帧。TPU1先加载第一帧的数据（(1)），而其他TPU等待数据传递。随着帧数据在流水线中的传递，TPU2在TPU1处理完成后开始处理帧(1)，同时TPU1开始处理帧(2)，以此类推。

• 等待缓冲释放：TPU1在处理帧(2)之后，需要等待缓冲区的释放才能继续处理后续帧。

• 线程关闭：当所有帧处理完毕后，各个TPU的处理线程依次关闭。

• TPU的优势：

  • 并行执行多个任务：通过同时处理多个任务，TPU可以显著提高处理速度。
  • 跨多个TPU分配模型：将大型模型分布到多个TPU上，能够让每个TPU的内部内存存储更多的模型参数，从而减少数据的传输时间，提高执行速度。

• 任务调度挑战： 在需要严格时间约束的系统中，必须开发一个调度模型，充分利用TPU的并行能力并减少任务重编程的开销（例如，加载DNN参数的时间）。文章提出了“非抢占性并行任务”（NPG）的调度模型，确保同一任务的多个线程能并行地在多个TPU上执行且不中断。

 图2展示了DNN推理在不同数量的TPU上运行时的性能表现。随着TPU数量的增加，推理时间逐渐减少。特别地，对于某些网络（如ResNet-152），当使用8个TPU时，推理速度可以达到12倍的加速。

• 核心贡献：

  • DNN推理的基准测试：文章展示了多个DNN模型在多TPU加速器上的推理表现。
  • NPG调度策略：文章提出了一种严格分区的NPG调度策略，用于多TPU系统中的DNN任务分配和调度。
  • 对比分析：文章将所提出的调度策略与其他两种最先进的全局NPG响应时间分析和联合调度策略进行了对比，验证了其在任务可调度性方面的优势。

二、 相关工作解析：

分区调度： 分区调度在运行时开销较低，但需要解决任务到处理器的“装箱问题”，这是一个NP-hard问题。为了应对这种挑战，针对顺序抢占式和非抢占式任务的近似方法被提出。不过，这些方法无法直接应用于NPG（非抢占性并行任务）任务，因为NPG任务还需要考虑处理器的分区和任务并行度的配置。

在分区调度中的“装箱问题”：

在分区调度中，处理器可以类比为“箱子”，而任务可以类比为“物品”。任务具有一定的资源需求（如CPU时间、内存、处理器核数等），而每个处理器都有固定的计算能力或资源限制。

具体来说，“任务到处理器的装箱问题”指的是如何将任务分配到不同的处理器上，确保每个处理器的资源利用率达到最优，同时避免资源过载和浪费。目标是在处理器不超载的前提下，尽可能高效地使用最少数量的处理器。

全局gang调度： gang调度（即同时在多个处理器上运行同一任务的多个线程）已被用于抢占式系统中，非抢占式系统中，Nelissen等人提出了针对moldable gang任务的响应时间分析模型。但本文中考虑的是sporadic的任务模型。

并行任务的分区调度： 严格分区调度的思想很早就被提出。本文在此基础上扩展了该方法，支持处理器分区、任务分配和并行度选择。与此相对的另一种方法是“静态gang任务分配”，该方法允许不同并行度的任务共享处理器，但无法应用于本文讨论的NPG任务。

联合调度： 联合调度将重负载任务分配给专门的处理器，并使用全局调度器来处理剩余的轻负载任务。而严格分区调度则只使用单处理器调度器，并且每个处理器分区可以被多个任务共享。本文会将NPG严格分区调度与现有的联合调度策略进行性能对比。

三、 任务和平台模型解析：

1. 平台模型： 本文假设系统包含M个相同的处理器，任务集𝜏包含n个独立的NPG任务，每个任务都是sporadic，即它们在不同时间点被触发，且每次触发之间间隔至少为𝑇𝑖时间单位。每个任务𝜏𝑖会生成一个job。

2. 任务的并行执行： 每个任务𝜏𝑖可以在m个处理器上并行执行，执行时间不超过𝐶𝑖,𝑚时间单位。为了确保任务调度的可预见性，任务不会自我挂起，也不会因其他任务的原因被阻塞，除非是处理器资源的争用导致的。

3. 任务利用率： 每个任务的利用率定义为𝑈𝑖,𝑚 = 𝐶𝑖,𝑚 · 𝑚/𝑇𝑖，即任务在m个处理器上的并行度乘以其执行时间除以任务的触发间隔。此外，参考任务集的利用率是所有任务在并行度为1时的总利用率。

四、NPG严格分区调度策略

4.1.1 避免无限优先级反转：

• 在NPG全局调度中，存在一种现象叫做“二维阻塞”（2-D blocking），如图3所示。即当多个低优先级任务的执行时间相互重叠时，会阻碍高优先级任务的执行。
• 例子中，任务𝑱1具有最高优先级，但需要两个TPU，但由于𝑱2、𝑱5、𝑱7、𝑱3等低优先级任务只需要一个TPU，他们已经占用处理器，𝑱1不得不等待它们完成后才能开始执行。这种情况可以导致较长的优先级倒置时间。

为了避免无限的优先级反转，可以让所有任务在相同的并行度下运行，确保阻塞时间仅由最长的低优先级任务的执行时间决定。然而，不同任务可能需要不同的并行度，因此文章提出了将任务划分为不同的子集，每个子集中的任务拥有相同的并行度，并将这些子集分配给独立的处理器区间。

4.1.2 分区调度的优势：

• 分区调度与全局调度不同，它将任务静态分配给特定的处理器。
  • 在全局调度中（图4a），𝑱1的优先级较高，但由于低优先级的任务𝑱2、𝑱3、𝑱4已经在处理器上运行，𝑱1不得不等待，造成了较长的阻塞。
  • 在分区调度中（图4b），通过将𝑱1与低优先级任务分配到不同的处理器上，𝑱1受到的阻塞减少了一半。这是因为在分区调度中，高优先级任务与低优先级任务不会共享同一个处理器，从而减少了低优先级任务对高优先级任务的影响。

4.1.3 大分区的优势：

• 尽管增加分区数量可以减少低优先级任务的阻塞，但更少的大分区也有其优势。
  • 图5a显示了当使用更多的小分区时，任务𝑱1、𝑱2和𝑱3被分配到不同的处理器，导致任务的阻塞时间更长。
  • 图5b则展示了较少大分区的情况下，任务𝑱1、𝑱2、𝑱3可以在同一处理器上运行，阻塞情况得以减少，处理器的利用率也得到了优化。

NPG严格分区调度的主要思想是将处理器和任务划分为互不交叉的子集（分区），并且每个分区中只允许具有相同并行度的任务执行。通过这种分区方式，任务调度得以优化，减少了高优先级任务被低优先级任务阻塞的情况。

4.2.1 两个主要决策：

1. 处理器分区（Processor Partitioning）：

   • 将处理器集Π划分为多个不相交的处理器分区𝜌，每个分区包含𝑚个处理器。每个分区中的处理器只能运行特定并行度的任务。
   • 形式化表达：𝜌𝑚1 𝑗1和𝜌𝑚2 𝑗2是不同的处理器分区，并且每个处理器只能属于一个分区。

2. 任务分配（Task Assignment）：

   • 每个任务集𝜏被映射到不同的处理器分区，每个分区内的任务具有相同的并行度𝑚。
   • 每个任务只能分配到一个处理器分区，并且不同分区之间的任务不会共享处理器。这意味着每个分区在同一时间只能执行一个任务，因为任务需要占据该分区中的所有处理器。

4.2.2 调度策略：

• 在线调度器（Online Scheduler）：每个处理器分区中的任务由一个非抢占式固定优先级（NP-FP）调度器进行调度。这种调度策略意味着分区内的任务在开始执行后不会被中断，并且通过单处理器的调度分析来验证任务是否可调度。

• 可调度性测试：本文采用了非抢占式固定优先级的响应时间分析（NP-FP）来验证任务的可调度性。具体来说，分析会计算每个任务的最晚启动时间，确保任务可以在其截止时间之前完成。

• 公式中的每一部分：
  • Cj表示比任务𝜏𝑖优先级低的任务的执行时间。
  • l表示第几个工作。
  • hp(i)表示优先级高于任务𝜏𝑖的任务集，lp(i)表示优先级低于任务𝜏𝑖的任务集。
  • 通过分析任务在忙期（busy-period）内的所有工作来验证任务是否可调度。

五、NPG 严格分区启发式

5.1 总览

在这一部分，文章提出了一个启发式方法，名为NPG-SP∗，用于辅助NPG严格分区调度策略进行处理器分区和任务分配决策。该启发式方法的目的是解决如何高效地分配任务和处理器，以提高系统的资源利用率和任务的可调度性。

任务优先级分配：任务的优先级按照截止时间单调递减的原则分配。

主要步骤：

1. 初始化任务集：开始时，将任务集𝜏赋值给未分配任务集𝜏̅，准备对每个任务进行分区和调度。

2. 初始化处理器分区：将M个处理器分配给单独的分区，初始时每个分区包含一个处理器。

3. 初始化任务分配：任务的分配集𝜏(𝜌)初始化为空。

4. 核心循环：在有多个分区时，循环执行任务分配操作。算法会尝试通过BestVolumeBinPack对任务进行分配。

   • 如果任务集𝜏̅为空，表示所有任务都已成功分配，则调度成功，设置sched = True并退出循环。

   • 如果只剩一个分区，说明任务无法被进一步分配，设置sched = False并退出循环。

   • 如果还有多个分区但未能分配所有任务，则调用MergePartitions合并分区，尝试提高资源利用率。

5. 输出结果：返回处理器分区𝜌、任务分配结果𝜏(𝜌)以及可调度性结果sched。

算法框架：初始时将所有处理器分为大小为1的分区，并初始化任务分配为空集。算法通过迭代尝试将任务分配到现有分区。如果任务无法成功分配，则尝试合并分区以提高资源利用率。

• 如果所有任务都被成功分配，算法返回成功。
• 如果所有分区已经合并为一个大分区，且仍然有未分配的任务，则算法返回失败。

5.2 启发式算法的细节

BestVolumeBinPack 组件解析（算法2）：

目标：将任务分配到资源利用率最高的分区，以实现最优的资源使用。

1. 初始化未分配任务集：将未分配任务集𝜏̅初始化为空。

2. 按优先级排序任务：将任务按优先级从高到低排序，确保高优先级任务先被分配。

3. 分配任务：对于每个任务𝜏𝑖，按升序排序处理器分区，根据并行度和任务利用率（𝑈𝑖,𝑚）进行分配。

   • 如果找到可调度的分区，则将任务分配到该分区，并继续下一个任务。

   • 如果任务无法分配到任何分区，调用LocalSearch尝试通过移动已分配任务来腾出空间。

4. 返回未分配任务：如果任务仍然无法分配，则将其加入未分配任务集𝜏̅并返回。

LocalSearch 组件解析（算法3）：

目标：当某个任务无法直接分配时，通过移动其他任务来尝试腾出空间。

主要步骤：

1. 遍历每个分区：对于每个分区𝜌𝑚𝑗，检查该分区中的任务𝜏𝑖是否可以通过删除某个任务而腾出空间给新的任务𝜏𝑘。

2. 任务移动：如果发现某个任务可以移出当前分区，并且移出任务后新的任务可被成功调度，则执行任务交换操作，将任务从当前分区移至其他分区，并返回sched = True，表示任务调度成功。

3. 返回结果：如果尝试所有分区后仍无法为新任务分配空间，则返回sched = False，表示无法调度。

MergePartitions 组件解析（算法4）：

目标：当分区资源利用率不足时，尝试合并分区以提高资源利用率。

主要步骤：

1. 选择合并分区：选择两个任务利用率最低的分区𝜌𝑚1𝑗1和𝜌𝑚2𝑗2，将它们合并为一个较大的分区。

2. 重新分配任务：将这两个分区中的任务移出，加入未分配任务集𝜏̅，然后在下一个迭代中尝试将这些任务分配到新的分区。

3. 返回新的分区：最终返回合并后的新分区𝜌′。

5.3 时间复杂度

文章在这一部分分析了NPG-SP∗算法的时间复杂度。算法的两个主要组件是 BestVolumeBinPack 和 MergePartitions，它们的时间复杂度如下：

BestVolumeBinPack的时间复杂度：

• 任务分配：对于每个任务，检查它是否可以在每个分区上成功调度。这个过程的时间复杂度为O(𝑀Φ)，其中𝑀是处理器的数量，Φ是单处理器可调度性测试的复杂度。
• 局部搜索：如果任务无法分配，局部搜索需要对每个已分配的任务进行调度检查，最坏情况下复杂度为O(𝑛𝑀Φ)，𝑛是任务的数量。
• 总时间复杂度：因此，BestVolumeBinPack组件的总复杂度为O(𝑛²𝑀Φ)。

MergePartitions的时间复杂度：

• 选择分区：选择两个任务利用率最低的分区，时间复杂度为O(𝑀)。
• 重新分配任务：将任务从合并的分区移出，时间复杂度为O(𝑛)。
• 总时间复杂度：因此，MergePartitions组件每次迭代的时间复杂度为O(𝑛²𝑀Φ)。

NPG-SP∗的总时间复杂度：

• 因为最坏情况下算法有𝑀次迭代，所以NPG-SP∗的总复杂度为O(𝑛²𝑀²Φ)。
• 其中，Φ是单处理器的精确非抢占固定优先级调度分析的复杂度。该复杂度与任务数量和最大并行度成正比。

NPG-SP∗整体算法详细总结

1. 初始化任务与处理器分区

首先，算法接受一个任务集𝜏和可用的处理器数量M。初始时，所有处理器被分为M个分区，每个分区仅包含一个处理器。这意味着每个处理器在一开始是独立的，任务分配也将基于这种处理器划分。然后，算法会初始化任务分配集𝜏(𝜌)，其初始状态为空集，表示任务尚未分配到任何处理器。

2. 任务分配：调用BestVolumeBinPack

在初始化后，算法进入核心循环，调用BestVolumeBinPack组件来进行任务分配。BestVolumeBinPack的主要目标是尽可能地将任务分配到最合适的分区中，最大化资源的利用率。任务分配的具体过程如下：

• 任务排序：BestVolumeBinPack首先对所有任务按优先级从高到低进行排序，优先处理高优先级的任务。

• 分区选择：对于每个任务𝜏𝑖，算法会遍历所有分区，检查每个分区是否适合分配该任务。此时，BestVolumeBinPack会依据每个分区的资源利用率（𝑈𝑖,𝑚，即该分区执行任务的计算开销）对分区进行升序排序，并尝试将任务分配给资源最充裕的分区。

• 可调度性检查：当任务𝜏𝑖被尝试分配到某个分区𝜌𝑚𝑗时，BestVolumeBinPack会通过检查任务在该分区的可调度性来确定任务是否可以在给定的时间约束内执行。如果该任务可以成功分配到分区，则将其加入该分区的任务集合𝜏(𝜌𝑚𝑗)。

3. 局部搜索：调用LocalSearch

如果BestVolumeBinPack发现某个任务无法被成功分配到现有的分区，那么它会调用LocalSearch组件进行局部调整。LocalSearch尝试通过重新分配已经分配的任务，为当前任务腾出空间：

• 任务交换：LocalSearch会遍历每个分区，并尝试将分区内已分配的任务移出。算法会检查是否可以通过移出一个已分配的任务，为当前任务𝜏𝑘腾出空间。如果可以通过重新分配已分配的任务𝜏𝑖使得新的任务可调度，则执行任务交换操作，将任务𝜏𝑖移出当前分区，并将新任务𝜏𝑘分配进来。

• 返回结果：如果局部搜索成功调整任务，则返回成功，并继续处理下一个任务。如果局部搜索失败（即当前任务仍然无法调度），则任务会被标记为未分配，并加入未分配任务集𝜏̅。

4. 分区合并：调用MergePartitions

当任务无法通过现有的分区分配或局部调整解决时，算法会尝试通过MergePartitions组件合并分区来增加资源利用率。分区合并的过程如下：

• 选择分区：MergePartitions会选择资源利用率最低的两个分区，试图将这两个分区合并为一个更大的分区。通过合并处理器资源较少使用的分区，可以为更多任务腾出空间，增加系统的整体计算能力。

• 重新分配任务：在合并分区后，原分区中的任务将被移出，并加入未分配任务集中。随后，合并后的分区将被用来再次尝试分配这些任务。

5. 循环迭代与结束

在每次调用BestVolumeBinPack、LocalSearch和MergePartitions后，算法会检查是否所有任务都已成功分配。如果所有任务都已被分配且可调度，则算法返回成功。如果任务集𝜏̅为空，表示所有任务都可调度并完成分配，算法结束并返回处理器分区和任务分配结果。

如果处理器分区已合并为一个单一的大分区且仍然有未分配任务，则表示调度失败，算法返回失败，表示系统中存在无法调度的任务。