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
from torch.utils.cpp_extension import load_inline

cpp_src = r"""
#include <torch/extension.h>
#include <ATen/ops/linear.h>

torch::Tensor fused_post_cuda(torch::Tensor input, double scale_factor, double clamp_min, double clamp_max);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                            double scale_factor, double clamp_min, double clamp_max) {
    auto y = at::linear(x, weight, bias);
    return fused_post_cuda(y, scale_factor, clamp_min, clamp_max);
}
"""

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

__device__ __forceinline__ 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;
}

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

template<int BLOCK>
__global__ void fused_post_kernel(const float* __restrict__ inp, float* __restrict__ out,
                                  int rows, int cols, float s, float cmin, float cmax) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const float* base = inp + (long long)row * cols;

    float m = -3.402823466e38F;
    #pragma unroll 4
    for (int j = tid; j < cols; j += BLOCK) {
        float v = base[j] * s;
        v = fminf(fmaxf(v, cmin), cmax);
        m = fmaxf(m, v);
    }

    m = warp_reduce_max(m);
    __shared__ float smem[32];
    if ((tid & 31) == 0) smem[tid >> 5] = m;
    __syncthreads();

    if (tid < 32) {
        m = (tid < (BLOCK + 31) / 32) ? smem[tid] : -3.402823466e38F;
        m = warp_reduce_max(m);
        if (tid == 0) smem[0] = m;
    }
    __syncthreads();
    m = smem[0];

    float sum = 0.0f;
    #pragma unroll 4
    for (int j = tid; j < cols; j += BLOCK) {
        float v = base[j] * s;
        v = fminf(fmaxf(v, cmin), cmax);
        sum += __expf(v - m);
    }

    sum = warp_reduce_sum(sum);
    if ((tid & 31) == 0) smem[tid >> 5] = sum;
    __syncthreads();

    if (tid < 32) {
        sum = (tid < (BLOCK + 31) / 32) ? smem[tid] : 0.0f;
        sum = warp_reduce_sum(sum);
        if (tid == 0) {
            float lse = __logf(sum) + m;
            float sp = __logf(1.0f + __expf(lse));
            out[row] = lse * (lse * tanhf(sp));
        }
    }
}

torch::Tensor fused_post_cuda(torch::Tensor input, double scale_factor, double clamp_min, double clamp_max) {
    auto in = input.contiguous();
    int rows = (int)in.size(0);
    int cols = (int)in.size(1);
    auto out = torch::empty({rows, 1}, in.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_post_kernel<256><<<rows, 256, 0, stream>>>(
        in.data_ptr<float>(), out.data_ptr<float>(), rows, cols,
        (float)(scale_factor * 2.0), (float)clamp_min, (float)clamp_max);
    return out;
}
"""

_ext = load_inline(
    name="fused_forward_lse_mish_submit",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["fused_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_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):
        return _ext.fused_forward(
            x, self.matmul.weight, self.matmul.bias,
            float(self.scale_factor), float(self.clamp_min), float(self.clamp_max)
        )