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

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

__device__ __forceinline__ float hardmish_impl(float x) {
    if (x <= -3.0f) {
        return 0.0f;
    } else if (x >= 3.0f) {
        return x;
    } else {
        return x * (x + 3.0f) / 6.0f;
    }
}

__global__ void hardmish_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= size) return;
    output[idx] = hardmish_impl(input[idx]);
}

torch::Tensor hardmish_cuda(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 = 256;
    const int blocks = (total_size + threads_per_block - 1) / threads_per_block;

    hardmish_kernel<<<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;
}
"""

hardmish_cpp_source = """
torch::Tensor hardmish_cuda(torch::Tensor input);
"""

hardmish_module = load_inline(
    name="hardmish_final",
    cpp_sources=hardmish_cpp_source,
    cuda_sources=hardmish_source,
    functions=["hardmish_cuda"],
    extra_cuda_cflags=["-O2"],
    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.hardmish = hardmish_module.hardmish_cuda

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