NEON加速im2col+gemm的卷积推理

比眉伴天荒 2022-12-10 14:45 292阅读 0赞

# 前言 #

主流框架包括caffe、tf、darknet实现卷积的时候不是直接进行卷积计算，而是往往采用im2col+gemm的方式来进行加速，加速的思路是“空间换时间”提高读取数据过程cache的命中。本文主要利用ARM的NEON汇编指令来加速darknet框架的CPU卷积计算

# darknet相关 #

## 卷积实现原理 ##

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hzcXlj_size_5_color_FFFFFF_t_70_pic_center]  
[具体的im2col原理][im2col]  
卷积实现时，卷积核第一个元素与经过im2col操作后的feature map的第一行对应的的每个元素先后分别相乘，结果值分别存到不同位置，如图中所示。然后以此类推，卷积核第二个元素和第二行的每个元素相乘，结果加到刚才的对应位置上。操作完第一个卷积核后再继续操作第二个卷积核，直到所有卷积核都被计算完  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hzcXlj_size_16_color_FFFFFF_t_70_pic_center]

## 数据内存排布 ##

*  输入图像或者上一层卷积输出feature map的内存排布  
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hzcXlj_size_16_color_FFFFFF_t_70_pic_center 1]  
    即nchw排布，内存连续分布，先放第一行，再放下一行，放完整个通道然后放下一个通道以此类推
 *  im2col后的图像内存排布  
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hzcXlj_size_16_color_FFFFFF_t_70_pic_center 2]  
    每个卷积块的第一个元素先排在一起，然后是每个卷积块的第二个元素排在一起，以此类推
 *  卷积核内存排布  
    同图像内存排布

# NEON卷积对比实验 #

> darknet的矩阵乘法gemm代码:

void gemm_nn(int M, int N, int K, float ALPHA, 
            float *A, int lda, 
            float *B, int ldb,
            float *C, int ldc)
            //M=n,卷积核个数
            //N=输出feature map的w*h
            //K=k*k*c,单个卷积核的空间尺寸大小
            //ALPHA=1
            //*A=指向卷积核参数的指针
            //lda=K=k*k*c,单个卷积核的空间尺寸大小
            //*B=经过im2col的上一层feature map
            //ldb=N=输出feature map的w*h
            //*C=指向卷积结果的指针
            //ldc=N=输出feature map的w*h
    { 
        int i,j,k;
        #pragma omp parallel for
        for(i = 0; i < M; ++i){ //卷积核间的循环，M=卷积核个数
            for(k = 0; k < K; ++k){ //卷积核内的循环，K=单个卷积核的空间尺寸大小
                register float A_PART = ALPHA*A[i*lda+k];
                for(j = 0; j < N; ++j){ //N=输出feature map的w*h
                    C[i*ldc+j] += A_PART*B[k*ldb+j];
                }
            }
        }
    }

> NEON的矩阵乘法gemm代码

.text
            .syntax   unified
            .align   4
            .global   gemm_nn_neon
            .thumb
            .thumb_func
    
    gemm_nn_neon:
            @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
            @
            @  gemm_nn(int M, int N, int K, 
            @           @float ALPHA, 
            @           float *A, int lda, 
            @           float *B, int ldb,
            @           float *C, int ldc)
            @  r0: M=n,卷积核个数
            @  r1: N=输出feature map的w*h
            @  r2: K=k*k*c,单个卷积核的空间尺寸大小
            @      @ALPHA=1
            @  r3: *A=指向卷积核参数的指针
            @  
            @  r5: *B=经过im2col的上一层featuremap
            @  r6: ldb=N=输出feature map的w*h
            @  r7: *C=指向卷积结果的指针
            @
            @
            @  r13: sp
            @  r14: ld
            @  r15: pc
            @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
            push	        {fp, lr}                @ 进入函数的基本操作
    	add	        fp, sp, #4
    	sub	        sp, sp, #24
    	str	        r0, [fp, #-16]
    	str	        r1, [fp, #-20]
    	str	        r2, [fp, #-24]
    	str	        r3, [fp, #-28]
    .L_n_start:
    
    .L_n_loop:
            ldr             r2, [fp, #-24]          @ 读取单个卷积核的空间尺寸大小        
    .L_c_start:
            ldr             r5, [fp, #8]            @ 读取要进行卷积的featuremap的初始地址
    .L_c_loop:
            ldr             r1, [fp, #-20]          @ 读取要进行卷积的featuremap的初始大小        
            ldr             r7, [fp, #16]           @ 读取卷积结果的初始地址    
    .L_wh_start:        
            vld1.32         d0[0], [r3]!            @ 读卷积核
    .L_wh_loop:
            and             r4, r1, #3              @ r4 = w*h % 4; calculate the residual loop
            asr             r6, r1, #2              @ r6 = w*h >> 2 = w*h / 4; calculate the main loop
            cbz             r4, .L_check_wh_mainloop
    .L_wh_residualloop:        
            vld1.32         {d2[0]}, [r7]           @ 读取要保存的地址
            vmov.f32        d1, #0.0                @ 赋值为0
            vld1.32         {d1[0]}, [r5]!          @ 读要进行卷积的featuremap        
            vmla.f32        d2, d1, d0[0]           @ q2 += q1 * d0[0]
            vst1.32         {d2[0]}, [r7]!          @ 将结果从寄存器保存到内存中
    
            subs            r4, r4, #1              @ 需要进行的residual循环次数减一
            bgt             .L_wh_residualloop      @ 如果上一句subs没有更改状态寄存器的大于位，即r4-1仍然大于0,则继续循环
    .L_check_wh_mainloop:
            cbz             r6, .L_check_c_loop
    .L_wh_mainloop:
            vld1.32         {q2}, [r7]              @ 读取要保存的地址
            vld1.32         {q1}, [r5]!             @ 读要进行卷积的featuremap
            vmla.f32        q2, q1, d0[0]           @ q2 += q1 * d0[0]
            vst1.32         {q2}, [r7]!             @ 将结果从寄存器保存到内存中
    
            subs            r6, r6, #1              @ 需进行的循环数减1
            bgt             .L_wh_mainloop          @ 如果上一句subs没有更改状态寄存器的大于位，即r6-1仍然大于0，则继续循环
    .L_check_c_loop:
            subs            r2, r2, #1              @ 需要进行的c循环次数减一
            bgt             .L_c_loop               @ 如果上一句subs没有更改状态寄存器的大于位，即r2-1仍然大于0,则继续循环
    .L_c_return:
    
    .L_check_n_loop:
            str             r7, [fp, #16]           @ 将更新的卷积结果地址保存
            subs            r0, r0, #1              @ 需要进行的n循环次数减一        
            bgt             .L_n_loop               @ 如果上一句subs没有更改状态寄存器的大于位，即r0-1仍然大于0，则继续循环                
    .L_n_return:
    
    .L_return:
            mov             r0, #0
            add	        sp, sp, #24             @ sp移到进入该函数的初始位置
            pop	        {fp, pc}

在ARM A53的树莓派3B+上进行对比实验，跑一个yolo v3 tiny模型，不开openmp的单核推理一张416\*416的图片平均需要30s左右，用了neon的单核推理平均只需要12s左右，速度为原来的约2.5倍。但是为什么没有达到理论上的4倍速度提升呢，也许是因为一些非卷积层操作没有因为NEON得到提升？下次把每个层花费的时间打印出来看看，待续。

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hzcXlj_size_5_color_FFFFFF_t_70_pic_center]: /images/20221123/2f955bb9455b43e69e29626ddf7108ea.png
[im2col]: https://blog.csdn.net/Mrhiuser/article/details/52672824
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hzcXlj_size_16_color_FFFFFF_t_70_pic_center]: /images/20221123/bc163f34b07e442497ee1c5845d1052a.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hzcXlj_size_16_color_FFFFFF_t_70_pic_center 1]: /images/20221123/1a6d38ad5ac247d8b4cdcc0d147516b0.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2hzcXlj_size_16_color_FFFFFF_t_70_pic_center 2]: /images/20221123/0a21d9c82a304190adf2f8c8c7ee0131.png