wehub-resource-sync
409 lines
13 KiB
Common Lisp
409 lines
13 KiB
Common Lisp
#ifdef MNN_SUPPORT_FP16
|
|
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
|
#endif
|
|
|
|
/*
|
|
\
|
|
#define OPWM 64 // The outputsize-per-workgroup in dimension M
|
|
#define OPWN 128 // The outputsize-per-workgroup in dimension N
|
|
#define CPWK 8 // The cachesize-per-workgroup in dimension K
|
|
#define OPTM 4 // The outputsize-per-thread in dimension M
|
|
#define OPTN 8 // The outputsize-per-thread in dimension N
|
|
*/
|
|
#define TPWM (OPWM/OPTM) // The threadsize-per-workgroup in dimension M
|
|
#define TPWN (OPWN/OPTN) // The threadsize-per-workgroup in dimension N
|
|
#define LPTA ((CPWK*OPWM)/(TPWM*TPWN)) // Loads-num-per-thread for A
|
|
#define LPTB ((CPWK*OPWN)/(TPWM*TPWN)) // Loads-num-per-thread for B
|
|
|
|
// vetorize + pragma unroll
|
|
__kernel void matmul_local_buf(const int M, const int N, const int K,
|
|
__global const FLOAT* A,
|
|
#if (defined USE_LOW_BIT_WEIGHT_INT8)
|
|
__global const char* B,
|
|
__global const float* dequantScale,
|
|
__global const float* dequantOffset,
|
|
#elif (defined USE_LOW_BIT_WEIGHT_INT4)
|
|
__global const uchar* B,
|
|
__global const float* dequantScale,
|
|
__global const float* dequantOffset,
|
|
#else
|
|
__global const FLOAT* B,
|
|
#endif
|
|
#ifdef BIAS
|
|
__global const FLOAT* bias,
|
|
#endif
|
|
__global FLOAT* C) {
|
|
|
|
// Local thread id
|
|
const int lidm = get_local_id(0); // Local row ID
|
|
const int lidn = get_local_id(1); // Local col ID
|
|
// group id
|
|
const int offsetM = get_group_id(0) * OPWM; // Work-group offset M
|
|
const int offsetN = get_group_id(1) * OPWN; // Work-group offset N
|
|
|
|
// Local memory for work-group cache of A and B
|
|
__local FLOAT Alocal[CPWK][OPWM];
|
|
__local FLOAT Blocal[OPWN][CPWK+2];
|
|
|
|
// Allocate register space
|
|
COMPUTE_FLOAT sum[OPTM][OPTN];
|
|
|
|
// Initialise the accumulation registers
|
|
for (int wm=0; wm<OPTM; wm++) {
|
|
for (int wn=0; wn<OPTN; wn++) {
|
|
sum[wm][wn] = 0.0f;
|
|
}
|
|
}
|
|
|
|
// Loop over all tiles
|
|
const int numLoops = K/CPWK;
|
|
int lid = lidn*TPWM + lidm;
|
|
|
|
for (int t=0; t<numLoops; t++) {
|
|
// Load one work-group of A and B into local memory
|
|
for (int la=0; la<LPTA; la++) {
|
|
int id = la*TPWN*TPWM + lid;
|
|
int row = id % OPWM;
|
|
int col = id / OPWM;
|
|
int tiledIndex = CPWK*t + col;
|
|
#ifdef TRANSPOSE_A
|
|
// [K, M]
|
|
Alocal[col][row] = A[tiledIndex*M + (offsetM + row)];
|
|
#else
|
|
// [M, K]
|
|
Alocal[col][row] = A[(offsetM + row)*K + tiledIndex];
|
|
#endif
|
|
}
|
|
|
|
for (int la=0; la<LPTB; la++) {
|
|
int id = la*TPWN*TPWM + lid;
|
|
int row = id % OPWN;
|
|
int col = id / OPWN;
|
|
int tiledIndex = CPWK*t + col;
|
|
#ifdef TRANSPOSE_B
|
|
// [N, K]
|
|
Blocal[row][col] = B[(offsetN + row)*K + tiledIndex];
|
|
#else
|
|
// [K, N]
|
|
Blocal[row][col] = B[tiledIndex*N + offsetN + row];
|
|
#endif
|
|
}
|
|
|
|
// Synchronise to make sure the tile is loaded
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
// Loop over the values of a single tile
|
|
|
|
// Perform the computation
|
|
FLOAT4 A_k0, B_k0[OPTN];
|
|
{
|
|
int row = lidm;
|
|
int col = lidn;
|
|
|
|
A_k0.s0 = Alocal[0][row];
|
|
A_k0.s1 = Alocal[1][row];
|
|
A_k0.s2 = Alocal[2][row];
|
|
A_k0.s3 = Alocal[3][row];
|
|
|
|
#pragma unroll
|
|
for (int wn=0; wn<OPTN; wn++) {
|
|
B_k0[wn].s0 = Blocal[col][0];
|
|
B_k0[wn].s1 = Blocal[col][1];
|
|
B_k0[wn].s2 = Blocal[col][2];
|
|
B_k0[wn].s3 = Blocal[col][3];
|
|
sum[0][wn] += dot(A_k0, B_k0[wn]);
|
|
col += TPWN;
|
|
}
|
|
|
|
#pragma unroll
|
|
for(int wm=1; wm<OPTM; wm++) {
|
|
row += TPWM;
|
|
A_k0.s0 = Alocal[0][row];
|
|
A_k0.s1 = Alocal[1][row];
|
|
A_k0.s2 = Alocal[2][row];
|
|
A_k0.s3 = Alocal[3][row];
|
|
for (int wn=0; wn<OPTN; wn++) {
|
|
sum[wm][wn] += dot(A_k0, B_k0[wn]);
|
|
}
|
|
}
|
|
}
|
|
|
|
{
|
|
int col = lidn;
|
|
for (int wn=0; wn<OPTN; wn++) {
|
|
B_k0[wn].s0 = Blocal[col][4];
|
|
B_k0[wn].s1 = Blocal[col][5];
|
|
B_k0[wn].s2 = Blocal[col][6];
|
|
B_k0[wn].s3 = Blocal[col][7];
|
|
col += TPWN;
|
|
}
|
|
int row = lidm;
|
|
for (int wm=0; wm<OPTM; wm++) {
|
|
A_k0.s0 = Alocal[4][row];
|
|
A_k0.s1 = Alocal[5][row];
|
|
A_k0.s2 = Alocal[6][row];
|
|
A_k0.s3 = Alocal[7][row];
|
|
for (int wn=0; wn<OPTN; wn++) {
|
|
sum[wm][wn] += dot(A_k0, B_k0[wn]);
|
|
}
|
|
row += TPWM;
|
|
}
|
|
}
|
|
// Synchronise before loading the next tile
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
}
|
|
|
|
// Store the final results in C
|
|
for (int wm=0; wm<OPTM; wm++) {
|
|
int globalRow = offsetM + lidm + wm*TPWM;
|
|
for (int wn=0; wn<OPTN; wn++) {
|
|
int globalCol = offsetN + lidn + wn*TPWN;
|
|
#ifdef BIAS
|
|
sum[wm][wn] += bias[globalCol];
|
|
#endif
|
|
C[globalRow*N + globalCol] = sum[wm][wn];
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
// double buffer
|
|
__kernel void matmul_local_double_buf(const int M, const int N, const int K,
|
|
__global const FLOAT* A,
|
|
#if (defined USE_LOW_BIT_WEIGHT_INT8)
|
|
__global const char* B,
|
|
__global const float* dequantScale,
|
|
__global const float* dequantOffset,
|
|
#elif (defined USE_LOW_BIT_WEIGHT_INT4)
|
|
__global const uchar* B,
|
|
__global const float* dequantScale,
|
|
__global const float* dequantOffset,
|
|
#else
|
|
__global const FLOAT* B,
|
|
#endif
|
|
#ifdef BIAS
|
|
__global const FLOAT* bias,
|
|
#endif
|
|
__global FLOAT* C) {
|
|
|
|
// Local thread id
|
|
const ushort lidm = get_local_id(0); // Local row ID
|
|
const ushort lidn = get_local_id(1); // Local col ID
|
|
// group id
|
|
const ushort offsetM = get_group_id(0) * OPWM; // Work-group offset M
|
|
const ushort offsetN = get_group_id(1) * OPWN; // Work-group offset N
|
|
|
|
// Local memory for work-group cache of A and B
|
|
__local FLOAT AlocalR[CPWK][OPWM];
|
|
__local FLOAT BlocalR[OPWN][CPWK+2];
|
|
__local FLOAT AlocalC[CPWK][OPWM];
|
|
__local FLOAT BlocalC[OPWN][CPWK+2];
|
|
|
|
// Allocate register space
|
|
COMPUTE_FLOAT sum[OPTM][OPTN];
|
|
|
|
// Initialise the accumulation registers
|
|
for (ushort wm=0; wm<OPTM; wm++) {
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
sum[wm][wn] = 0.0f;
|
|
}
|
|
}
|
|
|
|
// Loop over all tiles
|
|
const ushort numLoops = K/CPWK;
|
|
ushort lid = lidn*TPWM + lidm;
|
|
|
|
for (ushort t=0; t<numLoops; t++) {
|
|
// Load one work-group of A and B into local memory
|
|
for (ushort la=0; la<LPTA; la++) {
|
|
ushort id = la*TPWN*TPWM + lid;
|
|
ushort row = id % OPWM;
|
|
ushort col = id / OPWM;
|
|
ushort tiledIndex = CPWK*t + col;
|
|
#ifdef TRANSPOSE_A
|
|
// [K, M]
|
|
AlocalR[col][row] = A[tiledIndex*M + (offsetM + row)];
|
|
#else
|
|
// [M, K]
|
|
AlocalR[col][row] = A[(offsetM + row)*K + tiledIndex];
|
|
#endif
|
|
}
|
|
|
|
for (ushort la=0; la<LPTB; la++) {
|
|
ushort id = la*TPWN*TPWM + lid;
|
|
ushort row = id % OPWN;
|
|
ushort col = id / OPWN;
|
|
ushort tiledIndex = CPWK*t + col;
|
|
#ifdef TRANSPOSE_B
|
|
// [N, K]
|
|
BlocalR[row][col] = B[(offsetN + row)*K + tiledIndex];
|
|
#else
|
|
// [K, N]
|
|
BlocalR[row][col] = B[tiledIndex*N + offsetN + row];
|
|
#endif
|
|
}
|
|
|
|
// Synchronise to make sure the tile is loaded
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
// Loop over the values of a single tile
|
|
|
|
// Perform the computation
|
|
FLOAT4 A_k0, B_k0[OPTN];
|
|
{
|
|
ushort row = lidm;
|
|
ushort col = lidn;
|
|
|
|
A_k0.s0 = AlocalR[0][row];
|
|
A_k0.s1 = AlocalR[1][row];
|
|
A_k0.s2 = AlocalR[2][row];
|
|
A_k0.s3 = AlocalR[3][row];
|
|
|
|
#pragma unroll
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
B_k0[wn].s0 = BlocalR[col][0];
|
|
B_k0[wn].s1 = BlocalR[col][1];
|
|
B_k0[wn].s2 = BlocalR[col][2];
|
|
B_k0[wn].s3 = BlocalR[col][3];
|
|
sum[0][wn] += dot(A_k0, B_k0[wn]);
|
|
col += TPWN;
|
|
}
|
|
|
|
#pragma unroll
|
|
for(ushort wm=1; wm<OPTM; wm++) {
|
|
row += TPWM;
|
|
A_k0.s0 = AlocalR[0][row];
|
|
A_k0.s1 = AlocalR[1][row];
|
|
A_k0.s2 = AlocalR[2][row];
|
|
A_k0.s3 = AlocalR[3][row];
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
sum[wm][wn] += dot(A_k0, B_k0[wn]);
|
|
}
|
|
}
|
|
}
|
|
|
|
{
|
|
int col = lidn;
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
B_k0[wn].s0 = BlocalR[col][4];
|
|
B_k0[wn].s1 = BlocalR[col][5];
|
|
B_k0[wn].s2 = BlocalR[col][6];
|
|
B_k0[wn].s3 = BlocalR[col][7];
|
|
col += TPWN;
|
|
}
|
|
ushort row = lidm;
|
|
for (ushort wm=0; wm<OPTM; wm++) {
|
|
A_k0.s0 = AlocalR[4][row];
|
|
A_k0.s1 = AlocalR[5][row];
|
|
A_k0.s2 = AlocalR[6][row];
|
|
A_k0.s3 = AlocalR[7][row];
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
sum[wm][wn] += dot(A_k0, B_k0[wn]);
|
|
}
|
|
row += TPWM;
|
|
}
|
|
}
|
|
|
|
t++;
|
|
// Loop over the values of a single tile
|
|
// Load one work-group of A and B into local memory
|
|
for (ushort la=0; la<LPTA; la++) {
|
|
ushort id = la*TPWN*TPWM + lid;
|
|
ushort row = id % OPWM;
|
|
ushort col = id / OPWM;
|
|
ushort tiledIndex = CPWK*t + col;
|
|
#ifdef TRANSPOSE_A
|
|
// [K, M]
|
|
AlocalC[col][row] = A[tiledIndex*M + (offsetM + row)];
|
|
#else
|
|
// [M, K]
|
|
AlocalC[col][row] = A[(offsetM + row)*K + tiledIndex];
|
|
#endif
|
|
}
|
|
|
|
for (ushort la=0; la<LPTB; la++) {
|
|
ushort id = la*TPWN*TPWM + lid;
|
|
ushort row = id % OPWN;
|
|
ushort col = id / OPWN;
|
|
ushort tiledIndex = CPWK*t + col;
|
|
#ifdef TRANSPOSE_B
|
|
// [N, K]
|
|
BlocalC[row][col] = B[(offsetN + row)*K + tiledIndex];
|
|
#else
|
|
// [K, N]
|
|
BlocalC[row][col] = B[tiledIndex*N + offsetN + row];
|
|
#endif
|
|
}
|
|
|
|
// Synchronise to make sure the tile is loaded
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
// Perform the computation
|
|
{
|
|
ushort row = lidm;
|
|
ushort col = lidn;
|
|
|
|
A_k0.s0 = AlocalC[0][row];
|
|
A_k0.s1 = AlocalC[1][row];
|
|
A_k0.s2 = AlocalC[2][row];
|
|
A_k0.s3 = AlocalC[3][row];
|
|
|
|
#pragma unroll
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
B_k0[wn].s0 = BlocalC[col][0];
|
|
B_k0[wn].s1 = BlocalC[col][1];
|
|
B_k0[wn].s2 = BlocalC[col][2];
|
|
B_k0[wn].s3 = BlocalC[col][3];
|
|
sum[0][wn] += dot(A_k0, B_k0[wn]);
|
|
col += TPWN;
|
|
}
|
|
|
|
#pragma unroll
|
|
for(ushort wm=1; wm<OPTM; wm++) {
|
|
row += TPWM;
|
|
A_k0.s0 = AlocalC[0][row];
|
|
A_k0.s1 = AlocalC[1][row];
|
|
A_k0.s2 = AlocalC[2][row];
|
|
A_k0.s3 = AlocalC[3][row];
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
sum[wm][wn] += dot(A_k0, B_k0[wn]);
|
|
}
|
|
}
|
|
}
|
|
|
|
{
|
|
ushort col = lidn;
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
B_k0[wn].s0 = BlocalC[col][4];
|
|
B_k0[wn].s1 = BlocalC[col][5];
|
|
B_k0[wn].s2 = BlocalC[col][6];
|
|
B_k0[wn].s3 = BlocalC[col][7];
|
|
col += TPWN;
|
|
}
|
|
ushort row = lidm;
|
|
for (ushort wm=0; wm<OPTM; wm++) {
|
|
A_k0.s0 = AlocalC[4][row];
|
|
A_k0.s1 = AlocalC[5][row];
|
|
A_k0.s2 = AlocalC[6][row];
|
|
A_k0.s3 = AlocalC[7][row];
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
sum[wm][wn] += dot(A_k0, B_k0[wn]);
|
|
}
|
|
row += TPWM;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Store the final results in C
|
|
for (ushort wm=0; wm<OPTM; wm++) {
|
|
ushort globalRow = offsetM + lidm + wm*TPWM;
|
|
for (ushort wn=0; wn<OPTN; wn++) {
|
|
ushort globalCol = offsetN + lidn + wn*TPWN;
|
|
#ifdef BIAS
|
|
sum[wm][wn] += bias[globalCol];
|
|
#endif
|
|
C[globalRow*N + globalCol] = sum[wm][wn];
|
|
}
|
|
}
|
|
}
|