当矩阵维数不是4的倍数时，如何避免AVX2的错误？

Posted 2023-02-16

【中文标题】当矩阵维数不是4的倍数时，如何避免AVX2的错误？

【英文标题】：How to avoid the error of AVX2 when the matrix dimension isn't multiples of 4?

【发布时间】：2018-07-22 06:25:25

【问题描述】：

我在 C 中使用 AVX2、FMA 制作了矩阵向量乘法程序。我使用带有 -mfma、-mavx 的 GCC ver7 编译。

但是，我收到错误“已释放对象的校验和不正确 - 对象可能在被释放后被修改。”

我认为如果矩阵维度不是 4 的倍数，则会产生错误。

我知道 AVX2 使用 ymm 寄存器，可以使用 4 个双精度浮点数。因此，如果矩阵是 4 的倍数，我可以毫无错误地使用 AVX2。

但是，这是我的问题。 如果矩阵不是 4 的倍数，我如何有效地使用 AVX2 ???

这是我的代码。

#include "stdio.h"
#include "math.h"
#include "stdlib.h"
#include "time.h"
#include "x86intrin.h"

void mv(double *a,double *b,double *c, int m, int n, int l)

    __m256d va,vb,vc;
    int k;
    int i;
    for (k = 0; k < l; k++) 
        vb = _mm256_broadcast_sd(&b[k]);
        for (i = 0; i < m; i+=4) 
            va = _mm256_loadu_pd(&a[m*k+i]);
            vc = _mm256_loadu_pd(&c[i]);
            vc = _mm256_fmadd_pd(vc, va, vb);
            _mm256_storeu_pd( &c[i], vc );
        
    

int main(int argc, char* argv[]) 

    // set variables
    int m;
    double* a;
    double* b;
    double* c;
    int i;
    int temp=0;
    struct timespec startTime, endTime;

    m=9;
    // main program

    // set vector or matrix
    a=(double *)malloc(sizeof(double) * m*m);
    b=(double *)malloc(sizeof(double) * m*1);
    c=(double *)malloc(sizeof(double) * m*1);

    for (i=0;i<m;i++) 
        a[i]=1;
        b[i]=1;
        c[i]=0.0;
    
    for (i=m;i<m*m;i++) 
        a[i]=1;
    

    // check start time
    clock_gettime(CLOCK_REALTIME, &startTime);
    mv(a, b, c, m, 1, m);
    // check end time
    clock_gettime(CLOCK_REALTIME, &endTime);

    free(a);
    free(b);
    free(c);
    return 0;

【参考方案1】：

您加载和存储 4 个 double 的向量，但您的循环条件仅检查 first 向量元素是否在边界内，因此您最多可以写入 3x8 = 24 个字节的外部对象当m 不是 4 的倍数时。

您需要在主循环中使用 i < (m-3) 之类的内容，以及用于处理最后部分数据向量的清理策略。使用 SIMD 进行向量化非常类似于展开：您必须检查在循环条件中执行多个未来元素是否可以。

标量清理循环运行良好，但我们可以做得更好。例如，在最后一个完整的 256 位向量（即最多 1 个）之后，在进行标量之前，尽可能多地执行 128 位向量。

在许多情况下（例如只写目标），在数组末尾结束的未对齐向量加载非常好（当m>=4 时）。如果m%4 != 0，它可以与您的主循环重叠，但这很好，因为您的输出数组不会与您的输入重叠，因此作为单个清理的一部分重做一个元素比避免它的分支更便宜。

但这在这里不起作用，因为你的逻辑是c[i+0..3] += ...，所以重做一个元素会出错。

// cleanup using a 128-bit FMA, then scalar if there's an odd element.
// untested

void mv(double *a,double *b,double *c, int m, int n, int l)

   /*  the loop below should actually work for m=1..3, but a separate strategy might be good.
    if (m < 4) 
        // maybe check m >= 2 and use __m128 vectors?
        // or vectorize differently?
    
   */


    for (int k = 0; k < l; k++) 
        __m256 vb = _mm256_broadcast_sd(&b[k]);
        int i;
        for (i = 0; i < (m-3); i+=4) 
            __m256d va = _mm256_loadu_pd(&a[m*k+i]);
            __m256d vc = _mm256_loadu_pd(&c[i]);
                    vc = _mm256_fmadd_pd(vc, va, vb);
            _mm256_storeu_pd( &c[i], vc );
        
        if (i<(m-1)) 
            __m128d lasta = _mm_loadu_pd(&a[m*k+i]);
            __m128d lastc = _mm_loadu_pd(&c[i]);
                    lastc = _mm_fmadd_pd(lastc, va, _mm256_castpd256_pd128(vb));
                _mm_storeu_pd( &c[i], lastc );
            // i+=2;  // last element only checks m odd/even, doesn't use i
        
        // if (i<m)
        if (m&1) 
            // odd number of elements, do the last non-vector one
            c[m-1] += a[m*k + m-1] * _mm256_cvtsd_f64(vb);
        

    

我还没有详细了解 gcc/clang -O3 是如何编译它的。有时编译器会尝试在清理代码方面变得过于聪明（例如，尝试自动矢量化标量清理循环）。

其他策略可能包括使用 AVX 掩码存储执行最后最多 4 个元素：您需要在每个矩阵行的末尾使用相同的掩码，因此生成一次然后在每行的末尾使用它可以做个好人。见Vectorizing with unaligned buffers: using VMASKMOVPS: generating a mask from a misalignment count? Or not using that insn at all。 （为了简化分支，您需要将其设置为使您的主循环仅转到 i < (m-4)，然后您始终运行清理。在 m%4 == 0 的情况下，掩码是全一所以你做最后的完整向量。）如果你不能安全地读取矩阵的末尾，你可能需要一个掩码加载以及掩码存储。

---

您还可以考虑对齐行以提高效率，或者考虑与行的逻辑长度分开的行跨度。 （即将行填充到 32 字节边界）。在行尾保留填充可以简化清理工作，因为您始终可以使用写入填充的整个向量。

---

特殊情况m==2：您希望将 2 个元素广播到 __m256d 的两个 128 位通道中，而不是从 b[] 广播一个元素，所以一个 256- bit FMA 可以一次做 2 行。