OpenMP+AVX加速矩阵运算

1 调整计算顺序

一般的方阵乘法的实现：

void MatrixMul(int **A, int **B, int **C, int size) {for (int i = 0; i < size; i++)for (int j = 0; j < size; j++)for (int k = 0; k < size; k++)C[i][j] += A[i][k] * B[k][j];
}

根据线性代数的知识容易知道，C[i][j] 的计算结果和 ijk 的循环顺序没有关系。仔细观察上面的代码可以发现，最内层循环的循环变量是 k，那么此时 A 是逐行访问的，而 B 则是逐列访问的。根据高维数组的存储方式不难知道，在这样的情况下，矩阵 A 是的访问在内存上是连续的，而 B 的访问在内存上是不连续的。

为了方便使用 AVX 并行计算，以及访问 cache，可以调整循环下标的顺序为 ikj，这样矩阵 A 和 B 的访问在内存上都是连续的了，考虑到二维数组和一维数组的内存布局是一样的门卫了方便，可以使用一维数组来处理：

void conventional_MatMul(float *A, float *B, float *C, int size) {for (int i = 0; i < size; i++) {for (int j = 0; j < size; j++) {for (int k = 0; k < size; k++) {*(C + i * size + j) += *(A + i * size + k) * *(B + k * size + j);}}}
}

2 AVX 优化

对于最内存层循环，

for (int j = 0; j < size; j++)*(C + i * size + j) += *(A + i * size + k) * *(B + k * size + j);

计算分为两步进行： *(A + i * size + k) 与 B 的第 k 行相乘；相乘的结果与 C 的第 i 行相加。

2.1 乘法优化

考虑到最内层循环里面，*(A + i * size + k)是一个固定的数，与 B 的每一行中的每一个数相乘，因此存放 *(A + i * size + k) 的寄存器的每个单元都要写入相同的*(A + i * size + k)，可以使用 _mm512_set1_ps方法进行复制型的加载。

__m512 ra = _mm512_set1_ps(*(A + i * size + k));

B 的加载就想对比较方便，现在 avx 的内存对齐和不对齐的性能差不多，就不做内存对齐，使用 _mm512_loadu_ps：

_mm512_loadu_ps(B + k * size + j)；

乘法方面注意要使用 _mm512_mul_ps 来得到每一位相乘的结果。

rb = _mm512_mul_ps(ra, rb);

2.2 加法优化

加法和乘法类似，先加载矩阵 C：

_mm512_loadu_ps(C + k * size + j)；

和乘法的结果相加。

rc = _mm512_add_ps(rb, rc);

最后写回C，

_mm512_storeu_ps(C + i * size + j, rc);

2.3 循环步长

使用 AVX512，一次可以处理 512 位，即 16 个 float32，因此循环步长设置为 16。

#include <iostream>
#include <immintrin.h>
#include <cstdlib> 
#include <ctime>void gen_matrix(float *m, int size) {srand (static_cast <unsigned> (time(0)));for (int i = 0; i < size; i++) {for (int j= 0; j < size; j++)*(m + i * size + j) = static_cast <float> (rand()) / static_cast <float> (RAND_MAX);}
}void conventional_MatMul(float *A, float *B, float *C, int size) {for (int i = 0; i < size; i++) {for (int j = 0; j < size; j++) {for (int k = 0; k < size; k++) {*(C + i * size + j) += *(A + i * size + k) * *(B + k * size + j);}}}
}void avx512_MatMul(float *A, float *B, float *C, int size) {for (int i = 0; i < size; i++) {for (int k = 0; k < size; k++) {int j = 0;__m512 ra = _mm512_set1_ps(*(A + i * size + k));for (; j <= size; j += 16) {__m512 rc;rc = _mm512_add_ps(_mm512_mul_ps(ra, _mm512_loadu_ps(B + k * size + j)), rc);_mm512_storeu_ps(C + i * size + j, rc);}}}
}void avx256_MatMul(float *A, float *B, float *C, int size)
{for (int i = 0; i < size; i++) {for (int k = 0; k < size; k++) {int j = 0;__m256 ra = _mm256_set1_ps(*(A + i * size + k));for (; j <= size; j += 8) {__m256 rc;rc = _mm256_add_ps(_mm256_mul_ps(ra, _mm256_loadu_ps(B + k * size + j)), rc);_mm256_storeu_ps(C + i * size + j, rc);}}}
}int main() {int size = 2000;float* A = new float[size * size]();float* B = new float[size * size]();float* C1 = new float[size * size]();float* C2 = new float[size * size]();float* C3 = new float[size * size]();gen_matrix(A, size);gen_matrix(B, size);std::cout << "conventional MatMul: " << std::endl;clock_t begin_time1 = clock();conventional_MatMul(A, B, C1, size);float seconds1 = float(clock() - begin_time1) / CLOCKS_PER_SEC;std::cout << "time: " << seconds1 << "s\n" << std::endl;std::cout <<"avx256 MatMul: " << std::endl;clock_t begin_time2 = clock();avx256_MatMul(A, B, C2, size);float seconds2 = float(clock() - begin_time2) / CLOCKS_PER_SEC;std::cout << "time: " << seconds2 << "s\n" << std::endl;std::cout <<"avx512 MatMul: " << std::endl;clock_t begin_time3 = clock();avx512_MatMul(A, B, C3, size);float seconds3 = float(clock() - begin_time3) / CLOCKS_PER_SEC;std::cout << "time: " << seconds3 << "s\n" << std::endl;delete[] A;delete[] B;delete[] C1;delete[] C2;delete[] C3;return 0;
}

编译运行，

g++ mm_avx.cpp -march=native -o mm_avx
./mm_avx# conventional MatMul:
# time: 41.7601s# avx256 MatMul:
# time: 5.04586s# avx512 MatMul:
# time: 3.0,171s

对比 2000*2000 的矩阵乘法，使用 avx256 的速度大概是常规的方式的 8 倍多，使用 avx512 的速度大概接近常规的方式的 13 倍。

3 OpenMP 优化

相较于 AVX，OpenMP 在使用方面就简单了许多，在这里只需要在最外层循环上放加上一句#pragma omp parallel for即可，在末尾加上num_threads(N)设置使用的线程数为 N。

#include <iostream>
#include <immintrin.h>
#include <cstdlib> 
#include <ctime>
#include <omp.h>
#include <chrono>void gen_matrix(float *m, int size) {srand (static_cast <unsigned> (time(0)));for (int i = 0; i < size; i++) {for (int j= 0; j < size; j++)*(m + i * size + j) = static_cast <float> (rand()) / static_cast <float> (RAND_MAX);}
}void conventional_MatMul(float *A, float *B, float *C, int size) {for (int i = 0; i < size; i++) {for (int j = 0; j < size; j++) {for (int k = 0; k < size; k++) {*(C + i * size + j) += *(A + i * size + k) * *(B + k * size + j);}}}
}void avx512_MatMul(float *A, float *B, float *C, int size, int thread_num) {#pragma omp parallel for num_threads(thread_num)for (int i = 0; i < size; i++) {for (int k = 0; k < size; k++) {int j = 0;__m512 ra = _mm512_set1_ps(*(A + i * size + k));for (; j <= size; j += 16) {__m512 rc;rc = _mm512_add_ps(_mm512_mul_ps(ra, _mm512_loadu_ps(B + k * size + j)), rc);_mm512_storeu_ps(C + i * size + j, rc);}}}
}void avx256_MatMul(float *A, float *B, float *C, int size, int thread_num) {#pragma omp parallel for num_threads(thread_num)for (int i = 0; i < size; i++) {for (int k = 0; k < size; k++) {int j = 0;__m256 ra = _mm256_set1_ps(*(A + i * size + k));for (; j <= size; j += 8) {__m256 rc;rc = _mm256_add_ps(_mm256_mul_ps(ra, _mm256_loadu_ps(B + k * size + j)), rc);_mm256_storeu_ps(C + i * size + j, rc);}}}
}int main() {int size = 2000;int thread_num = 32;float* A = new float[size * size]();float* B = new float[size * size]();float* C1 = new float[size * size]();float* C2 = new float[size * size]();float* C3 = new float[size * size]();gen_matrix(A, size);gen_matrix(B, size);std::cout << "conventional MatMul: " << std::endl;auto start1 = std::chrono::steady_clock::now();conventional_MatMul(A, B, C1, size);auto end1 = std::chrono::steady_clock::now();std::chrono::duration<double, std::micro> elapsed1 = end1 - start1;float seconds1 = float(elapsed1.count()) / 1000000;std::cout << "time: " << seconds1 << "s\n" << std::endl;std::cout <<"avx256 MatMul: " << std::endl;auto start2 = std::chrono::steady_clock::now();avx256_MatMul(A, B, C2, size, thread_num);auto end2 = std::chrono::steady_clock::now();std::chrono::duration<double, std::micro> elapsed2 = end2 - start2;float seconds2 = float(elapsed2.count()) / 1000000;std::cout << "time: " << seconds2 << "s\n" << std::endl;std::cout <<"avx512 MatMul: " << std::endl;auto start3 = std::chrono::steady_clock::now();avx512_MatMul(A, B, C3, size, thread_num);auto end3 = std::chrono::steady_clock::now();std::chrono::duration<double, std::micro> elapsed3 = end3 - start3;float seconds3 = float(elapsed3.count()) / 1000000;std::cout << "time: " << seconds3 << "s\n" << std::endl;delete[] A;delete[] B;delete[] C1;delete[] C2;delete[] C3;return 0;
}