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

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

__device__ __forceinline__ float sigmoid(float x) {
    // 针对1.702系数优化的sigmoid计算
    return 1.0f / (1.0f + expf(-1.702f * x));
}

__global__ void swiglu_kernel(
    const float* __restrict__ x1,  // 主分支
    const float* __restrict__ x2,  // 门控分支
    float* __restrict__ output,
    const int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= size) return;
    output[idx] = x1[idx] * sigmoid(x2[idx]);
}

torch::Tensor swiglu_cuda(torch::Tensor x1, torch::Tensor x2) {
    TORCH_CHECK(x1.is_cuda() && x2.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x1.dtype() == torch::kFloat32 && x2.dtype() == torch::kFloat32, "must be float32");
    TORCH_CHECK(x1.sizes() == x2.sizes(), "x1 and x2 must have same shape");

    x1 = x1.contiguous();
    x2 = x2.contiguous();
    const int total_size = x1.numel();
    auto output = torch::empty_like(x1);

    const int threads_per_block = 256;
    const int blocks = (total_size + threads_per_block - 1) / threads_per_block;

    swiglu_kernel<<<blocks, threads_per_block>>>(
        x1.data_ptr<float>(),
        x2.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;
}
"""

swiglu_cpp_source = """
torch::Tensor swiglu_cuda(torch::Tensor x1, torch::Tensor x2);
"""

swiglu_module = load_inline(
    name="swiglu_final",
    cpp_sources=swiglu_cpp_source,
    cuda_sources=swiglu_source,
    functions=["swiglu_cuda"],
    extra_cuda_cflags=["-O2"],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, in_features):
        super().__init__()
        torch.manual_seed(42)
        self.linear1 = nn.Linear(in_features, in_features)
        self.linear2 = nn.Linear(in_features, in_features)
        self.swiglu = swiglu_module.swiglu_cuda

    def forward(self, x):
        x1 = self.linear1(x)
        x2 = self.linear2(x)
        return self.swiglu(x1, x2)