#ifndef __MATMUL_KERNEL_CUH__
#define __MATMUL_KERNEL_CUH__

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

#define TILE_SIZE 32
#define PADDING 1

//op12
// Adding +1 padding to the inner dimension shifts each row by one bank, distributing the 32 column reads across all 32 banks.

__global__ void matmul_kernel(const float* __restrict__ inp, float* __restrict__ out, const float* __restrict__ weights, 
                            const float* __restrict__ biases, int B, int T, int OC, int C ) {
    // Implement this
    // reshape (BT X C) * (OC X C)^T -> (BT X OC)

    int bt = blockIdx.y * TILE_SIZE + threadIdx.y;
    int oc = blockIdx.x * TILE_SIZE + threadIdx.x;
    int BT  = B*T;

    __shared__ float A_smem_tile[TILE_SIZE][TILE_SIZE + PADDING];
    __shared__ float B_smem_tile[TILE_SIZE][TILE_SIZE + PADDING];

    float acc = 0.0f;
    int num_tiles = (C + TILE_SIZE - 1) / TILE_SIZE;

    for(int t = 0; t < num_tiles; t++){
        int c = t * TILE_SIZE + threadIdx.x;
        int c_ = t * TILE_SIZE + threadIdx.y;
        A_smem_tile[threadIdx.y][threadIdx.x] = (bt < BT && c < C) ? inp[bt * C + c] : 0.0f;
        B_smem_tile[threadIdx.y][threadIdx.x] = (oc < OC && c_ < C) ? weights[oc * C + c_] : 0.0f;

        __syncthreads();

        #pragma unroll
        for(int i = 0; i < TILE_SIZE; i++){
            acc += A_smem_tile[threadIdx.y][i] * B_smem_tile[i][threadIdx.x];
        }
        __syncthreads();
    }
    if(bt < BT && oc < OC){
        if(biases != nullptr){
            acc += biases[oc];
        }
        out[bt * OC + oc] = acc;
    }
}

void matmul_forward(float* out, const float* inp, const float* weight, const float* bias,
                    int B, int T, int C, int OC){
    dim3 block(TILE_SIZE, TILE_SIZE);
    dim3 grid((OC + TILE_SIZE - 1) / TILE_SIZE, (B * T + TILE_SIZE - 1) / TILE_SIZE);
    matmul_kernel<<<grid, block>>>(inp, out, weight, bias, B, T, OC, C);
}

#endif