#ifndef __MATMUL_KERNEL_CUH__
#define __MATMUL_KERNEL_CUH__

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

using namespace nvcuda;

#define WMMA_M 16
#define WMMA_N 16
#define WMMA_K 8

#define WARPS_M 4
#define WARPS_N 2
#define BM 64
#define BN 32
#define BK WMMA_K

__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 BT = B*T;

    // block output
    int start_row = blockIdx.x * BM;
    int start_col = blockIdx.y * BN;

    // which warp
    int warp_id = threadIdx.x / 32;
    int warp_row = warp_id / WARPS_N;
    int warp_col = warp_id % WARPS_N;

    int warp_out_row = start_row + warp_row * WMMA_M;
    int warp_out_col = start_col + warp_col * WMMA_N;

    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;

    wmma::fill_fragment(c_frag, 0.0f);

    for(int k = 0; k < C; k += WMMA_K){
        if(warp_out_row < BT && k < C){
            wmma::load_matrix_sync(a_frag, inp + warp_out_row * C + k, C);

            #pragma unroll
            for(int i = 0; i< a_frag.num_elements; i++){
                a_frag.x[i] = wmma::__float_to_tf32(a_frag.x[i]);
            }
        }

        if (warp_out_col < OC && k < C) {
            wmma::load_matrix_sync(b_frag, weight + warp_out_col * C + k, C);
            #pragma unroll
            for (int i = 0; i < b_frag.num_elements; i++) {
                b_frag.x[i] = wmma::__float_to_tf32(b_frag.x[i]);
            }
        }

        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
    }

    //store  in s mem, add bais durinthe writw out  for performance
    __shared__ float out_tile[BM * BN];

    int row_smem = warp_row * WMMA_M;
    int col_smem = warp_col * WMMA_N;
    wmma::store_matrix_sync(out_tile + row_smem * BN + col_smem, c_frag, BN, wmma::mem_row_major);
    __syncthreads();

    for(int i = threadIdx.x; i < BM*BN; i += blockDim.x){
        int r_l = i / BN;
        int c_l = i % BN;
        int r_g = start_row + r_l;
        int c_g = start_col + c_l;

        if(r_g < BT && c_g < OC){
            if(bias != NULL){
                out_tile[i] += bias[c_g];
            }
            out[r_g * OC + c_g] = out_tile[i];
        }
    }
}

// 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;
    dim3 block(256);
    dim3 grid((BT + BM - 1) / BM, (OC + BN - 1) / BN);
    matmul_forward_kernel<<<grid, block>>>(out, inp, weight, bias,B,T, C, OC);

}

#endif // __MATMUL_KERNEL_CUH__