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>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <mutex>
#include <vector>

#define CUBLAS_CHECK(status) do { cublasStatus_t s=(status); TORCH_CHECK(s == CUBLAS_STATUS_SUCCESS, "cuBLASLt error status=", (int)s, " line=", __LINE__); } while(0)

static cublasLtHandle_t ltHandle = nullptr;
static std::once_flag handle_once;
static void init_handle(){ CUBLAS_CHECK(cublasLtCreate(&ltHandle)); }

struct LtPlan {
    int m, n, k, fast;
    bool relu;
    cublasLtMatmulDesc_t op;
    cublasLtMatrixLayout_t aD, bD, cD;
    cublasLtMatmulAlgo_t algo;
};
static std::vector<LtPlan> plans;
static std::mutex plan_mutex;

static LtPlan* get_plan(int m, int n, int k, bool relu, bool fast, size_t workspaceSize) {
    std::call_once(handle_once, init_handle);
    std::lock_guard<std::mutex> guard(plan_mutex);
    for (auto &p : plans) {
        if (p.m==m && p.n==n && p.k==k && p.relu==relu && p.fast==(int)fast) return &p;
    }
    LtPlan p; p.m=m; p.n=n; p.k=k; p.relu=relu; p.fast=(int)fast;
    CUBLAS_CHECK(cublasLtMatmulDescCreate(&p.op, fast ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F, CUDA_R_32F));
    cublasOperation_t ta=CUBLAS_OP_T, tb=CUBLAS_OP_N;
    CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(p.op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)));
    CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(p.op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)));
    cublasLtEpilogue_t epi = relu ? CUBLASLT_EPILOGUE_RELU_BIAS : CUBLASLT_EPILOGUE_BIAS;
    CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(p.op, CUBLASLT_MATMUL_DESC_EPILOGUE, &epi, sizeof(epi)));

    // Row-major tensors are interpreted as column-major transposes:
    // W storage -> (k,n), X storage -> (k,m), Y storage -> (n,m).
    // This computes Y^T = W * X^T, equivalent to PyTorch linear Y = X * W^T.
    CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&p.aD, CUDA_R_32F, k, n, k));
    CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&p.bD, CUDA_R_32F, k, m, k));
    CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&p.cD, CUDA_R_32F, n, m, n));

    cublasLtMatmulPreference_t pref;
    CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&pref));
    CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize)));
    cublasLtMatmulHeuristicResult_t heur = {};
    int ret = 0;
    CUBLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(ltHandle, p.op, p.aD, p.bD, p.cD, p.cD, pref, 1, &heur, &ret));
    TORCH_CHECK(ret > 0, "no cuBLASLt matmul algorithm");
    p.algo = heur.algo;
    CUBLAS_CHECK(cublasLtMatmulPreferenceDestroy(pref));
    plans.push_back(p);
    return &plans.back();
}

static inline void lt_linear(torch::Tensor X, torch::Tensor W, torch::Tensor bias, torch::Tensor Y,
                             bool relu, bool fast, torch::Tensor workspace) {
    int m=(int)X.size(0), k=(int)X.size(1), n=(int)W.size(0);
    size_t ws = (size_t)workspace.numel();
    LtPlan* plan = get_plan(m, n, k, relu, fast, ws);
    const void* bp = bias.data_ptr<float>();
    CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(plan->op, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bp, sizeof(bp)));
    float alpha=1.0f, beta=0.0f;
    CUBLAS_CHECK(cublasLtMatmul(ltHandle, plan->op,
        &alpha, W.data_ptr<float>(), plan->aD, X.data_ptr<float>(), plan->bD,
        &beta, Y.data_ptr<float>(), plan->cD, Y.data_ptr<float>(), plan->cD,
        &plan->algo, workspace.data_ptr(), ws, at::cuda::getCurrentCUDAStream()));
}

torch::Tensor mlp_forward_buf(torch::Tensor x, torch::Tensor w1, torch::Tensor b1,
                              torch::Tensor w2, torch::Tensor b2, torch::Tensor w3, torch::Tensor b3,
                              torch::Tensor y1, torch::Tensor y2, torch::Tensor y3, torch::Tensor workspace) {
    lt_linear(x,  w1, b1, y1, true,  false, workspace);
    lt_linear(y1, w2, b2, y2, true,  true,  workspace);
    lt_linear(y2, w3, b3, y3, false, false, workspace);
    return y3;
}
'''

cuda_src = r'''
extern "C" __global__ void mlp_presence_kernel() { }
'''

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

class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_layer_sizes, output_size):
        super(ModelNew, self).__init__()
        layers = []
        cur = input_size
        for h in hidden_layer_sizes:
            layers.append(nn.Linear(cur, h))
            layers.append(nn.ReLU())
            cur = h
        layers.append(nn.Linear(cur, output_size))
        self.network = nn.Sequential(*layers)
        self._buf1 = None
        self._buf2 = None
        self._buf3 = None
        self._workspace = None

    def _ensure_buffers(self, x):
        b = x.shape[0]
        h1 = self.network[0].weight.shape[0]
        h2 = self.network[2].weight.shape[0]
        o = self.network[4].weight.shape[0]
        if self._buf1 is None or self._buf1.device != x.device or self._buf1.dtype != x.dtype or self._buf1.shape != (b, h1):
            self._buf1 = torch.empty((b, h1), device=x.device, dtype=x.dtype)
            self._buf2 = torch.empty((b, h2), device=x.device, dtype=x.dtype)
            self._buf3 = torch.empty((b, o), device=x.device, dtype=x.dtype)
            self._workspace = torch.empty((64 * 1024 * 1024,), device=x.device, dtype=torch.uint8)

    def forward(self, x):
        x = x.contiguous()
        self._ensure_buffers(x)
        return _ext.mlp_forward_buf(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,
                                    self._buf1, self._buf2, self._buf3, self._workspace)