在Medium的上一篇博文中,我们研究了一个简单的路径跟踪器,它演示了如何编写Poplar/C++程序,该程序利用Graphcore的智能处理器(IPU)上可用的大量tile。IPU不仅由于其大量独立内核而具有大规模并行性,而且每个tile内也有显著的并行性可供使用。在本文中,我们将优化一个简单的Poplar小代码,以利用这种tile级的并行性。
IPU上的tile级并行性
一个IPU由许多被称为tile的内核组成,单个GC200芯片包含1472个可运行完全独立程序的tile。每个tile中,我们可以通过两种其他方式并行工作。每个tile有6个硬件节点线程,它们以循环(桶)调度方式调度,每个上下文依次执行一条指令。虽然这是一系列指令,每个线程一个指令,但任何存储访问都以与其他节点上下文并行的方式继续,因此从每个上下文中看到的读/写都是单周期的。这些节点线程中的每一个都执行自己的指令流,从而使IPU能够同时执行8832个MIMD程序。tile中的另一种并行形式是向量SIMD浮点算术:每个tile上的每个节点每个周期可以发出2 x fp32或4 x fp16向量化算术运算。
优化小代码以使用节点线程
我们可以通过两种方式为每个tile上的线程创建工作。第一个是在计算图编程级别:如果我们在同一个计算集内创建没有数据依赖关系的顶点,那么Poplar的计算图编译器将安排这些顶点在节点之间并行运行。这就是我们在上一篇博文(此处为计算图构建代码)中并行化路径跟踪小代码的方式。这种技术的优点是我们可以在计算图级别描述它,从而简化小代码/内核实施。缺点是我们可能最终将顶点状态复制六次,这可能会对每个tile的存储消耗产生影响。
幸运的是,Poplar提供了用于C++小代码的替代Vertex类:MultiVertex。在MultiVertex中,节点ID被直接传递给小代码的计算方法。这样做的好处是现在可以在节点之间共享顶点状态,并且节点级别的并行性在低级别代码中是明确的。缺点是代码复杂性增加,因为工作需要在小代码中明确拆分,并且程序员需要注意一些线程安全问题(详细信息请参阅Vertex 编程指南)。
MultiVertex
在最初的路径跟踪代码中,我们没有费心在节点上并行化光线生成,因为它不是瓶颈。尽管这是优化工作的一个不太好的候选项,但由于小代码的简单性,它提供了一个很好的关于如何使用MultiVertex的示例。让我们看看如何将GenerateCameraRays Vertex更改为MultiVertex。
以下是原始顶点:

请注意,括号<<>>之间的注释代码片段要么在别处列出以避免杂乱,要么为了简洁而完全省略。例如,顶点状态将对Vertex和MultiVertex通用,并在此处单独列出:

将Vertex更改为MultiVertex是比较容易的部分。我们简单地从不同的基继承并更改计算方法的签名,得到如下结果:

您可以看到我们现在可以直接在计算方法中访问workerId。我们还可以调用 numWorkers()来获取可用的节点线程数量(IPU的未来设计可以自由更改节点线程的数量,因此我们避免了对六个节点线程进行硬编码,以实现向前兼容性)。现在我们需要决定如何在节点之间分配处理。让我们看看原始顶点的Lambda片段<< Lambda to consume random numbers>>:

每次通过消耗下一个条目以调用此函数时,函数都会遍历噪声缓冲区。我们在MultiVertex中可以进行的最简单的更改是在不同的抵销(线程的workerId)中启动每个线程,然后在每次调用时通过节点计数增加索引:

现在我们消耗随机数的速度提高了6倍。那很简单!光线生成循环稍微复杂一些,但原理是相同的。来看看原始Vertex的代码片段<<Ray generation loop>>:

在这里,如何在节点之间拆分两个循环,我们有了更多选择。因为IPU tile直接访问其私有SRAM,我们不一定需要为了实现缓存一致性而确保节点线程访问连续的存储位置。但是,稍后我们将尝试对计算进行向量化,以便每个线程的整行处理对那些SIMD操作是友好的。考虑到这一点,我们将与之前相同的原则应用于行索引:每个节点处理一行,并且它们从节点计数开始抵销。我们还需要以相同的方式抵销输出索引。这导致MultiVertex中出现以下循环:

我们根本不需要修改内部循环,我们已经完成了到MultiVertex的转换。小心一点,确保索引正确,这是一种非常直截了当的方法,可以将此类代码的速度提高6倍。
向量化内循环
几乎所有现代架构都具有某种形式的SIMD指令级并行性,IPU也不例外。Popc基于LLVM,并会在可能的情况下自动向量化代码,但我们也可以在IPU上手动向量化,而无需编写汇编代码。在header <ipu_vector_math>中有一些向量化效用程序函数,我们还将使用<ipu_memory_intrinsics>,它让我们能够访问一些补充向量操作的特殊存储指令。
与在节点线程上并行化相比,手动向量化小代码将是一个更具侵入性的变化(因此实际上我们应该首先通过分析来确保值得付出努力)。但是对于这个插图,让我们直接跳进去看看MultiVertex的SIMD版本的新结构:

我们首先纳入我们需要的标题。请注意,到目前为止,我们一直在使用普通的C++。但是,popc(编译C++ 小代码)是一个多目标编译器:它默认为三种架构生成代码:ipu1、ipu2和CPU。CPU目标允许使用IPUModel设备。为了保持对IPUModel/CPU目标的支持,我们需要使用#ifdef __IPU__保护IPU特定代码(注意,这包括导入IPU特定header)。
接下来,我们使用不同的名称GenerateCameraRaysSIMD创建一个新顶点,这不是绝对必要的,但它允许更广泛的应用程序通过在计算图构建时选择顶点名称来启用/禁用向量化。SIMD顶点对antiAliasNoise和光线顶点状态进行了修改,但所有其他状态保持不变。最重要的修改是确保数据对齐到8字节边界(这将允许我们使用特殊的加载/存储指令)。另一个指定了更紧凑的指针格式(我们可以这样做,因为我们知道来自其他状态的数据的大小)。
上面的最终结构更改是将Lambda函数和库调用内联到light::pixelToRay,并将内部循环展开两倍(因此您可以看到上面的列索引现在增加了2而不是1)。这些更改会损害可读性,但我们预计手动向量化代码会非常冗长。
现在我们已经看到了新的结构,我们可以看看具体的片段。首先在<< Setup pionters >>中看看如何访问噪声缓冲区和输出光线:

第一行为每个节点创建一个指向起始行的指针:这是一个指向4个半值(fp16)的向量类型的指针。为了抵销它,我们使用workerId索引到噪声缓冲区并乘以4(因为我们将指针转换为向量类型,每个节点将在四个值的块中消耗噪声)。在下一行,我们设置了一个恒定的half4值(这只是将噪声比例因子广播到四宽的半向量)。设置行索引和输出指针也相对简单:和之前一样,startRow被workerId抵销,而新的输出指针被workerId乘以输出光线成分的数量(注意每个像素有两个,所以numRayComponentsInRow是列乘以二)抵销。
在我们进入内部循环之前的下一个设置代码片段是<< Pre-compute values related to camera model >>(与相机模型相关的预计算值)。因为我们内联了light::pixelToRay,所以我们看到其中使用的许多值都是常数(与相机的视野相关)。因此,我们可以预先计算一些用于内循环的向量值:

我们不会详细介绍这个简单的视野相机模型。需要注意的重要的一点是,我们可以在这里预先计算许多向量值,并且ipu::命名空间中有对向量类型进行操作的函数(例如,tan将在half2向量上逐元素地应用)。另一件要注意的事情是我们将startCol索引预置为一半。转换可能相对昂贵,因此最好在外循环中进行转换并保持内循环只执行浮点算术和存储操作。我们还预先计算了unrollInc,将使用它来把half2向量的第一个元素递增1,但保持第二个元素不变。
最后,我们准备好查看内部循环。由于我们已经完成了所有设置,这实际上非常简单:

您可以看到我们可以使用带有向量类型的标准算术运算符,因此代码非常易读。这个循环唯一值得注意的部分是特殊的加载存储功能。IPU能够发出一条同时递增指针的加载存储指令。我们使用它们一次读取和写入8字节(四个half2向量),这就是为什么我们需要在之前的顶点状态中指定8字节对齐。
这样我们就完成了。您可以在Graphcore的Github示例中查看完整代码:codelets.cpp。请注意,我们本可以使用half4向量类型将内部循环向量化为另一个因子2,但我们将其作为练习留给读取器。
检查汇编
在此级别进行优化时,我们可能想做的最后一件事是查看为该内部循环生成的代码,以检查其是否有效。我们可以通过使用以下选项调用Popc来做到这一点:popc -O3 -S –target=ipu2 codelets.cpp -o simd.S,然后我们可以在输出文件中找到内部循环,如下所示:
.LBB1_5: # =>This Inner Loop
Header: Depth=1
{
add $m2, $m2, 2
f16v2add $a0, $a5, $a5
}
ld32 $a7, $m11, $m15, 14 # 4-byte Folded
Reload
{
ld64step $a2:3, $m15, $m0+=, 6
f16v2add $a5, $a5, $a7
}
ld32 $a6, $m11, $m15, 11 # 4-byte Folded
Reload
f16v2sub $a0, $a0, $a6
{
st32 $a0, $m11, $m15, 15
mov $a4, $a1
}
ld64 $a0:1, $m11, $m15, 6 # 8-byte Folded
Reload
f16v4mul $a2:3, $a0:1, $a2:3
{
ld32 $a0, $m11, $m15, 15
mov $a1, $a4
}
f16v2mul $a0, $a1, $a0
f16v2add $a4, $a5, $a5
f16v2sub $a4, $a4, $a6
f16v2mul $a4, $a1, $a4
f16v2add $a5, $a5, $a7
f16v2add $a2, $a0, $a2
f16v2add $a3, $a4, $a3
st64step $a2:3, $m15, $m1+=, 1
ld32 $m4, $m7, $m15, 5
ld32 $m4, $m4, $m15, 0
cmpult $m5, $m2, $m4
brnz $m5, .LBB1_5
我们不会在这里详细介绍指令集架构,但有一些有趣的事情需要注意。首先,我们可以看到有许多以f16v2为前缀的指令,这些是我们希望看到的2 x fp16向量算术指令。其次,可以看到我们使用的8字节(64位)加载/存储如下所示:
{
ld64step $a2:3, $m15, $m0+=, 6
f16v2add $a5, $a5, $a7
}
这实际上是一个单一的指令包:IPU tile可以在同一周期内双重发出浮点和存储指令。因此,如果我们计算内循环中发出的指令,则有21个(不是25个)。
总结
我们已经看到如何使用两种正交技术优化IPU小代码。首先,只需极少的改动,我们就可以通过利用tile上的所有节点线程来获得6倍的加速。其次,如有必要,我们可以通过利用内置向量类型和效用函数(无需借助汇编代码),使用手动向量化和专门的存储操作。Vertex编程的完整细节请见Graphcore的Vertex编程指南。