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


B, C = 128, 512
BLOCK_SIZE = 256


class ModelNew(nn.Module):
    def __init__(self, block_size=256):
        super().__init__()
        self.block_size = block_size
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor softsign_fused_cuda(torch::Tensor input);
        """

        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <cmath> 

        #define BLOCK_SIZE {self.block_size}

        __global__ void softsign_fused_kernel(
            const float* __restrict__ input_data,
            float* __restrict__ output_data,
            int N_total
        ) {{
            for (int i = blockIdx.x * blockDim.x + threadIdx.x; 
                 i < N_total; 
                 i += gridDim.x * blockDim.x) 
            {{
                const float x = input_data[i];
                output_data[i] = x / (1.0f + fabsf(x));
            }}
        }}

        torch::Tensor softsign_fused_cuda(torch::Tensor input) {{
            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
            input = input.contiguous();

            const int64_t N_total = input.numel();
            auto output = torch::empty_like(input);

            dim3 block_dim(BLOCK_SIZE);
            dim3 grid_dim((N_total + BLOCK_SIZE - 1) / BLOCK_SIZE);

            softsign_fused_kernel<<<grid_dim, block_dim>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                N_total
            );

            return output;
        }}
        """

        self.op = load_inline(
            name='softsign_opt',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['softsign_fused_cuda'],
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if not input.is_cuda: input = input.cuda()
        return self.op.softsign_fused_cuda(input.contiguous())