90 lines
2.4 KiB
Plaintext
90 lines
2.4 KiB
Plaintext
#ifndef __MATMUL_KERNEL_CUH__
|
|
#define __MATMUL_KERNEL_CUH__
|
|
|
|
#include <cuda_runtime.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
|
|
#define TILE_T 128
|
|
#define TILE_U 8
|
|
#define TILE_S 16
|
|
|
|
__global__ void matmul_forward_kernel(float* out, const float* inp, const float* weight,
|
|
const float* bias,int B,int T, int C, int OC) {
|
|
// Implement this
|
|
const int BT = B * T;
|
|
|
|
// getting which pair this thread is going to process
|
|
const int row = blockIdx.x * TILE_T + threadIdx.x;
|
|
// OC
|
|
const int col = blockIdx.y * TILE_U;
|
|
|
|
//register tile
|
|
float accumilation[TILE_U] = {0.0f};
|
|
|
|
//init the tile with bais
|
|
if(bias != NULL && row < BT){
|
|
#pragma unroll
|
|
for(int i = 0; i < TILE_U; i++){
|
|
if(col + i < OC){
|
|
accumilation[i] = bias[col + i];
|
|
}
|
|
}
|
|
}
|
|
|
|
__shared__ float N_tile[TILE_S][TILE_U];
|
|
|
|
for(int tile_start = 0; tile_start < C; tile_start += TILE_S){
|
|
const int s_l = threadIdx.x / TILE_U;
|
|
const int u_l = threadIdx.x % TILE_U;
|
|
|
|
const int c_g = tile_start + s_l;
|
|
const int oc_g = col + u_l;
|
|
|
|
if(c_g < C && oc_g < OC){
|
|
N_tile[s_l][u_l] = weight[oc_g * C + c_g];
|
|
}
|
|
else{
|
|
N_tile[s_l][u_l] = 0.0f;
|
|
}
|
|
__syncthreads();
|
|
|
|
if(row < BT){
|
|
#pragma unroll
|
|
for(int s = 0 ; s < TILE_S; s++){
|
|
const int c = tile_start + s;
|
|
|
|
float m = 0.0f;
|
|
if(c < C){
|
|
m = inp[row * C + c] ;
|
|
}
|
|
#pragma unroll
|
|
for(int u = 0; u < TILE_U; u++){
|
|
accumilation[u] +=m *N_tile[s][u];
|
|
}
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
if(row < BT){
|
|
#pragma unroll
|
|
for(int u = 0; u < TILE_U; u++){
|
|
if(col + u < OC){
|
|
out[row*OC+col+u] = accumilation[u];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Launch kernel here
|
|
void matmul_forward(float* out, const float* inp, const float* weight, const float* bias,
|
|
int B, int T, int C, int OC) {
|
|
// Implement this
|
|
const int BT = B*T;
|
|
dim3 block(TILE_T);
|
|
dim3 grid((BT+TILE_T-1)/TILE_T,(OC+TILE_U-1)/TILE_U);
|
|
|
|
matmul_forward_kernel<<<grid, block>>>(out, inp, weight, bias,B,T, C, OC);
|
|
|
|
}
|
|
|
|
#endif // __MATMUL_KERNEL_CUH__ |