import torch
from torch.utils.cpp_extension import load_inline

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

// VECTORIZED版本：4元素向量化处理（最终优化版本）
__global__ void logbeta_vectorized_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ z,
    int size
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    
    // 每次处理4个元素
    int vec_size = size / 4;
    
    for (int i = tid; i < vec_size; i += stride) {
        int base_idx = i * 4;
        
        // 加载4个元素对
        float x0 = x[base_idx + 0];
        float x1 = x[base_idx + 1];
        float x2 = x[base_idx + 2];
        float x3 = x[base_idx + 3];
        
        float y0 = y[base_idx + 0];
        float y1 = y[base_idx + 1];
        float y2 = y[base_idx + 2];
        float y3 = y[base_idx + 3];
        
        // 计算4个log beta值（使用标准库函数确保精度）
        z[base_idx + 0] = lgammaf(x0) + lgammaf(y0) - lgammaf(x0 + y0);
        z[base_idx + 1] = lgammaf(x1) + lgammaf(y1) - lgammaf(x1 + y1);
        z[base_idx + 2] = lgammaf(x2) + lgammaf(y2) - lgammaf(x2 + y2);
        z[base_idx + 3] = lgammaf(x3) + lgammaf(y3) - lgammaf(x3 + y3);
    }
    
    // 处理剩余元素
    int remainder = size % 4;
    if (tid == 0 && remainder > 0) {
        int start_idx = vec_size * 4;
        for (int i = start_idx; i < size; i++) {
            float xi = x[i];
            float yi = y[i];
            z[i] = lgammaf(xi) + lgammaf(yi) - lgammaf(xi + yi);
        }
    }
}

torch::Tensor logbeta_cuda(
    torch::Tensor x,
    torch::Tensor y,
    std::string mode = "vectorized"
) {
    auto x_contig = x.contiguous();
    auto y_contig = y.contiguous();
    auto z = torch::empty_like(x_contig);
    int size = x_contig.numel();
    
    if (mode == "vectorized") {
        // 向量化版本
        const int block_size = 256;
        int num_blocks = (size / 4 + block_size - 1) / block_size;
        int grid_size = std::min(num_blocks, 65535);
        
        logbeta_vectorized_kernel<<<grid_size, block_size>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            z.data_ptr<float>(),
            size
        );
    }
    
    return z;
}
"""

logbeta_cpp_source = """
torch::Tensor logbeta_cuda(torch::Tensor x, torch::Tensor y, std::string mode);
"""

# 编译CUDA代码
logbeta = load_inline(
    name="logbeta_vectorized",
    cpp_sources=logbeta_cpp_source,
    cuda_sources=logbeta_source,
    functions=["logbeta_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-std=c++17",
        "-maxrregcount=64"
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, mode="vectorized"):
        super(ModelNew, self).__init__()
        self.mode = mode
        self.logbeta = logbeta  # The module containing the kernel

    def forward(self, x, y):
        return self.logbeta.logbeta_cuda(x, y, self.mode)
