38 lines
1.3 KiB
Plaintext
38 lines
1.3 KiB
Plaintext
#ifndef __MATMUL_KERNEL_CUH__
|
|
#define __MATMUL_KERNEL_CUH__
|
|
|
|
#include <cuda_runtime.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
|
|
__global__ void matmul_forward_kernel(float* __restrict__ out, const float* __restrict__ inp, const float* __restrict__ weight,
|
|
const float* __restrict__ bias,int B,int T, int C, int OC) {
|
|
// Implement this
|
|
int b = blockDim.x*blockIdx.x + threadIdx.x; // B
|
|
int t = blockDim.y*blockIdx.y + threadIdx.y; // T
|
|
int oc = blockDim.z*blockIdx.z + threadIdx.z; // OC
|
|
|
|
if (b >= B || t >= T || oc >= OC) {
|
|
return;
|
|
}
|
|
|
|
float sum = (bias != NULL) ? bias[oc] : 0.0f;
|
|
const float* __restrict__ inp_token_ptr = inp + b * T * C + t * C;
|
|
const float* __restrict__ w_r = weight + oc * C;
|
|
for(int i = 0; i < C; i++){
|
|
sum += inp_token_ptr[i] * w_r[i] ;
|
|
}
|
|
out[b * T * OC + t * OC + oc] = sum;
|
|
}
|
|
|
|
// 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
|
|
dim3 block(8,8,8);
|
|
dim3 grid((B+block.x-1)/block.x, (T+block.y-1)/block.y, (OC+block.z-1)/block.z);
|
|
|
|
matmul_forward_kernel<<<grid, block>>>(out, inp, weight, bias,B,T, C, OC);
|
|
|
|
}
|
|
|
|
#endif // __MATMUL_KERNEL_CUH__ |