#ifndef __MATMUL_KERNEL_CUH__
#define __MATMUL_KERNEL_CUH__

#include <cuda_runtime.h>
#include "../utils/cuda_utils.cuh"
#include "cutlass/cutlass.h"
#include "cutlass/gemm/device/gemm.h"
#include "/work/hdd/bche/hshah6/cutlass/include/cutlass/epilogue/thread/linear_combination_gelu.h"

//referenced: https://github.com/NVIDIA/cutlass/blob/main/examples/14_ampere_tf32_tensorop_gemm/ampere_tf32_tensorop_gemm.cu

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

inline void matmul_forward(float* __restrict__ out, const float* __restrict__ inp, const float* __restrict__ weight, 
                                      const float* __restrict__ bias,int B,int T,  int C, int OC) {
    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_filler_kernel<<<(total + 255) / 256, 256>>>(out, bias, BT, OC);
    }
}

// Launch kernel here
inline void matmul_forward_fused_gelu(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;

    int total = BT * OC;
    bias_add_filler_kernel<<<(total + 255) / 256, 256>>>(out, bias, BT, OC);

    using ElementOutput = float;
    using ElementAccumulator = float;
    using ElementCompute = float;

    using EpilogueOp = cutlass::epilogue::thread::LinearCombinationGELU<ElementOutput, 128/ cutlass::sizeof_bits<ElementOutput>::value, ElementAccumulator, ElementCompute>;

    using Gemm = cutlass::gemm::device::Gemm<
                        float, cutlass::layout::RowMajor,
                        float, cutlass::layout::ColumnMajor,
                        ElementOutput, cutlass::layout::RowMajor,
                        ElementAccumulator,
                        cutlass::arch::OpClassTensorOp,
                        cutlass::arch::Sm80,
                        cutlass::gemm::GemmShape<128, 128, 32>,
                        cutlass::gemm::GemmShape<64, 64, 32>,
                        cutlass::gemm::GemmShape<16, 8, 8>,
                        EpilogueOp,
                        cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
                        2 // Stages
                    >;

    Gemm gemm_op;

    typename Gemm::Arguments args{ 
        {BT, OC, C}, // problem size
        {inp, C}, // A matrix
        {weight, C}, // B matrix
        {out, OC}, // C matrix
        {out, OC}, // D matrix (output)
        {1.0f, 1.0f} // alpha, beta for the epilogue
    };

    cutlass::Status status = gemm_op(args);
}

#endif // __MATMUL_KERNEL_CUH__