稀疏矩阵加法_使用avx2加速稀疏矩阵推理

原始代码是循环累加和，计算效率较低，运行速度慢，可以考虑使用avx2指令集，实现计算加速。

原始计算过程如下：weight为稀疏矩阵，内存不连续，存储方式为value-index(coo、crs、csc等)，

稀疏矩阵的存储​www.jianshu.com

SPARSE_GROUP_SIZE表示模型进行稀疏训练时，稀疏块的大小，可以设为4/8/16等，从weight中每次取SPARSE_GROUP_SIZE大小的数据与该层的输入进行乘法再累加，将累积和作为该层的输出，所有参数都为float32类型。

int weightPos = 0;
const float * __restrict x_ptr = x.data();
float * __restrict y_ptr = y.data();
for(int i=0; i<nGroups; ++i)
{
    float sum = 0;
    int col = SPARSE_GROUP_SIZE*colIdx[i];
    //scalar product loop. compiler should optimize it.
    for(int k=0; k<SPARSE_GROUP_SIZE; ++k)
    {
        sum += weight[weightPos++]*x_ptr[col++];
    }
    y_ptr[rowIdx[i]] += sum;
}

avx2中，一个寄存器可以处理256bit的数据，即8个float32的数据。

_mm_cvtss_f32​docs.microsoft.com

_mm256_mul_ps表示对两个float类型的向量进行相乘；

_mm256_hadd_ps表示水平方向上对两个float类型向量做加法；

_mm256_extractf128_ps表示从256中取128的数据，即4个f32；

_mm_add_ss将两个向量相加；

_mm_cvtss_f32从128的寄存器中取第一个f32的数据；

//SPARSE_GROUP_SIZE = 8;
for(int i=0; i<nGroups; ++i)
{
    float sum2 = 0;
    int col = SPARSE_GROUP_SIZE*colIdx[i];
    
    __m256 res0 = _mm256_mul_ps(*(__m256*)(weight + weightPos),*(__m256*)(x_ptr + col));

    res0 = _mm256_hadd_ps(res0, res0);
    res0 = _mm256_hadd_ps(res0, res0);
    __m128 acc1 = _mm256_extractf128_ps(res0, 0);
    __m128 acc2 = _mm256_extractf128_ps(res0, 1);
    acc1 = _mm_add_ss(acc1, acc2);
    sum2 = _mm_cvtss_f32(acc1);

    y_ptr[rowIdx[i]] += sum2;
    weightPos += 8;
}

当稀疏度达到90%以上，将原始稀疏矩阵推理换成avx2后，模型的整体推理速度有了大幅度的提高。

Intel 内部指令 --- AVX和AVX2学习笔记​blog.csdn.net