TPU编程竞赛

《第一届TPU编程竞赛》代码思路分享-进阶

sophgo logo
算能SOPHGO
智算赋能数字世界
作者：南京大学-王子翊 杨雪松
发布日期： 2022-04-19 12:00

乒乓缓冲

按照正常流程编程的话，比如按N切分，硬件所执行的流程为：搬入数据->计算数据->搬出数据->搬入数据->计算数据->....->搬出数据。而数据的搬运和数据的计算分别由GDMA和BDC控制，因此，当执行GDMA时，BDC为闲置状态，执行BDC时，GDMA为闲置状态。

因此，一个简单的想法是利用乒乓缓冲，达到同时使用GDMA和BDC的目的。

以上图为例：

(1) 搬入第一块的输入。

(2)计算第一块，同时搬入第二块的输入。

(3)搬出第一块的输出，搬入第一块的输入，同时计算第二块。

(4)搬出第二块的输出，搬入第二块的输入，计算第一块，之后重复(3)(4)，直到运行结束，可以看到，除了开始和结束的搬运，之后所有时间内，GDMA和BDC都是同时工作的。

伪代码以 output=input+1.0 为例，通过奇偶性来控制访问的地址：

local_addr_t input_addr[2], output_addr[2];

for(int i=0; i<blocks+2; i++)
{
okk_parallel_start(); //标志开始并行
if(i<blocks)
S2L(input_addr[i%2], param->input_addr + i*input_skip_bytes); // 搬运数据

if(i>0 && i<blocks+1; i++)
add_one(output_addr[(i-1)%2], input[(i-1)%2]); // 计算上次搬运的数据

if(i>1)
L2S(param->output_addr + (i-2)*output_skip_bytes, output_addr[(i-2)%2]); //搬出上次计算好的数据
okk_parallel_end();
}

包含在okk_parallel_start()和okk_parallel_end()中的操作会被并行执行。更详细的代码见算丰文档。

赛题分析

《第一届TPU编程竞赛》代码思路分享-入门
《第一届TPU编程竞赛》代码思路分享-进阶
《第一届TPU编程竞赛》代码思路分享-Softmax
《第一届TPU编程竞赛》代码思路分享-Matmul
《第一届TPU编程竞赛》代码思路分享-Conv2D和Depthwise2D