TPU编程竞赛

由于Depthwise2D和Conv2D极其相似，因此主要讨论Conv2D，Depthwise2D可以仿照该逻辑编写代码）

算子结构体

typedef struct {
int N, IC, OC, H, W; //batchsize, 输入channel，输出channel，输入H，输入W
int kernel_h, kernel_w; //卷积核大小
int pad_top, pad_bottom, pad_left, pad_right; //四周padding
int stride_h, stride_w; //卷积核步长
int dilation_h, dilation_w; //空洞卷积
unsigned long long output_addr;
unsigned long long input_addr;
unsigned long long kernel_addr;
} __attribute__((packed)) param_t;

测试样例分析

.N = 4, .IC = 2, .OC = 1, .H = 1, .W = 3,
.N = 4, .IC = 3, .OC = 16, .H = 640, .W = 640
.N = 4, .IC = 3, .OC = 16, .H = 512, .W = 960
.N = 4, .IC = 3, .OC = 16, .H = 1080, .W = 192
.N = 4, .IC = 3, .OC = 24, .H = 384, .W = 384
.N = 4, .IC = 3, .OC = 96, .H = 227, .W = 227
.N = 4, .IC = 3, .OC = 192, .H = 127, .W = 127
.N = 4, .IC = 18, .OC = 36, .H = 192, .W = 256
.N = 4, .IC = 72, .OC = 72, .H = 160, .W = 160
.N = 4, .IC = 128, .OC = 256, .H = 50, .W = 50,
.N = 4, .IC = 160, .OC = 192, .H = 30, .W = 30,
.N = 4, .IC = 256, .OC = 512, .H = 128, .W = 128
.N = 4, .IC = 512, .OC = 512, .H = 28, .W = 28,
.N = 4, .IC = 1024, .OC = 546, .H = 10, .W = 10,
.N = 4, .IC = 2048, .OC = 256, .H = 28, .W = 28,
.N = 4, .IC = 4032, .OC = 672, .H = 11, .W = 11,

以上是本次竞赛卷积算子host端测试用例的部分基础信息，padding，stride，dilation，kernel_size等信息暂时不涉及。从测试样例可以看出，Conv2D测试样例的难点主要是

• 输入的IC太大
• 输入的H太大

解决了这两个问题，Conv2D算子基本就都能跑通了。

按N切分

在进行复杂的数据切分前，可以先完成按N切分的函数，该函数可以通过一部分用例。在阅读完ok_device_conv2d_demo.c示例代码后，读者应对卷积算子API的参数准备和设置有了一定了解。那么不难写出Naive版本的按N切分（伪）代码：

// 按照demo计算IC_new、kernel_h_ext等参数，需要注意的是output_shape等参数的.N = 1
local_addr_t output_addr = 0;
local_addr_t kernel_addr = ...; // 仿照demo
local_addr_t input_addr = ...; // 仿照demo
S2L(kernel_addr, param->kernel_addr); // 在循环外将kernel参数搬入local memory对应位置（kernel在不同batch运算时不会改变）
int batch_input_offset = 1 * input_shape.c * input_shape.h * input_shape.w;
int batch_output_offset = 1 * output_shape.c * output_shape.h * output_shape.w; // 计算每个batch在global memory的占用大小，用于S2L和L2S搬运地址的计算
for(int i = 0; i < param->H, i++){ // 每次计算一个batch
S2L(input_addr, param->input_addr + i * batch_input_offset);
conv2d(input_addr, kernel_addr, output_addr);
L2S(param->output_addr + i * batch_output_offset, output_addr);
}

在此基础上结合OKKernel的关于pingpong计算的文档，可以改写为并行计算模式（前提是在local memory中，output_size *2 + input_size *2 + kernel_size < LOCAL_MEM_SIZE/2）

// 按照demo计算IC_new、kernel_h_ext等参数，需要注意的是output_shape等参数的.N = 1
local_addr_t output_addr[2] = {0,...};
local_addr_t kernel_addr = ...; // 仿照demo
local_addr_t input_addr[2] = {..., ...}; // 仿照demo
S2L(kernel_addr, param->kernel_addr); // 在循环外将kernel参数搬入local memory对应位置（kernel在不同batch运算时不会改变）
int batch_input_offset = 1 * input_shape.c * input_shape.h * input_shape.w;
int batch_output_offset = 1 * output_shape.c * output_shape.h * output_shape.w; // 计算每个batch在global memory的占用大小，用于S2L和L2S搬运地址的计算
for(int i = 0; i < param->H+2, i++){ // 每次计算一个batch
parallel_start();
if (i < S)
S2L(input_addr[i%2], param->input_addr + i * batch_input_offset);
if (i > 0 && i < S + 1)
conv2d(input_addr[(i+1)%2], kernel_addr, output_addr[(i+1)%2]);
if (i > 1)
L2S(param->output_addr + (i-2) * batch_output_offset, output_addr[i%2]);
parallel_end();
}

按H拆分

只按照N进行切分，会遇到两个问题：

• 即使不用pingpong，也会出现一个batch塞不下的情况
• 虽然不用pingpong可以计算，但是却不满足使用pingpong的要求

因此需要考虑从卷积的HW维度进行切分，这里选择切分H维度（由于输入输出的W都是连续存储的，按照W切分可能会造成GDMA性能下降）。同时，由于卷积算子运算时，卷积核是不变动的，因此整个计算流程不需要做太大改动。

按H拆分的一个难点就是如何高效地、准确地进行拆分（因为有padding、stride和kernel_szie参数等都在影响滑动窗口变化的规律），一个较为简单的解决办法是：从输出特征图反推输入特征图的切分索引。

如上图所示，考虑一个最简单的情况，输入大小为 7 X 7，卷积核为3 X 3，padding=0，stride=1，dilation=1，输出 5 X 5，图中演示了在第一个滑动窗口上运算的示例。以行（H）为单位观察卷积的特点：

观察另一组参数以行（H）为单位的特点：输入大小为 7 X 7，卷积核为 3 X 3 ，padding=0，stride=2，dilation=1，输出 3 X 3 。

对比观察两个参数设置，可以发现无论参数如何设置，从输出来看，它的output_h是依次按顺序得到的，而输入的滑动窗口偏移则和其他参数相关。而从输出的行索引可以很简单的得到输入窗口的索引，可以得到以下索引推导公式（dilation可以转化为更大的kernel_size，不做讨论，公式中的输入索引可以理解为在添加了padding_top/bottom之后的特征图上的索引）：

如果padding_top不为0，那么真实的索引就是：

其他情况也可以如此推导。

那么可以写出按照N和H切分的（伪）代码。

// 按照demo计算IC_new、kernel_h_ext等参数，需要注意的是output_shape等参数的.N = 1
local_addr_t output_addr = 0;
int output_step = output_h*2;
do{
output_step/=2;// 从output_h（完整的batch）开始尝试切分
local_addr_t kernel_addr = output_addr+...; // 仿照demo
local_addr_t input_addr = kernel_addr+...; // 仿照demo
}while(input_addr + input_aligned_size > LOCAL_MEM_SIZE);// 当前切分H的step太大，继续衰减
S2L(kernel_addr, param->kernel_addr); // 在循环外将kernel参数搬入local memory对应位置（kernel在不同batch运算时不会改变）
int batch_input_offset = 1 * input_shape.c * input_shape.h * input_shape.w;
int batch_output_offset = 1 * output_shape.c * output_shape.h * output_shape.w;
// 计算每个batch在global memory的占用大小，用于S2L和L2S搬运地址的计算
for(int i = 0; i < param->H, i++){ // 每次计算一个batch
for(int output_h_idx = 0; output_h_idx < output_h, output_h_idx+=output_step){
int input_h_idx = ...;//可利用上面公式的完整推导计算
int output_h_size = ...;// 求本次计算得出的输出大小（<=output_step），对于首尾的特殊情况需要单独处理
int input_h_size = ...;//同理
S2L(input_addr, param->input_addr + i * batch_input_offset + input_h_idx * param->W * sizeof(float));// 参数省略，但需要注意与input_h_size保持一致，padding在首尾和输入特征图的中间是不同的，需要谨慎考虑
conv2d(input_addr, kernel_addr, output_addr);
L2S(param->output_addr + i * batch_output_offset + output_h_idx * output_w ( sizeof(float), output_addr);//与18行相同，需要谨慎对待。
}
}

该代码逻辑也可以改写为PingPong格式，逻辑上与按N切分相似，可自行推导。

Depthwise2D

Depthwise2D的大体逻辑与Conv2D类似，但有一个部分可以尝试改写Depthwise逻辑，如下：

.N = 4, .C = 3, .H = 224, .W = 224, .kernel_h = 11, .kernel_w = 11
.N = 4, .C = 3, .H = 256, .W = 256, .kernel_h = 7, .kernel_w = 7
.N = 4, .C = 3, .H = 640, .W = 640, .kernel_h = 3, .kernel_w = 3,
.N = 4, .C = 96, .H = 150, .W = 150, .kernel_h = 3, .kernel_w = 3,
...

以上是不完整的本次竞赛测试用例，第四个case需要按照常规的方法去完成。但观察第一到第三个case可以发现，如果按照N来切分，那么每次由于BM1684按照C维度分配local memory，就会造成大量的NPU被浪费（一共64块，但是只有三块被使用）。

Depthwise卷积由于特点是只对每个channel进行卷积，如图所示（一个batch，通道数为3）：

尝试同时对两个Batch进行Depthwise卷积：

好像也没有问题，只是这里卷积核需要复制一份，那么回到测试用例1-3来说，可以写出下面的代码：

//搬运输入
S2L(local_inbatch1_addr, global_inbatch1_addr); // 搬运第一个batch到0-2号NPU
S2L(local_inbatch2_addr, global_inbatch2_addr); // 搬运第二个batch到3-5号NPU，相隔3个LOCAL_MEM_SIZE距离
S2L(local_inbatch3_addr, global_inbatch3_addr); // 搬运第三个batch到6-8号NPU
S2L(local_inbatch4_addr, global_inbatch4_addr); // 搬运第四个batch到9-11号NPU
//搬运卷积核
S2L(local_kernel1_addr, global_kernel_addr); // 搬运到0-2号NPU
S2L(local_kernel2_addr, global_kernel_addr); // 搬运到3-5号NPU
S2L(local_kernel3_addr, global_kernel_addr); // 搬运到6-8号NPU
S2L(local_kernel4_addr, global_kernel_addr); // 搬运到9-11号NPU
//搬运输出
L2S(global_outbatch1_addr, local_outbatch1_addr); // 从0-2号NPU搬回
L2S(global_outbatch2_addr, local_outbatch2_addr); // 从3-5号NPU搬回
L2S(global_outbatch3_addr, local_outbatch3_addr); // 从6-8号NPU搬回
L2S(global_outbatch4_addr, local_outbatch4_addr); // 从9-11号NPU搬回

就可以一次完成推理，但实际上N*C的数据可以视为1*NC的数据，因此上面的S2L和L2S也可以优化为：

//搬运输入
S2L(local_inbatch_addr, global_inbatch_addr); // 搬运一个抽象的通道数为N*C的batch到0-11号NPU
//搬运卷积核
S2L(local_kernel1_addr, global_kernel_addr); // 搬运到0-2号NPU
S2L(local_kernel2_addr, global_kernel_addr); // 搬运到3-5号NPU
S2L(local_kernel3_addr, global_kernel_addr); // 搬运到6-8号NPU
S2L(local_kernel4_addr, global_kernel_addr); // 搬运到9-11号NPU
//搬运输出
L2S(global_outbatch_addr, local_outbatch_addr); // 从0-11号NPU搬回

以上就是Conv2D和Depthwise2D的一些分析思路和伪代码，由于本人在这方面研究并不深入，在性能上难以优化到最佳，希望可以作为一个简单的示例帮助参赛选手快速完成一个算子原型的实现。关于利用PingPong完成加速只提供了一个示例，参赛选手可以在非PingPong的代码基础上完成PingPong代码的实现，需要注意的是要处理好边界情况。