IPU程序的优化之一：初识内存布局

2022/07/07

IPU程序的优化之一：初识内存布局

作者:

现代处理器，计算速度远比内存的读写速度快。因此，不仅仅是IPU，其他各种处理器，都对内存布局有一定的要求，只有内存的布局满足一定的条件，才能充分发挥出应有的算力。不同的是，内存的布局在不同的处理器上有不同的表现形式。以GPU为例，在神经网络中，我们通常会将tensor的格式定义成NCHW，考虑到GPU算力的发挥，在数据加载指令数一定的情况下，我们需要追求更多的计算指令数，即前者和后者的比值越低，算力的发挥越充分。同时，考虑到连续访存，以及cache映射的效率，会尽可能的将tensor的格式定义成NC4HW4（或相似结构，比如NHC4W4）。

这部分内容涉及到代码示例，以及一些性能分析，受篇幅限制，本文不再赘述，感兴趣的话可以参考这个链接了解更多。

这种NC4HW4就可以理解成一种内存布局，它对于算力的发挥至关重要，可以说，它是后续一切优化工作的起点。

本文后续章节会反复提到PopVision。PopVision是Graphcore为开发者提供的一套性能诊断工具，对于分析一段IPU程序是否高效有着重要的作用，如欲了解更多详情，请访问这个链接。

什么是IPU的内存布局

内存布局对于充分发挥IPU的算力性能有特别重要的意义。首先，我们需要清楚什么是IPU上的内存布局。

以MK2为例，一个IPU有1472个tile，每个tile上有624K的内存（片上），一共约900M的内存，如图2所示。

现在，我们假设有一个float32的tensor，其shape是[64, 128, 256]。很显然，一个tile无法放下这么多的数据，必须将这个tensor的数据以合理方式分配到各个tile上，其分配的结果就是我们要说的内存布局。

以传统的CPU/GPU为例，其内存是中央控制式的，即一颗芯片可以访问/共享全部的内存空间。相反，在IPU内部，每个tile都是独立的，因此IPU对内存的访问是标准的分布式控制，即每个tile各有其独立的内存空间。在IPU程序中，若要跨tile对内存进行访问，就要涉及到tile间的数据交换，这样就会产生一定的消耗。如何尽可能降低这部分的消耗，也是我们需要重点关注的问题。

我们以一个简单的问题y=1-x为例，通过PopVision，观察一下IPU程序都有哪些消耗：

从图3可以看出，为了执行一个简单的程序，IPU需要执行这几个部分：

DoExchange，tile间的数据交换；
PreArrange, 针对输入数据，tile内部对内存的布局进行重排；
DeExecute, 实际要执行的程序；
PostArrage，如有需要，针对输出数据，tile内部对内存的布局进行重排。

除了实际要执行的程序，其他的消耗都是由内存布局引起的，从上图可知，内存布局对IPU程序的重要性。

为了降低由内存布局引起的各种消耗，通常我们需要从两个方面考虑问题：

以原地访存的方式，合理地在tile上分配内存和vertex。这样可以避免tile间的数据交换；
避免不连续的内存的拷贝，比如transpose/subsample/reverse等等，这也是为了避免tile间的数据交换以及一些tile内部的数据重排。

TileMapping

对于一个给定的计算任务，在IPU编程模型中，我们需要将输入数据、计算程序，和输出数据合理地映射（或称分配）到各个具体的tile上，这里的“映射”就是TileMapping。

图2描绘了IPU程序需要将哪些对象映射到具体的tile上，其中“Tensors”（包括t1和t2）表示输入数据，“Vertices”（蓝色部分的圆圈）表示要执行的计算程序，最后的t3和t4表示输出数据。

结合图4，关于内存布局我们需要清楚以下几点：

输入数据（input tensor），实际程序中，输入数据的布局是由上一个op决定的，或者是在程序开始处决定。作为op的开发者，我们无法确定输入数据的布局是什么样的。如果按照开发者的期望，对输入数据进行TileMapping，就会造成一个tensor的多重mapping，导致同一段数据被多个tile引用，会对程序造成极大的隐患，这点在后文将有详细的解释；
计算程序（vertex），vertex如何mapping取决于op的算法特性和开发者的设计，因此，由开发者决定vertex的TileMapping；
输出数据（output tensor），输出数据的映射由op的开发者决定，其映射的结果就是输出数据（下一个op的输入数据）的布局。

基于以上三点，为了合理规划vertex和输出数据的布局，如何确定输入数据的布局就非常重要，我们需要根据输入数据的布局来确定如何将vertex和输出数据映射到相应的tile上。

现在解释一下上述第一条，同一段数据被多个tile引用，会有什么不利影响呢？我们可以看一个例子。假设有一个float32的tensor，命名为A，其shape是[64, 32768]，要经过两个op，假设这两个op都仅仅做数据copy。初始化的时候，A按行进行TileMapping，这样就会映射到64个tile上（tile0-tile63）。我们再看一下两个op对其输入输出的映射步骤，正常情况下，第一个op是这样的：

Tensor B = graph.AddVarialble(poplar::FLOAT, {64, 32768}, “B”); //创建tensorB

for(int i= 0 ; I < 64; i ++)

  graph.setTileMapping(B[i], i);

prog.add(poplar::program::Copy(A, B)); //指定A往B进行copy

然后，第二个op是这样的：

Tensor C = graph.AddVarialble(poplar::FLOAT, {64, 32768}, “C”); //创建tensorC

for(int i= 0 ; i < 64; i ++)

  graph.setTileMapping(C[i], i);

prog.add(poplar::program::Copy(B, C));//指定B往C进行copy

为了说明什么是多重映射，我们假设在第二个op中对其输入（也就是B）进行了多重映射，假设又把B映射到了tile64-tile127，那么其代码像这样：

Tensor C = graph.AddVarialble(poplar::FLOAT, {64, 32768}, “C”); //创建tensorC

for(int i= 0 ; i < 64; i ++)

{

  graph.setTileMapping(B[i], i + 64);

  graph.setTileMapping(C[i], i);

}

prog.add(poplar::program::Copy(B, C)). //指定B往C进行copy

最后，我们需要通过PopVision观察一下二者的区别，如图5、图6所示：

从上图可以看到，B被多重（准确地说是2重）映射，在tile0-tile63和tile64-tile128都有一份副本，造成了内存的浪费。为了说明其危害，我们还可以做这样的假设：假设这个多重映射就发生在tile0-tile63，我们在B的高维度做transpose（仅仅是view层面）这样的操作。

Tensor C = graph.AddVarialble(poplar::FLOAT, {64, 32768}, “C”); //创建tensorC

Tensor BT = B.reshape({ 8, 8, 32768 });

BT = BT.dimShuffle({ 1, 0, 2 }); //在view层面，高维度进行transpose

BT = BT.reshape({ 64, 32768 })

for(int i= 0 ; i < 64; i ++)

{

  graph.setTileMapping(BT[i], i);

  graph.setTileMapping(C[i], i);

}

prog.add(poplar::program::Copy(B, C)). //指定B往C进行copy

综合图5、图6、图7可以看出，开发者在实践中要尽量避免这种多重映射的发生。

原地访存

原地访存，很容易联想到“inplace”。本文所述的“原地访存”跟通常所说的“inplace”是有区别的。一般而言，inplace是指，一段内存既是输入也是输出。本文所述的“原地访存”是指在原有输入数据的布局基础上对数据进行读取。

举个例子，假设有一个一维的tensorA，其shape是[32]，其中0-15在tile0上，16-31在tile1上，vertex在tile2上。显然，vertex跟数据不在同样的tile上，其对数据的读取必然涉及到tile间的数据交换，这种方式就不是原地访存。

相反，如果在tile0上配置一个vertex读取tensorA的0-15，在tile1上配置一个vertex读取tensorA的16-31，那么这种数据的读取方式就是原地访存。因为vertex和tensorA的各个切片是在相同的tile上的，这样就不会产生tile之间的数据交换，对于IPU程序的性能发挥有很大的帮助。

假设tensor是以NCHW来定义的。在许多程序中，数据的处理会不自觉地按照最后一个维度进行的，即一行一行地处理。这样我们就会按照逐行的方式确定vertex和输出数据的布局。1.3中已经说明了输入数据的布局是由上一个op决定的，开发者无法确定，输入数据有可能是这样的：

布局不是按照逐行的方式；
不是一个tensor实体，可能是由不同的tensor拼接起来的；
内存不是连续的。

如果op开发者仍然按照自己期望的逐行的方式确定自己的vertex和输出数据的布局，那么就会产生一定的消耗。举个例子，假设有A和B两个tensor，每个tensor的shape都是[32, 64, 128, 128]。每个tensor都按照线性的方式将数据均匀地映射到每个tile上，然后两者在最后一个维度进行拼接（concat），最终形成tensorC。那么，这个C的shape就是[32, 64, 128, 256]。我们假设按照逐行的方式对数据进行处理（计算任务就是简单的y=1-x），即vertex和输出数据都是按照逐行的方式进行tile的映射的。最后，我们通过PopVision，观察其执行的情况，如图8。

Cycles(DoExchange + PreArrange + Execute) = 117448

很显然，tile间的数据交换和tile内部的PreArrange占据了整个程序的大部分执行时间，实际程序的执行时间不到整个程序执行时间的一半，这造成了算力的极大浪费。

实际上，如果我们能够充分利用输入数据的布局来规划vertex和输出数据的布局，这些浪费是可以避免的。事实上，像y=1-x这样的计算任务，我们确实可以达成原地访存的目的。IPU编程的基本接口提供了一些关于分析tensor布局的接口，具体详情可以参考这个链接，这里简单介绍一下：

graph.getTileMapping：获取待定tensor在每个tile上的view层面的数据分布；
graph.getSortedContiguousRegions：根据上述view层面的数据分布，获取每个tile上物理上的分布，并且按照连续区块分组；
graph.reorderToSimplify：重排以简化待定tensor的view层面的结构，尽可能将物理上的连续区域合并，这个接口一般对于顺序无关（比如elementwise）的操作有用，本例即可适用。

以上述引用的例子为例，为了达成原地访存的目的，我们可以利用这些接口做一些准备工作：