bert深度优化

背景

在当下cv, nlp领域，一般都会使用attention结构，对attention的优化就显得额外重要，这边文章就nvidia对attention的优化策略进行分析。

相关知识

PTX指令

因为这里我们需要利用到一些PTX指令，所以简单做一点介绍：

        asm volatile( \
            "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 \n" \
            "    {%0, %1, %2, %3}, \n" \
            "    {%4, %5, %6, %7}, \n" \
            "    {%8, %9}, \n" \
            "    {%0, %1, %2, %3}; \n" \
                    : "+f"(  elt(0)), "+f"(  elt(1)), "+f"(  elt(2)), "+f"(  elt(3))
                    :  "r"(a.reg(0)),  "r"(a.reg(1)),  "r"(a.reg(2)),  "r"(a.reg(3))
                    ,  "r"(b.reg(0)),  "r"(b.reg(1)));

上面是一段矩阵乘加的ptx代码，计算公式是D(16x8) = A(16x16) * B(16x8) + C(16*8)。具体解释如下：（一个寄存器可以存一个fp32）

mma.sync.aligned.m8n8k4.alayout.blayout.dtype.f16.f16.ctype d, a, b, c;

• 矩阵A，数据类型是.fp16, 一个向量表达式中包含4个寄存器，也就是4个fp16 *2
• 矩阵B，数据类型是.fp16, 一个向量表达式中包含2个寄存器，也就是2个fp16 *2
• 矩阵C，数据类型是.fp32, 一个向量表达式中包含4个寄存器，也就是4个fp32
• 矩阵D，数据类型是.fp32, 一个向量表达式中包含4个寄存器，也就是4个fp32
  nvidia的文档链接：Matrix Shape, PTX ISA reference