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 = 1024
input_size = 8192
hidden_layer_sizes = [1024] * 16  # deep network with wider layers
output_size = 8192

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 <vector>
at::Tensor mlp_lt_forward_buf(at::Tensor x, std::vector<at::Tensor> weights, std::vector<at::Tensor> biases, std::vector<at::Tensor> outs, at::Tensor workspace);
'''
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 <vector>
#include <stdexcept>
#include <string>
__global__ void _kernelbench_marker_kernel() {}
#define CHECK_CUBLAS(x) do{cublasStatus_t st=(x); if(st!=CUBLAS_STATUS_SUCCESS) throw std::runtime_error(std::string("cublasLt ")+std::to_string((int)st));}while(0)
static cublasLtHandle_t geth(){static cublasLtHandle_t h=[](){cublasLtHandle_t t; cublasLtCreate(&t); return t;}(); return h;}
at::Tensor mlp_lt_forward_buf(at::Tensor x,std::vector<at::Tensor> weights,std::vector<at::Tensor> biases,std::vector<at::Tensor> outs,at::Tensor workspace){
 c10::cuda::CUDAGuard guard(x.device()); auto y=x; int64_t L=weights.size(); auto h=geth(); auto stream=at::cuda::getCurrentCUDAStream(); float alpha=1,beta=0; size_t ws=workspace.numel();
 for(int64_t i=0;i<L;i++){auto w=weights[i], b=biases[i], out=outs[i]; int64_t m=y.size(0), k=y.size(1), n=w.size(0);
  cublasLtMatmulDesc_t op; CHECK_CUBLAS(cublasLtMatmulDescCreate(&op,CUBLAS_COMPUTE_32F,CUDA_R_32F)); cublasOperation_t ta=CUBLAS_OP_T,tb=CUBLAS_OP_N; CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(op,CUBLASLT_MATMUL_DESC_TRANSA,&ta,sizeof(ta))); CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(op,CUBLASLT_MATMUL_DESC_TRANSB,&tb,sizeof(tb)));
  cublasLtEpilogue_t epi=(i==L-1)?CUBLASLT_EPILOGUE_BIAS:CUBLASLT_EPILOGUE_RELU_BIAS; CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(op,CUBLASLT_MATMUL_DESC_EPILOGUE,&epi,sizeof(epi))); const void* bp=b.data_ptr<float>(); CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(op,CUBLASLT_MATMUL_DESC_BIAS_POINTER,&bp,sizeof(bp)));
  cublasLtMatrixLayout_t A,B,C,D; CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&A,CUDA_R_32F,k,n,k)); CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&B,CUDA_R_32F,k,m,k)); CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&C,CUDA_R_32F,n,m,n)); CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&D,CUDA_R_32F,n,m,n));
  cublasLtMatmulPreference_t pref; CHECK_CUBLAS(cublasLtMatmulPreferenceCreate(&pref)); CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute(pref,CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,&ws,sizeof(ws))); cublasLtMatmulHeuristicResult_t heur; int ret=0; CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(h,op,A,B,C,D,pref,1,&heur,&ret)); if(!ret) throw std::runtime_error("no algo");
  CHECK_CUBLAS(cublasLtMatmul(h,op,&alpha,w.data_ptr<float>(),A,y.data_ptr<float>(),B,&beta,out.data_ptr<float>(),C,out.data_ptr<float>(),D,&heur.algo,workspace.data_ptr(),ws,stream));
  cublasLtMatmulPreferenceDestroy(pref); cublasLtMatrixLayoutDestroy(A);cublasLtMatrixLayoutDestroy(B);cublasLtMatrixLayoutDestroy(C);cublasLtMatrixLayoutDestroy(D);cublasLtMatmulDescDestroy(op); y=out;
 } return y;}
'''
_ext=load_inline(name='mlp_lt_buf_final2',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['mlp_lt_forward_buf'],verbose=False,extra_cflags=['-O3'],extra_cuda_cflags=['-O3'],extra_ldflags=['-lcublasLt','-lcublas'])
class ModelNew(nn.Module):
 def __init__(self,input_size,hidden_layer_sizes,output_size):
  super().__init__(); layers=[]; cur=input_size; self.linears=nn.ModuleList(); self._out_sizes=[]
  for h in hidden_layer_sizes:
   lin=nn.Linear(cur,h); layers += [lin, nn.ReLU()]; self.linears.append(lin); self._out_sizes.append(h); cur=h
  lin=nn.Linear(cur,output_size); layers.append(lin); self.linears.append(lin); self._out_sizes.append(output_size); self.network=nn.Sequential(*layers); self._weights=[l.weight for l in self.linears]; self._biases=[l.bias for l in self.linears]; self._outs=None; self._workspace=None
 def _ensure_buffers(self,x):
  if self._outs is None or self._outs[0].device != x.device or self._outs[0].dtype != x.dtype or self._outs[0].shape[0] != x.shape[0]:
   self._outs=[torch.empty((x.shape[0],s),device=x.device,dtype=x.dtype) for s in self._out_sizes]
   self._workspace=torch.empty((32*1024*1024,),device=x.device,dtype=torch.uint8)
 def forward(self,x):
  self._ensure_buffers(x)
  return _ext.mlp_lt_forward_buf(x,self._weights,self._biases,self._outs,self._workspace)