#ifndef __MATMUL_KERNEL_CUH__
#define __MATMUL_KERNEL_CUH__

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


__global__ void bias_add_kernel (float* out, const float* bias, int BT, int OC){
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    if(x < BT*OC){
        out[x] +=  bias[x % OC];
    }
}

// __global__ void matmul_forward_kernel(float* out, const float* inp, const float* weight, 
//                                       const float* bias,int B,int T,  int C, int OC) {
//     // no need for this kernel
    
// }

// 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
    int BT = B*T;
    float alpha = 1.0f;
    float beta = 0.0f;

    cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, BT, C, &alpha, weight, C, inp, C, &beta, out, OC);

    if (bias != NULL) {
        int total = BT * OC;
        bias_add_kernel<<<(total + 255) / 256, 256>>>(out, bias, BT, OC);
    }

}

#endif // __MATMUL_KERNEL_CUH__