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

# Swish + LayerNorm融合的CUDA实现
swish_layernorm_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

__device__ __forceinline__ float fast_sigmoid(float x) {
    return 1.0f / (1.0f + expf(-x));
}

__device__ __forceinline__ float swish(float x) {
    return x * fast_sigmoid(x);
}

__global__ void swish_layernorm_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int batch_size,
    const int seq_len,
    const int hidden_size,
    const float eps
) {
    int batch_idx = blockIdx.x;
    int seq_idx = blockIdx.y;
    
    if (batch_idx >= batch_size || seq_idx >= seq_len) return;

    int token_start = (batch_idx * seq_len + seq_idx) * hidden_size;

    extern __shared__ float sdata[];

    float sum = 0.0f;
    float sum_sq = 0.0f;
    
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        int idx = token_start + i;
        float val = input[idx];
        float swish_val = swish(val);

        output[idx] = swish_val;
        
        sum += swish_val;
        sum_sq += swish_val * swish_val;
    }
    
    sdata[threadIdx.x] = sum;
    sdata[threadIdx.x + blockDim.x] = sum_sq;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            sdata[threadIdx.x] += sdata[threadIdx.x + stride];
            sdata[threadIdx.x + blockDim.x] += sdata[threadIdx.x + stride + blockDim.x];
        }
        __syncthreads();
    }
    
    float mean = sdata[0] / hidden_size;
    float var = sdata[blockDim.x] / hidden_size - mean * mean;
    float inv_std = rsqrtf(var + eps);
    
    // 第二遍：应用LayerNorm
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        int idx = token_start + i;
        float normalized = (output[idx] - mean) * inv_std;
        output[idx] = normalized * weight[i] + bias[i];
    }
}

torch::Tensor swish_layernorm_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias,
    float eps = 1e-5f
) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
    TORCH_CHECK(input.dtype() == torch::kFloat32, "input must be float32");
    TORCH_CHECK(input.dim() == 3, "input must be a 3D tensor");
    
    const int batch_size = input.size(0);
    const int seq_len = input.size(1);
    const int hidden_size = input.size(2);
    
    auto output = torch::empty_like(input);
    
    // 使用warp大小（32）作为线程块大小
    const int threads_per_block = 256;
    const int shared_mem_size = threads_per_block * 2 * sizeof(float);
    
    dim3 blocks(batch_size, seq_len);
    
    swish_layernorm_kernel<<<blocks, threads_per_block, shared_mem_size>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        seq_len,
        hidden_size,
        eps
    );
    
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error("CUDA error: " + std::string(cudaGetErrorString(err)));
    }
    
    return output;
}
"""

swish_layernorm_cpp_source = """
torch::Tensor swish_layernorm_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, float eps);
"""

swish_layernorm_module = load_inline(
    name="swish_layernorm_final",
    cpp_sources=swish_layernorm_cpp_source,
    cuda_sources=swish_layernorm_source,
    functions=["swish_layernorm_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, hidden_size):
        super(ModelNew, self).__init__()
        torch.manual_seed(42)
        self.hidden_size = hidden_size
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        self.eps = 1e-5
        self.swish_layernorm = swish_layernorm_module.swish_layernorm_cuda
    
    def forward(self, x):
        return self.swish_layernorm(x, self.weight, self.bias, self.eps)
