使用 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 的延迟说明。

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