import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a matrix multiplication, scales the result, adds a residual connection, clamps the output,
    applies LogSumExp, and finally applies the Mish activation function.
    """
    def __init__(self, input_size, hidden_size, scale_factor, clamp_min, clamp_max):
        super(Model, self).__init__()
        self.matmul = nn.Linear(input_size, hidden_size)
        self.scale_factor = scale_factor
        self.clamp_min = clamp_min
        self.clamp_max = clamp_max

    def forward(self, x):
        """
        Args:
            x: Input tensor of shape (batch_size, input_size).

        Returns:
            Output tensor of shape (batch_size, hidden_size).
        """
        x = self.matmul(x)
        x = x * self.scale_factor
        x = x + x
        x = torch.clamp(x, self.clamp_min, self.clamp_max)
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = x * torch.nn.functional.mish(x)  # Mish activation
        return x

batch_size = 1024
input_size = 8192
hidden_size = 8192
scale_factor = 2.0
clamp_min = -10.0
clamp_max = 10.0

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

def get_init_inputs():
    return [input_size, hidden_size, scale_factor, clamp_min, clamp_max]

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


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor fused_post_forward(torch::Tensor y, double scale_factor, double clamp_min, double clamp_max);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>
#include <math.h>

static __inline__ __device__ float warp_reduce_max(float v) {
    unsigned mask = 0xffffffffu;
    v = fmaxf(v, __shfl_down_sync(mask, v, 16));
    v = fmaxf(v, __shfl_down_sync(mask, v, 8));
    v = fmaxf(v, __shfl_down_sync(mask, v, 4));
    v = fmaxf(v, __shfl_down_sync(mask, v, 2));
    v = fmaxf(v, __shfl_down_sync(mask, v, 1));
    return v;
}

static __inline__ __device__ float warp_reduce_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}

static __inline__ __device__ float block_reduce_max(float v) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_reduce_max(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    v = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -FLT_MAX;
    if (wid == 0) v = warp_reduce_max(v);
    return v;
}

static __inline__ __device__ float block_reduce_sum(float v) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_reduce_sum(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    v = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) v = warp_reduce_sum(v);
    return v;
}

__global__ void fused_post_kernel(const float* __restrict__ y,
                                  float* __restrict__ out,
                                  int rows, int cols,
                                  float scale, float cmin, float cmax) {
    int row = blockIdx.x;
    if (row >= rows) return;
    const float* base = y + ((long long)row) * cols;

    float tmax = -FLT_MAX;
    for (int col = threadIdx.x; col < cols; col += blockDim.x) {
        float v = base[col] * scale;
        v = v + v;
        v = fminf(fmaxf(v, cmin), cmax);
        tmax = fmaxf(tmax, v);
    }
    float m = block_reduce_max(tmax);
    __syncthreads();

    float tsum = 0.0f;
    for (int col = threadIdx.x; col < cols; col += blockDim.x) {
        float v = base[col] * scale;
        v = v + v;
        v = fminf(fmaxf(v, cmin), cmax);
        tsum += expf(v - m);
    }
    float s = block_reduce_sum(tsum);

    if (threadIdx.x == 0) {
        float lse = m + logf(s);
        float sp;
        if (lse > 20.0f) {
            sp = lse;
        } else if (lse < -20.0f) {
            sp = expf(lse);
        } else {
            sp = log1pf(expf(lse));
        }
        float mish = lse * tanhf(sp);
        out[row] = lse * mish;
    }
}

torch::Tensor fused_post_forward(torch::Tensor y, double scale_factor, double clamp_min, double clamp_max) {
    TORCH_CHECK(y.is_cuda(), "input must be CUDA");
    TORCH_CHECK(y.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(y.dim() == 2, "input must be 2D");
    if (!y.is_contiguous()) y = y.contiguous();
    const int rows = (int)y.size(0);
    const int cols = (int)y.size(1);
    auto out = torch::empty({rows, 1}, y.options());
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_post_kernel<<<rows, threads, 0, stream>>>(
        y.data_ptr<float>(), out.data_ptr<float>(), rows, cols,
        (float)scale_factor, (float)clamp_min, (float)clamp_max);
    return out;
}
"""

_ext = load_inline(
    name="fused_linear_post_lse_mish_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_post_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_size, scale_factor, clamp_min, clamp_max):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(input_size, hidden_size)
        self.scale_factor = scale_factor
        self.clamp_min = clamp_min
        self.clamp_max = clamp_max

    def forward(self, x):
        y = F.linear(x, self.matmul.weight, self.matmul.bias)
        return _ext.fused_post_forward(y, float(self.scale_factor), float(self.clamp_min), float(self.clamp_max))