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/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>

void post_lse_mish_bias_launcher(const torch::Tensor& y, const torch::Tensor& bias, torch::Tensor& out, float scale2, float cmin, float cmax);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale_factor, double clamp_min, double clamp_max) {
    const int64_t B64 = x.size(0), K64 = x.size(1), H64 = weight.size(0);
    auto y = torch::empty({B64, H64}, x.options());

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, at::cuda::getCurrentCUDAStream());
    cublasMath_t old_mode;
    cublasGetMathMode(handle, &old_mode);
    cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH);

    float alpha = 1.0f, beta = 0.0f;
    // Row-major y[B,H] is the same storage as column-major C[H,B].
    // weight[H,K] row-major is column-major A[K,H], so op(A)=T gives HxK.
    // x[B,K] row-major is column-major B[K,B].
    cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N,
                (int)H64, (int)B64, (int)K64,
                &alpha,
                weight.data_ptr<float>(), (int)K64,
                x.data_ptr<float>(), (int)K64,
                &beta,
                y.data_ptr<float>(), (int)H64);
    cublasSetMathMode(handle, old_mode);

    auto out = torch::empty({B64, 1}, x.options());
    post_lse_mish_bias_launcher(y, bias, out, (float)(2.0 * scale_factor), (float)clamp_min, (float)clamp_max);
    return out;
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <math.h>

__global__ void post_lse_mish_bias_kernel(const float* __restrict__ y,
                                          const float* __restrict__ bias,
                                          float* __restrict__ out,
                                          int H, float scale2, float cmin, float cmax) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const float* base = y + (size_t)row * H;

    float local_max = -3.402823466e+38F;
    for (int j = tid; j < H; j += blockDim.x) {
        float v = (base[j] + bias[j]) * scale2;
        v = fminf(fmaxf(v, cmin), cmax);
        local_max = fmaxf(local_max, v);
    }

    __shared__ float smem[256];
    smem[tid] = local_max;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        __syncthreads();
    }
    float m = smem[0];

    float local_sum = 0.0f;
    for (int j = tid; j < H; j += blockDim.x) {
        float v = (base[j] + bias[j]) * scale2;
        v = fminf(fmaxf(v, cmin), cmax);
        local_sum += expf(v - m);
    }

    smem[tid] = local_sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        float z = m + logf(smem[0]);
        float sp = log1pf(expf(z));
        out[row] = z * (z * tanhf(sp));
    }
}

void post_lse_mish_bias_launcher(const torch::Tensor& y, const torch::Tensor& bias, torch::Tensor& out, float scale2, float cmin, float cmax) {
    int B = (int)y.size(0);
    int H = (int)y.size(1);
    const int threads = 256;
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    post_lse_mish_bias_kernel<<<B, threads, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), H, scale2, cmin, cmax);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}
'''

_ext = load_inline(
    name='fused_lse_mish_tf32_final',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['fused_forward'],
    verbose=False,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_ldflags=['-lcublas'],
)

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),
        )