import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a Softsign activation.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Softsign activation to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of any shape.

        Returns:
            torch.Tensor: Output tensor with Softsign applied, same shape as input.
        """
        return x / (1 + torch.abs(x))

batch_size = 4096
dim = 393216

def get_inputs():
    x = torch.rand(batch_size, dim)
    return [x]

def get_init_inputs():
    return []  # No special initialization inputs needed

import torch
from torch.utils.cpp_extension import load_inline

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

template <typename scalar_t>
__global__ void softsign_kernel(const scalar_t* __restrict__ input,
                                scalar_t* __restrict__ output,
                                int64_t n) {
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = idx; i < n; i += stride) {
        scalar_t x = input[i];
        output[i] = x / (scalar_t(1.0) + fabsf(x));
    }
}

// Vectorized version using float4
template <>
__global__ void softsign_kernel<float>(const float* __restrict__ input,
                                       float* __restrict__ output,
                                       int64_t n) {
    int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    int64_t stride = (int64_t)blockDim.x * gridDim.x * 4;
    for (int64_t i = idx; i < n; i += stride) {
        float4 x4 = reinterpret_cast<const float4*>(input)[i / 4];
        float4 y4;
        y4.x = x4.x / (1.0f + fabsf(x4.x));
        y4.y = x4.y / (1.0f + fabsf(x4.y));
        y4.z = x4.z / (1.0f + fabsf(x4.z));
        y4.w = x4.w / (1.0f + fabsf(x4.w));
        reinterpret_cast<float4*>(output)[i / 4] = y4;
    }
}

torch::Tensor softsign_forward(torch::Tensor x) {
    auto y = torch::empty_like(x);
    int64_t n = x.numel();
    const int threads = 256;
    const int blocks = 2048;
    
    if (x.dtype() == torch::kFloat) {
        int64_t vec_n = n / 4 * 4;
        if (vec_n > 0) {
            softsign_kernel<float><<<blocks, threads>>>(
                x.data_ptr<float>(), y.data_ptr<float>(), vec_n);
        }
        if (n > vec_n) {
            softsign_kernel<float><<<1, threads>>>(
                x.data_ptr<float>() + vec_n, y.data_ptr<float>() + vec_n, n - vec_n);
        }
    } else if (x.dtype() == torch::kHalf) {
        int64_t vec_n = n / 8 * 8;
        softsign_kernel<at::Half><<<blocks, threads>>>(
            x.data_ptr<at::Half>(), y.data_ptr<at::Half>(), n);
    } else {
        AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16,
            x.scalar_type(), "softsign", [&]() {
                softsign_kernel<scalar_t><<<blocks, threads>>>(
                    x.data_ptr<scalar_t>(), y.data_ptr<scalar_t>(), n);
            });
    }
    return y;
}
"""

module = load_inline(
    name="softsign_cuda",
    cpp_sources="torch::Tensor softsign_forward(torch::Tensor x);",
    cuda_sources=cuda_source,
    functions=["softsign_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return module.softsign_forward(x)