import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, input_size, hidden_layer_sizes, output_size):
        """
        :param input_size: The number of input features
        :param hidden_layer_sizes: A list of ints containing the sizes of each hidden layer
        :param output_size: The number of output features
        """
        super(Model, self).__init__()
        
        layers = []
        current_input_size = input_size
        
        for hidden_size in hidden_layer_sizes:
            layers.append(nn.Linear(current_input_size, hidden_size))
            layers.append(nn.ReLU())
            current_input_size = hidden_size
        
        layers.append(nn.Linear(current_input_size, output_size))
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, x):
        """
        :param x: The input tensor, shape (batch_size, input_size)
        :return: The output tensor, shape (batch_size, output_size)
        """
        return self.network(x)

# Test code
batch_size = 128
input_size = 16384
hidden_layer_sizes = [32768, 32768]
output_size = 16384

def get_inputs():
    return [torch.rand(batch_size, input_size)]

def get_init_inputs():
    return [input_size, hidden_layer_sizes, output_size]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor mlp_forward(torch::Tensor x,
                          torch::Tensor w0, torch::Tensor b0,
                          torch::Tensor w1, torch::Tensor b1,
                          torch::Tensor w2, torch::Tensor b2);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <stdexcept>

static cublasLtHandle_t lt_handle = nullptr;
static inline void ck(cublasStatus_t s, const char* m) {
    if (s != CUBLAS_STATUS_SUCCESS) throw std::runtime_error(m);
}

__global__ void add_bias_kernel(float* __restrict__ y,
                                const float* __restrict__ b,
                                long long total, int N) {
    long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x;
    if (i < total) y[i] += b[i & (N - 1)];
}

static torch::Tensor linear_impl(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                                 bool relu, cudaStream_t stream) {
    const int64_t M64 = x.size(0), K64 = x.size(1), N64 = w.size(0);
    const int M = (int)M64, N = (int)N64, K = (int)K64;
    auto y = torch::empty({M64, N64}, x.options());

    cublasLtMatmulDesc_t op;
    cublasLtMatrixLayout_t ad, bd, cd;
    ck(cublasLtMatmulDescCreate(&op, CUBLAS_COMPUTE_32F, CUDA_R_32F), "op");

    cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N;
    ck(cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)), "ta");
    ck(cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)), "tb");

    cublasLtEpilogue_t epi = relu ? CUBLASLT_EPILOGUE_RELU_BIAS : CUBLASLT_EPILOGUE_DEFAULT;
    ck(cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_EPILOGUE, &epi, sizeof(epi)), "epi");
    if (relu) {
        const void* bp = bias.data_ptr<float>();
        ck(cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bp, sizeof(bp)), "bias");
    }

    // Row-major y[M,N] == column-major view [N,M].
    // Row-major w[N,K] == column-major view [K,N], transposed in GEMM.
    // Row-major x[M,K] == column-major view [K,M].
    ck(cublasLtMatrixLayoutCreate(&ad, CUDA_R_32F, K, N, K), "ad");
    ck(cublasLtMatrixLayoutCreate(&bd, CUDA_R_32F, K, M, K), "bd");
    ck(cublasLtMatrixLayoutCreate(&cd, CUDA_R_32F, N, M, N), "cd");

    float alpha = 1.0f, beta = 0.0f;
    cublasStatus_t st = cublasLtMatmul(lt_handle, op, &alpha,
                                       w.data_ptr<float>(), ad,
                                       x.data_ptr<float>(), bd,
                                       &beta,
                                       y.data_ptr<float>(), cd,
                                       y.data_ptr<float>(), cd,
                                       nullptr, nullptr, 0, stream);

    cublasLtMatrixLayoutDestroy(ad);
    cublasLtMatrixLayoutDestroy(bd);
    cublasLtMatrixLayoutDestroy(cd);
    cublasLtMatmulDescDestroy(op);
    ck(st, "matmul");

    if (!relu) {
        long long total = M64 * N64;
        int block = 256;
        int grid = (int)((total + block - 1) / block);
        add_bias_kernel<<<grid, block, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), total, N);
    }
    return y;
}

torch::Tensor mlp_forward(torch::Tensor x,
                          torch::Tensor w0, torch::Tensor b0,
                          torch::Tensor w1, torch::Tensor b1,
                          torch::Tensor w2, torch::Tensor b2) {
    TORCH_CHECK(x.is_cuda() && w0.is_cuda() && b0.is_cuda() && w1.is_cuda() && b1.is_cuda() && w2.is_cuda() && b2.is_cuda(), "CUDA tensors required");
    TORCH_CHECK(x.is_contiguous() && w0.is_contiguous() && b0.is_contiguous() && w1.is_contiguous() && b1.is_contiguous() && w2.is_contiguous() && b2.is_contiguous(), "contiguous tensors required");

    c10::cuda::CUDAGuard guard(x.device());
    if (!lt_handle) ck(cublasLtCreate(&lt_handle), "create");
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();

    auto y0 = linear_impl(x,  w0, b0, true,  stream);
    auto y1 = linear_impl(y0, w1, b1, true,  stream);
    auto y2 = linear_impl(y1, w2, b2, false, stream);
    return y2;
}
'''

_ext = load_inline(
    name='mlp_cublaslt_fused_submit',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['mlp_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcublasLt', '-lcublas'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_layer_sizes, output_size):
        super(ModelNew, self).__init__()
        layers = []
        current_input_size = input_size
        for hidden_size in hidden_layer_sizes:
            layers.append(nn.Linear(current_input_size, hidden_size))
            layers.append(nn.ReLU())
            current_input_size = hidden_size
        layers.append(nn.Linear(current_input_size, output_size))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return _ext.mlp_forward(
            x,
            self.network[0].weight, self.network[0].bias,
            self.network[2].weight, self.network[2].bias,
            self.network[4].weight, self.network[4].bias,
        )