使用AVX的全连接层(点积)

Question

我有下面的C++代码来执行完全连接层的乘法和累积步骤(没有偏倚)。基本上，我只是用向量(输入)和矩阵(权重)来做点积。我用AVX向量来加速手术。

const float* src = inputs[0]->buffer();
const float* scl = weights->buffer();
float* dst = outputs[0]->buffer();

SizeVector in_dims = inputs[0]->getTensorDesc().getDims();
SizeVector out_dims = outputs[0]->getTensorDesc().getDims();

const int in_neurons = static_cast<int>(in_dims[1]);
const int out_neurons = static_cast<int>(out_dims[1]);    

for(size_t n = 0; n < out_neurons; n++){
    float accum = 0.0;
    float temp[4] = {0,0,0,0};
    float *p = temp;

    __m128 in, ws, dp;

    for(size_t i = 0; i < in_neurons; i+=4){

        // read and save the weights correctly by applying the mask
        temp[0] = scl[(i+0)*out_neurons + n];
        temp[1] = scl[(i+1)*out_neurons + n];
        temp[2] = scl[(i+2)*out_neurons + n];
        temp[3] = scl[(i+3)*out_neurons + n];

        // load input neurons sequentially
        in = _mm_load_ps(&src[i]);

        // load weights
        ws = _mm_load_ps(p);

        // dot product
        dp = _mm_dp_ps(in, ws, 0xff);

        // accumulator
        accum += dp.m128_f32[0]; 
    }
    // save the final result
    dst[n] = accum.m128_f32[0];
}

这是可行的，但加速比远远超出了我的预期。正如您在下面看到的，使用x24的卷积层比我的自定义点产品层花费更少的时间。这是没有道理的，应该有更多的改进余地。当我尝试使用AVX时，我的主要错误是什么？(我对AVX编程很陌生，所以我不完全理解从哪里开始对代码进行全面优化)。

**Convolutional Convolutional Layer Fully Optimized (AVX)**
Layer: CONV3-32 
Input: 28x28x32 = 25K   
Weights: (3*3*32)*32 = 9K   
Number of MACs: 3*3*27*27*32*32 = 7M    
Execution Time on OpenVINO framework: 0.049 ms

**My Custom Dot Product Layer Far From Optimized (AVX)**
Layer: FC
Inputs: 1x1x512
Outputs: 576    
Weights: 3*3*64*512 = 295K  
Number of MACs: 295K    
Execution Time on OpenVINO framework: 0.197 ms

谢谢大家提前提供帮助！

Answer 1

附录:你所做的实际上是一个矩阵向量积。我们很清楚如何有效地实现这一点(尽管有了缓存和指令级的并行性，这并不是完全微不足道的事情)。答案的其余部分只显示了一个非常简单的向量化实现。

您可以通过增量n+=8和i+=1 (假设out_neurons为8的倍数，否则需要对最后一个元素执行一些特殊处理)来大大简化您的实现，也就是说，您可以同时积累8个dst值。

一个非常简单的实现假设FMA是可用的(否则使用乘法和加法)：

void dot_product(const float* src, const float* scl, float* dst,
                 const int in_neurons, const int out_neurons)
{
    for(size_t n = 0; n < out_neurons; n+=8){
        __m256 accum = _mm256_setzero_ps();

        for(size_t i = 0; i < in_neurons; i++){
            accum = _mm256_fmadd_ps(_mm256_loadu_ps(&scl[i*out_neurons+n]), _mm256_set1_ps(src[i]), accum);
        }
        // save the result
        _mm256_storeu_ps(dst+n ,accum);
    }
}

这仍然可以优化，例如，通过在内环中累积2、4或8个dst分组，这不仅可以保存一些广播操作( _mm256_set1_ps指令)，而且还可以补偿FMA指令的延迟。

如果你想玩代码：https://godbolt.org/z/mm-YHi