79 lines
2.2 KiB
C++
79 lines
2.2 KiB
C++
|
#include "mat_mul.h"
|
||
|
|
|
||
|
|
#include <cstdlib>
|
||
|
|
#include <cstdio>
|
||
|
|
#include <pthread.h>
|
||
|
|
#include <immintrin.h>
|
||
|
|
|
||
|
|
static float *A, *B, *C;
|
||
|
|
static int M, N, K;
|
||
|
|
static int num_threads;
|
||
|
|
#define BLOCK_SIZE (48)
|
||
|
|
#define VEC_NUM (1)
|
||
|
|
#define VEC_SIZE (64)
|
||
|
|
#define VECDATA_SIZE (VEC_SIZE>>2)
|
||
|
|
#define MIN(a,b) ((a>b)?b:a)
|
||
|
|
|
||
|
|
static void* mat_mul_thread(void *data) {
|
||
|
|
// TODO: parallelize & optimize matrix multiplication
|
||
|
|
int i,j,k;
|
||
|
|
int kk;
|
||
|
|
int pid = * (int *) data;
|
||
|
|
|
||
|
|
int slice = (M + num_threads - 1) / num_threads;
|
||
|
|
int si = pid * slice;
|
||
|
|
int ei = (pid + 1) * slice;
|
||
|
|
if (si >= M) return NULL;
|
||
|
|
if (ei >= M) ei = M;
|
||
|
|
|
||
|
|
for (kk = 0; kk < K; kk+=BLOCK_SIZE) {
|
||
|
|
for (i = si; i < ei; i+=1) {
|
||
|
|
for (k = kk; k < MIN(K,kk+BLOCK_SIZE); ++k) {
|
||
|
|
float Aik = A[i * K + k];
|
||
|
|
unsigned long spC, epC;
|
||
|
|
__m512 va = _mm512_set1_ps(Aik);
|
||
|
|
__m512 vb[VEC_NUM],vc[VEC_NUM];
|
||
|
|
spC = ((unsigned long)&C[i*N] + (VEC_SIZE*VEC_NUM-1)) & (VEC_SIZE*VEC_NUM-1);
|
||
|
|
epC = (unsigned long)&C[(i+1)*N] & (VEC_SIZE*VEC_NUM-1);
|
||
|
|
|
||
|
|
for (j = 0; (unsigned long)&C[i * N + j] < spC; j=j+1) {
|
||
|
|
C[i * N + j] += Aik * B[k * N + j];
|
||
|
|
}
|
||
|
|
for (; (unsigned long)&C[i * N + j] < epC; j+=VECDATA_SIZE*VEC_NUM) {
|
||
|
|
vb[0] = _mm512_load_ps(&B[k*N+j]);
|
||
|
|
vc[0] = _mm512_load_ps(&C[i*N+j]);
|
||
|
|
vc[0] = _mm512_fmadd_ps(va,vb[0],vc[0]);
|
||
|
|
_mm512_store_ps(&C[i*N+j],vc[0]);
|
||
|
|
/*
|
||
|
|
vb[1] = _mm512_load_ps(&B[k*N+j+VECDATA_SIZE]);
|
||
|
|
vc[1] = _mm512_load_ps(&C[i*N+j+VECDATA_SIZE]);
|
||
|
|
vc[1] = _mm512_fmadd_ps(va,vb[1],vc[1]);
|
||
|
|
_mm512_store_ps(&C[i*N+j+VECDATA_SIZE],vc[1]);*/
|
||
|
|
}
|
||
|
|
for (; j < N; j=j+1) {
|
||
|
|
C[i * N + j] += Aik * B[k * N + j];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return NULL;
|
||
|
|
}
|
||
|
|
|
||
|
|
void mat_mul(float *_A, float *_B, float *_C, int _M, int _N, int _K, int _num_threads) {
|
||
|
|
A = _A, B = _B, C = _C;
|
||
|
|
M = _M, N = _N, K = _K;
|
||
|
|
num_threads = _num_threads;
|
||
|
|
|
||
|
|
// TODO: create '_num_threads' pthreads
|
||
|
|
pthread_t thread[_num_threads];
|
||
|
|
int tid[_num_threads];
|
||
|
|
for (int i = 0; i < _num_threads; i++) {
|
||
|
|
tid[i] = i;
|
||
|
|
pthread_create(&thread[i], NULL, mat_mul_thread, &tid[i]);
|
||
|
|
}
|
||
|
|
|
||
|
|
for (int i = 0; i < _num_threads; i++)
|
||
|
|
pthread_join(thread[i], NULL);
|
||
|
|
}
|