#ifndef __MATMUL_KERNEL_CUH__
#define __MATMUL_KERNEL_CUH__

#include <cuda_runtime.h>
#include "../utils/cuda_utils.cuh"

__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
    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* inp_token_ptr = inp + b * T * C + t * C;
    for(int i = 0; i < C; i++){
        sum += inp_token_ptr[i] * weight[oc*C + 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__