#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__