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>

void bias_relu_launcher(torch::Tensor inp, torch::Tensor bias, torch::Tensor out);

// Forward for the fixed MLP pattern: all but the last layer are Linear+ReLU,
// last layer is Linear.  The heavy work is still dispatched to cuBLAS through
// ATen, while the elementwise bias+ReLU epilogue is fused in a small CUDA
// kernel and the whole network is driven from one Python/C++ call.
torch::Tensor mlp_forward(torch::Tensor x, std::vector<torch::Tensor> weights, std::vector<torch::Tensor> biases) {
    TORCH_CHECK(weights.size() == biases.size(), "weights/biases length mismatch");
    TORCH_CHECK(weights.size() >= 1, "expected at least one layer");
    auto h = x;
    const int64_t L = (int64_t)weights.size();
    for (int64_t i = 0; i < L - 1; ++i) {
        // h: [B, in], weight: [out, in] -> tmp: [B, out]
        auto tmp = at::matmul(h, weights[i].t());
        auto y = at::empty_like(tmp);
        bias_relu_launcher(tmp, biases[i], y);
        h = y;
    }
    // final linear, use the same addmm formulation as nn.Linear
    auto out = at::addmm(biases[L-1], h, weights[L-1].t());
    return out;
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

__global__ void bias_relu_float4_kernel(const float* __restrict__ inp,
                                        const float* __restrict__ bias,
                                        float* __restrict__ out,
                                        int64_t n4,
                                        int cols) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n4) return;
    float4 v = reinterpret_cast<const float4*>(inp)[idx];
    int col4 = ((int)(idx % (cols / 4))) * 4;
    float4 b = reinterpret_cast<const float4*>(bias)[col4 / 4];
    v.x = fmaxf(v.x + b.x, 0.0f);
    v.y = fmaxf(v.y + b.y, 0.0f);
    v.z = fmaxf(v.z + b.z, 0.0f);
    v.w = fmaxf(v.w + b.w, 0.0f);
    reinterpret_cast<float4*>(out)[idx] = v;
}

__global__ void bias_relu_scalar_kernel(const float* __restrict__ inp,
                                        const float* __restrict__ bias,
                                        float* __restrict__ out,
                                        int64_t n,
                                        int cols) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    int col = (int)(idx % cols);
    float v = inp[idx] + bias[col];
    out[idx] = v > 0.0f ? v : 0.0f;
}

void bias_relu_launcher(torch::Tensor inp, torch::Tensor bias, torch::Tensor out) {
    c10::cuda::CUDAGuard device_guard(inp.device());
    TORCH_CHECK(inp.is_cuda() && bias.is_cuda() && out.is_cuda(), "CUDA tensors expected");
    TORCH_CHECK(inp.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(inp.is_contiguous() && out.is_contiguous() && bias.is_contiguous(), "contiguous tensors expected");
    const int64_t n = inp.numel();
    const int cols = (int)inp.size(-1);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;
    if ((cols % 4) == 0 && ((uintptr_t)inp.data_ptr<float>() % 16) == 0 && ((uintptr_t)out.data_ptr<float>() % 16) == 0 && ((uintptr_t)bias.data_ptr<float>() % 16) == 0) {
        int64_t n4 = n / 4;
        int blocks = (int)((n4 + threads - 1) / threads);
        bias_relu_float4_kernel<<<blocks, threads, 0, stream>>>(inp.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), n4, cols);
    } else {
        int blocks = (int)((n + threads - 1) / threads);
        bias_relu_scalar_kernel<<<blocks, threads, 0, stream>>>(inp.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), n, cols);
    }
}
'''

_ext = load_inline(
    name='mlp_fused_bias_relu_ext_v1',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['mlp_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    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)
        # Cache the Linear modules; Parameter objects remain the authoritative
        # module parameters and are used directly in forward.
        self._linear_indices = [i for i, m in enumerate(self.network) if isinstance(m, nn.Linear)]

    def forward(self, x):
        weights = []
        biases = []
        for i in self._linear_indices:
            m = self.network[i]
            weights.append(m.weight)
            biases.append(m.bias)
        return _ext.mlp_forward(x, weights, biases)