c - 在 C 中使用 AVX 实现矩阵运算

Question

我正在尝试使用 AVX 实现以下操作:

for (i=0; i<N; i++) {
  for(j=0; j<N; j++) {
    for (k=0; k<K; k++) {
      d[i][j] += 2 * a[i][k] * ( b[k][j]- c[k]);
    }
  }
}

for (int i=0; i<N; i++){
   f+= d[ind[i]][ind[i]]/2;
}    

其中 d 是 NxN 矩阵，a 是 NxK，b 是 KxN，c 是长度为 K 的 vector 。它们都是 double 。当然，所有数据都是对齐的，我正在使用 #pragma vector aligned 来帮助编译器 (gcc)。

我知道如何对一维数组使用 AVX 扩展，但对我来说用矩阵来做有点棘手。目前，我有以下内容，但没有得到正确的结果:

    for (int i=0; i< floor (N/4); i++){
        for (int j=0; j< floor (N/4); j++){
            __m256d D, A, B, C;
            D = _mm256_setzero_pd();
            #pragma vector aligned
            for (int k=0; k<K_MAX; k++){
                A = _mm256_load_pd(a[i] + k*4);
                B = _mm256_load_pd(b[k] + j*4);
                C = _mm256_load_pd(c + 4*k);
                B = _mm256_sub_pd(B, C);
                A = _mm256_mul_pd(A, B);
                D = _mm256_add_pd(_mm256_set1_pd(2.0), A);
                _mm256_store_pd(d[i] + j*4, D);
            }

        }
    }


    for (int i=0; i<N; i++){
        f+= d[ind[i]][ind[i]]/2;
    }    

希望有人能告诉我错在哪里

提前致谢。

注意:OpenMP我不太愿意介绍，只是用SIMD Intel指令

Answer 1

假设 N 和 K 数字都相对较大(比硬件 vector 大小 4 大得多)，这是一种矢量化主循环的方法。未经测试。

主要思想是矢量化中间循环而不是内部循环。这样做有两个原因。

1. 这避免了横向操作。当仅对内部循环进行向量化时，我们将不得不计算 vector 的水平和。

2. 当加载 4 个连续的 k 值时，b[k][j] 加载具有不幸的 RAM 访问模式，需要 4 个单独的加载指令，或者收集负载，这两种方法都比较慢。为 4 个连续的 j 值加载元素是一个全 vector 加载指令，非常高效，尤其是当您对齐输入时。

    const int N_aligned = ( N / 4 ) * 4;
    for( int i = 0; i < N; i++ )
    {
        int j = 0;
        for( ; j < N_aligned; j += 4 )
        {
            // Load 4 scalars from d
            __m256d dv = _mm256_loadu_pd( &d[ i ][ j ] );

            // Run the inner loop which only loads from RAM but never stores any data
            for( int k = 0; k < K; k++ )
            {
                __m256d av = _mm256_broadcast_sd( &a[ i ][ k ] );
                __m256d bv = _mm256_loadu_pd( &b[ k ][ j ] );
                __m256d cv = _mm256_broadcast_sd( &c[ k ] );

                // dv += 2*av*( bv - cv )
                __m256d t1 = _mm256_add_pd( av, av );   // 2*av
                __m256d t2 = _mm256_sub_pd( bv, cv );   // bv - cv
                dv = _mm256_fmadd_pd( t1, t2, dv );
            }
            // Store the updated 4 values
            _mm256_storeu_pd( &d[ i ][ j ], dv );
        }

        // Handle remainder with scalar code
        for( ; j < N; j++ )
        {
            double ds = d[ i ][ j ];
            for( int k = 0; k < K; k++ )
                ds += 2 * a[ i ][ k ] * ( b[ k ][ j ] - c[ k ] );
            d[ i ][ j ] = ds;
        }
    }

如果您想进一步优化，请尝试将内部循环展开一个小因子，例如 2，使用 2 个用 _mm256_setzero_pd() 初始化的独立累加器，将它们添加到循环之后。可能是在某些处理器上，此版本在 FMA 指令的延迟上停止，而不是使加载端口或 ALU 饱和。多个独立的累加器有时会有所帮助。