import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a matrix multiplication (Gemm), applies Sigmoid,
    another Gemm, and computes LogSumExp over features.
    """
    def __init__(self, input_size, hidden_size, output_size):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(input_size, hidden_size)
        self.linear2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.linear1(x)
        x = torch.sigmoid(x)
        x = self.linear2(x)
        x = torch.logsumexp(x, dim=1)  # compute LogSumExp over features per sample
        return x

batch_size = 16384
input_size = 2048
hidden_size = 4096
output_size = 1024

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

def get_init_inputs():
    return [input_size, hidden_size, 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 forward_cuda(torch::Tensor x, 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 <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void sigmoid_inplace_kernel(float* __restrict__ data, int64_t n) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = idx; i < n; i += stride) {
        float v = data[i];
        data[i] = 1.0f / (1.0f + expf(-v));
    }
}

__inline__ __device__ float warp_reduce_max(float v) {
    for (int offset = 16; offset > 0; offset >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffff, v, offset));
    return v;
}
__inline__ __device__ float warp_reduce_sum(float v) {
    for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(0xffffffff, v, offset);
    return v;
}

__global__ void logsumexp_kernel(const float* __restrict__ input, float* __restrict__ output, int rows, int cols) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const float* rowptr = input + (int64_t)row * cols;
    __shared__ float smem[32];
    int lane = tid & 31;
    int warp = tid >> 5;

    float m = -INFINITY;
    for (int c = tid; c < cols; c += blockDim.x) m = fmaxf(m, rowptr[c]);
    m = warp_reduce_max(m);
    if (lane == 0) smem[warp] = m;
    __syncthreads();
    float bm = -INFINITY;
    if (tid < (blockDim.x >> 5)) bm = smem[tid];
    bm = warp_reduce_max(bm);
    if (tid == 0) smem[0] = bm;
    __syncthreads();
    bm = smem[0];

    float s = 0.0f;
    for (int c = tid; c < cols; c += blockDim.x) s += expf(rowptr[c] - bm);
    s = warp_reduce_sum(s);
    if (lane == 0) smem[warp] = s;
    __syncthreads();
    float bs = 0.0f;
    if (tid < (blockDim.x >> 5)) bs = smem[tid];
    bs = warp_reduce_sum(bs);
    if (tid == 0) output[row] = logf(bs) + bm;
}

torch::Tensor forward_cuda(torch::Tensor x, torch::Tensor w1, torch::Tensor b1, torch::Tensor w2, torch::Tensor b2) {
    auto h = at::addmm(b1, x, w1.t());
    int64_t hn = h.numel();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int threads = 256;
    int blocks = (int)std::min<int64_t>((hn + threads - 1) / threads, 65535);
    sigmoid_inplace_kernel<<<blocks, threads, 0, stream>>>(h.data_ptr<float>(), hn);
    auto y = at::addmm(b2, h, w2.t());
    auto out = torch::empty({y.size(0)}, y.options());
    logsumexp_kernel<<<(int)y.size(0), 256, 0, stream>>>(y.data_ptr<float>(), out.data_ptr<float>(), (int)y.size(0), (int)y.size(1));
    return out;
}
'''

_ext = load_inline(name='gemm_sigmoid_gemm_lse_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['forward_cuda'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(ModelNew, self).__init__()
        self.linear1 = nn.Linear(input_size, hidden_size)
        self.linear2 = nn.Linear(hidden_size, output_size)
    def forward(self, x):
        return _ext.forward_cuda(x, self.linear1.weight, self.linear1.bias, self.linear2.weight, self.linear2.bias)