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


UNROLL_FACTOR = 4

tanhshrink_source_optimized = f"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

__device__ __forceinline__ float tanhshrink_impl(float x) {{
    return x - tanhf(x);
}}


__global__ void tanhshrink_kernel_unroll{UNROLL_FACTOR}(
    const float* __restrict__ input,
    float* __restrict__ output,
    int size
) {{
   
    int idx = blockIdx.x * blockDim.x * {UNROLL_FACTOR} + threadIdx.x;

    #pragma unroll
    for (int i = 0; i < {UNROLL_FACTOR}; ++i) {{
        int current_idx = idx + i * blockDim.x;
        if (current_idx < size) {{
            output[current_idx] = tanhshrink_impl(input[current_idx]);
        }}
    }}
}}

torch::Tensor tanhshrink_cuda_optimized(torch::Tensor input) {{
    TORCH_CHECK(input.is_cuda(), "input must be CUDA tensor");
    TORCH_CHECK(input.dtype() == torch::kFloat32, "input must be float32");

    input = input.contiguous();
    int total_size = input.numel();
    auto output = torch::empty_like(input);

    const int threads_per_block = 512; 
    
    const int blocks = (total_size + threads_per_block * {UNROLL_FACTOR} - 1) / (threads_per_block * {UNROLL_FACTOR});

    tanhshrink_kernel_unroll{UNROLL_FACTOR}<<<blocks, threads_per_block>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        total_size
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {{
        throw std::runtime_error("CUDA error: " + std::string(cudaGetErrorString(err)));
    }}
    return output;
}}
"""

tanhshrink_cpp_source_optimized = """
torch::Tensor tanhshrink_cuda_optimized(torch::Tensor input);
"""

tanhshrink_module_optimized = load_inline(
    name="tanhshrink_optimized",
    cpp_sources=tanhshrink_cpp_source_optimized,
    cuda_sources=tanhshrink_source_optimized,
    functions=["tanhshrink_cuda_optimized"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, in_features):
        super().__init__()
        torch.manual_seed(42)
        self.linear = nn.Linear(in_features, in_features)
        self.tanhshrink = tanhshrink_module_optimized.tanhshrink_cuda_optimized

    def forward(self, x):
        x = self.linear(x)
        return self.tanhshrink(x)