import torch
import math
from torch.utils.cpp_extension import load_inline

# 简化但高效的Softmax CUDA实现
softmax_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <float.h>

// 简化但高效的Softmax内核
__global__ void softmax_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int batch_size,
    int seq_length) {
    
    int batch_idx = blockIdx.x;
    int tid = threadIdx.x;
    
    // 每个block处理一个batch
    if (batch_idx >= batch_size) return;
    
    extern __shared__ float shared_mem[];
    float* shared_max = shared_mem;
    float* shared_sum = &shared_mem[blockDim.x];
    
    // 第一步：每个线程计算局部最大值
    float thread_max = -FLT_MAX;
    for (int i = tid; i < seq_length; i += blockDim.x) {
        float val = input[batch_idx * seq_length + i];
        thread_max = fmaxf(thread_max, val);
    }
    
    // 块内归约求最大值
    #pragma unroll
    for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
        thread_max = fmaxf(thread_max, __shfl_xor_sync(0xFFFFFFFF, thread_max, offset));
    }
    
    if (tid == 0) {
        shared_max[0] = thread_max;
    }
    __syncthreads();
    
    float row_max = shared_max[0];
    
    // 第二步：计算指数和
    float thread_sum = 0.0f;
    for (int i = tid; i < seq_length; i += blockDim.x) {
        float val = input[batch_idx * seq_length + i];
        float exp_val = expf(val - row_max);
        output[batch_idx * seq_length + i] = exp_val;
        thread_sum += exp_val;
    }
    
    // 块内归约求和
    #pragma unroll
    for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
        thread_sum += __shfl_xor_sync(0xFFFFFFFF, thread_sum, offset);
    }
    
    if (tid == 0) {
        shared_sum[0] = thread_sum;
    }
    __syncthreads();
    
    float row_sum = shared_sum[0];
    
    // 第三步：归一化
    for (int i = tid; i < seq_length; i += blockDim.x) {
        output[batch_idx * seq_length + i] /= row_sum;
    }
}"

torch::Tensor softmax_cuda_forward(torch::Tensor input, int dim) {
    // 检查输入维度
    if (input.dim() != 2) {
        throw std::runtime_error("Softmax currently only supports 2D tensors");
    }
    
    int batch_size = input.size(0);
    int seq_length = input.size(1);
    
    // 使用简化的内核
    dim3 blocks(batch_size);
    dim3 threads(256);  // 优化块大小
    size_t shared_mem_size = threads.x * 2 * sizeof(float);
    
    auto output = torch::zeros_like(input);
    
    softmax_kernel<<<blocks, threads, shared_mem_size>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        seq_length
    );
    
    return output;
}
"""

softmax_cpp_source = """
torch::Tensor softmax_cuda_forward(torch::Tensor input, int dim);
"""

# 编译内联CUDA代码
softmax_cuda = None
if torch.cuda.is_available() and torch.version.cuda is not None:
    try:
        softmax_cuda = load_inline(
            name="softmax_cuda",
            cpp_sources=softmax_cpp_source,
            cuda_sources=softmax_source,
            functions=["softmax_cuda_forward"],
            verbose=True
        )
    except Exception as e:
        print(f"[警告] Softmax CUDA 扩展构建失败，将回退到 PyTorch 实现。错误: {e}")

class ModelNew(torch.nn.Module):
    def __init__(self, dim=-1):
        super(ModelNew, self).__init__()
        self.dim = dim

    def forward(self, x):
        # 若扩展可用则调用 CUDA 内核，否则回退到 PyTorch softmax
        if softmax_cuda is not None:
            return softmax_cuda.softmax_cuda_forward(x, self.dim)
        else:
            return torch.softmax(x, self.dim)