#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

//removing bias overhead and fusing with cutlass
// __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];
// }

//using big cutlass for kernels for BT>=64
using GemmDefault = cutlass::gemm::device::Gemm<
    float, cutlass::layout::RowMajor,
    float, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<128,128,32>,
    cutlass::gemm::GemmShape<64,64,32>,
    cutlass::gemm::GemmShape<16,8,8>,
    cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    3 // Stages
>;

//small cutlass for kernels for BT<64, we can use smaller threadblocks and warps to increase occupancy
using GemmsmallDefault = cutlass::gemm::device::Gemm<
    float, cutlass::layout::RowMajor,
    float, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<32,128,32>,
    cutlass::gemm::GemmShape<32,64,32>,
    cutlass::gemm::GemmShape<16,8,8>,
    cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    2
>;

using GemmMediumDefault = cutlass::gemm::device::Gemm<
    float, cutlass::layout::RowMajor,
    float, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<64,128,32>,
    cutlass::gemm::GemmShape<32,64,32>,
    cutlass::gemm::GemmShape<16,8,8>,
    cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    3
>;

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;

    if(bias != NULL){
        if(BT <= 32){
            GemmsmallDefault gemm_op;
            typename GemmsmallDefault::Arguments args{ 
                {BT, OC, C}, // problem size
                {inp, C}, // A matrix
                {weight, C}, // B matrix
                {bias, 0}, // C matrix
                {out, OC}, // D matrix (output)
                {1.0f, 1.0f} // alpha, beta for the epilogue
            };
            gemm_op(args);
        }
        else if(BT <= 128){
            GemmMediumDefault gemm_op;
            typename GemmMediumDefault::Arguments args{ 
                {BT, OC, C}, // problem size
                {inp, C}, // A matrix
                {weight, C}, // B matrix
                {bias, 0}, // C matrix
                {out, OC}, // D matrix (output)
                {1.0f, 1.0f} // alpha, beta for the epilogue
            };
            gemm_op(args);
        }
        else {
            GemmDefault gemm_op;
            typename GemmDefault::Arguments args{ 
                {BT, OC, C}, // problem size
                {inp, C}, // A matrix
                {weight, C}, // B matrix
                {bias, 0}, // C matrix
                {out, OC}, // D matrix (output)
                {1.0f, 1.0f} // alpha, beta for the epilogue
            };
            gemm_op(args);
        }
    }
    else{
        if(BT <= 32){
            GemmsmallDefault gemm_op;
            typename GemmsmallDefault::Arguments args{ 
                {BT, OC, C}, // problem size
                {inp, C}, // A matrix
                {weight, C}, // B matrix
                {out, OC}, // C matrix (output)
                {out, OC}, // D matrix (output)
                {1.0f, 0.0f} // alpha, beta for the epilogue
            };
            gemm_op(args);
        }
        else if(BT <= 128){
            GemmDefault gemm_op;
            typename GemmDefault::Arguments args{ 
                {BT, OC, C}, // problem size
                {inp, C}, // A matrix
                {weight, C}, // B matrix
                {out, OC}, // C matrix (output)
                {out, OC}, // D matrix (output)
                {1.0f, 0.0f} // alpha, beta for the epilogue
            };
            gemm_op(args);
        }
        else{
            GemmDefault gemm_op;
            typename GemmDefault::Arguments args{ 
                {BT, OC, C}, // problem size
                {inp, C}, // A matrix
                {weight, C}, // B matrix
                {out, OC}, // C matrix (output)
                {out, OC}, // D matrix (output)
                {1.0f, 0.0f} // alpha, beta for the epilogue
            };
            gemm_op(args);
        }
    }

    
}

using Gemmgelu = cutlass::gemm::device::Gemm<
    float, cutlass::layout::RowMajor,
    float, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<128, 128, 32>,
    cutlass::gemm::GemmShape<64, 64, 32>,
    cutlass::gemm::GemmShape<16, 8, 8>,
    cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    3>;

using GemmgeluSmall = cutlass::gemm::device::Gemm<
    float, cutlass::layout::RowMajor,
    float, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<32, 128, 32>,
    cutlass::gemm::GemmShape<32, 64, 32>,
    cutlass::gemm::GemmShape<16, 8, 8>,
    cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    2>;

using GemmgeluMedium = cutlass::gemm::device::Gemm<
    float, cutlass::layout::RowMajor,
    float, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<64, 128, 32>,
    cutlass::gemm::GemmShape<32, 64, 32>,
    cutlass::gemm::GemmShape<16, 8, 8>,
    cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    3>;

// 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;

    if (BT <= 32) {
        GemmgeluSmall op;
        typename GemmgeluSmall::Arguments args{
            {BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
        op(args);
    } else if (BT <= 128) {
        GemmgeluMedium op;
        typename GemmgeluMedium::Arguments args{
            {BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
        op(args);
    } else {
        Gemmgelu op;
        typename Gemmgelu::Arguments args{
            {BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
        op(args);
    }
}

#endif // __MATMUL_KERNEL_CUH__