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

# Mish激活函数的CUDA实现
mish_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

// Mish函数实现
__device__ __forceinline__ float mish_impl(float x) {
    float softplus_val = (x > 20.0f) ? x : ((x < -20.0f) ? expf(x) : log1pf(expf(x)));
    float tanh_val = tanhf(softplus_val);
    return x * tanh_val;
}

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

torch::Tensor mish_cuda(torch::Tensor input) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.dtype() == torch::kFloat32, "input must be float32");
    
    input = input.contiguous();
    const 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;
    
    mish_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;
}
"""

mish_cpp_source = """
torch::Tensor mish_cuda(torch::Tensor input);
"""

mish_module = load_inline(
    name="mish_final",
    cpp_sources=mish_cpp_source,
    cuda_sources=mish_source,
    functions=["mish_cuda"],
    extra_cuda_cflags=["-O2"],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, in_features):
        super(ModelNew, self).__init__()
        torch.manual_seed(42)
        self.linear = nn.Linear(in_features, in_features)
        self.mish = mish_module.mish_cuda
    
    def forward(self, x):
        x = self.linear(x)
        return self.mish(x)
