import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a gemm, swish, divide, clamp, tanh, and clamp operations.
    """
    def __init__(self, in_features, out_features, bias=True):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=bias)

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_features).
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_features).
        """
        x = self.gemm(x)
        x = x * torch.sigmoid(x)  # Swish activation
        x = x / 2.0
        x = torch.clamp(x, min=-1.0, max=1.0)  # Clamp between -1 and 1
        x = torch.tanh(x)  # Tanh activation
        x = torch.clamp(x, min=-1.0, max=1.0)  # Clamp between -1 and 1
        return x

batch_size = 1024
in_features = 8192
out_features = 8192

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

def get_init_inputs():
    return [in_features, out_features]

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

cpp_sources = r"""
#include <torch/extension.h>
torch::Tensor swish_clamp_tanh_inplace(torch::Tensor input);
"""

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

__device__ __forceinline__ float tanh_pade(float x) {
    float x2 = x * x;
    float num = ((x2 + 378.0f) * x2 + 17325.0f) * x2 + 135135.0f;
    float den = ((28.0f * x2 + 3150.0f) * x2 + 62370.0f) * x2 + 135135.0f;
    return x * (num / den);
}

__device__ __forceinline__ float activate_val(float x) {
    float s = 1.0f / (1.0f + expf(-x));
    float y = 0.5f * x * s;
    y = fminf(1.0f, fmaxf(-1.0f, y));
    return tanh_pade(y);
}

__global__ void act_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) {
        data[i] = activate_val(data[i]);
    }
}

torch::Tensor swish_clamp_tanh_inplace(torch::Tensor input) {
    const int64_t n = input.numel();
    if (n == 0) return input;
    const int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    act_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(input.data_ptr<float>(), n);
    return input;
}
"""

_ext = load_inline(
    name="swish_clamp_tanh_ext_final",
    cpp_sources=cpp_sources,
    cuda_sources=cuda_sources,
    functions=["swish_clamp_tanh_inplace"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=bias)

    def forward(self, x):
        y = F.linear(x, self.gemm.weight, self.gemm.bias)
        return _ext.swish_clamp_tanh_inplace(y)