使用 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

使用 AVX 的全连接层（点积）

Fully Connected Layer (dot product) using AVX

c++

performance

avx

dot-product

openvino