AVX优化矩阵乘法

1.看了OpenBlas的矩阵乘法优化 尝试写下AVX版本的矩阵优化

2.在单线程情况下 单精度 1000X1000的矩阵乘1000X1000 运行100次 取平均值

Pytorch-Mkl 23.ms
Numpy(应该也是用的MKL) 23.ms
我的 38.ms

以下贴上代码

void addDot8x8Pack(int k, float *A, float *B, float *C, int n) {
	int p;

	float *bp0_pntr = B;

	__m256 Zero1 = _mm256_setzero_ps();
	__m256 Zero2 = _mm256_setzero_ps();
	__m256 Zero3 = _mm256_setzero_ps();
	__m256 Zero4 = _mm256_setzero_ps();
	__m256 Zero5 = _mm256_setzero_ps();
	__m256 Zero6 = _mm256_setzero_ps();
	__m256 Zero7 = _mm256_setzero_ps();
	__m256 Zero8 = _mm256_setzero_ps();



	for (p = 0; p < k; p++) {

		__m256 a0 = _mm256_set1_ps(*(A + p));
		__m256 a1 = _mm256_set1_ps(*(A + n + p));
		__m256 a2 = _mm256_set1_ps(*(A + 2 * n + p));
		__m256 a3 = _mm256_set1_ps(*(A + 3 * n + p));
		__m256 a4 = _mm256_set1_ps(*(A + 4 * n + p));
		__m256 a5 = _mm256_set1_ps(*(A + 5 * n + p));
		__m256 a6 = _mm256_set1_ps(*(A + 6 * n + p));
		__m256 a7 = _mm256_set1_ps(*(A + 7 * n + p));

		__m256 bp = _mm256_load_ps(bp0_pntr);

		Zero1 = _mm256_fmadd_ps(a0, bp, Zero1);
		Zero2 = _mm256_fmadd_ps(a1, bp, Zero2);
		Zero3 = _mm256_fmadd_ps(a2, bp, Zero3);
		Zero4 = _mm256_fmadd_ps(a3, bp, Zero4);
		Zero5 = _mm256_fmadd_ps(a4, bp, Zero5);
		Zero6 = _mm256_fmadd_ps(a5, bp, Zero6);
		Zero7 = _mm256_fmadd_ps(a6, bp, Zero7);
		Zero8 = _mm256_fmadd_ps(a7, bp, Zero8);


		bp0_pntr += 8;

	}

	_mm256_store_ps(C, Zero1);
	_mm256_store_ps(C + n, Zero2);
	_mm256_store_ps(C + 2 * n, Zero3);
	_mm256_store_ps(C + 3 * n, Zero4);
	_mm256_store_ps(C + 4 * n, Zero5);
	_mm256_store_ps(C + 5 * n, Zero6);
	_mm256_store_ps(C + 6 * n, Zero7);
	_mm256_store_ps(C + 7 * n, Zero8);

}

void PackedBMatrix(int j, int n, int k, float *input, float *output) {

	for (int i = 0; i < k; i++) {
		
		memcpy(output + i * 8, input + i * n+j, sizeof(float) * 8);
	}

}

void testBlas3(float *A, float *B, float *C, int m, int n, int k) {
	int i, j;
	float *PackB = new float[8 * k];
	//float *PackA = new float[8 * k];
	for (j = 0; j < n; j += 8) {
		PackedBMatrix(j, n, k, B, PackB);
		for (i = 0; i < m; i += 8) {		
			//addDot8x8R(k, A + i * k, B + j, C + i * n + j, n);
			addDot8x8Pack(k, A + i * k, PackB, C + i * n + j, n);
		}
	}
	delete[]PackB;
}

#include <iostream>
int main(){
	float *A = new float[1000 * 1000];
	float *B = new float[1000 * 1000];

	for (int i = 0; i < 1000 * 1000; i++) {
		A[i] = 1.0f;
		B[i] = 1.0f;
	}
	clock_t v2,v3;
	v2 = clock();
	float *C = nullptr;
	for (int loop = 0; loop < 100; loop++) {
		C = new float[1000 * 1000];
		testBlas3(A, B, C, 1000,1000,1000);
		delete[] C;
	}	
	v3 = clock();
	cout << (v3 - v2)/100.0 << endl;

}

采取8*8分块方式
具体原理讲解参考

OpenBLAS项目与矩阵乘法优化 | AI 研习社
OpenBlas矩阵乘法优化