import torch
from torch.utils.cpp_extension import load_inline

logsumexp_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#define BLOCK_SIZE 256

// 融合的LogSumExp Kernel，处理任意长度的行
__global__ void logsumexp_kernel(const float* input, float* output, int rows, int cols) {
    extern __shared__ float sdata[];

    int row = blockIdx.x;
    int tid = threadIdx.x;

    // --- 第一阶段：在共享内存中找最大值 ---
    float max_val = -INFINITY;
    for (int i = tid; i < cols; i += blockDim.x) {
        max_val = fmaxf(max_val, input[row * cols + i]);
    }
    sdata[tid] = max_val;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
        }
        __syncthreads();
    }
    float row_max = sdata[0];
    __syncthreads(); // 确保所有线程都拿到了row_max

    // --- 第二阶段：在共享内存中求和 ---
    float sum = 0.0f;
    for (int i = tid; i < cols; i += blockDim.x) {
        sum += expf(input[row * cols + i] - row_max);
    }
    sdata[tid] = sum;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    if (tid == 0) {
        float row_sum = sdata[0];
        // 最终结果 log(sum) + max
        output[row] = logf(row_sum) + row_max;
    }
}

torch::Tensor logsumexp_cuda(torch::Tensor input) {
    const auto rows = input.size(0);
    const auto cols = input.size(1);

    auto output = torch::empty({rows}, input.options());

    dim3 grid_size(rows);
    dim3 block_size(BLOCK_SIZE);
    // 共享内存大小需要能容纳BLOCK_SIZE个float
    size_t shared_mem_size = BLOCK_SIZE * sizeof(float);

    logsumexp_kernel<<<grid_size, block_size, shared_mem_size>>>(
        input.data_ptr<float>(), output.data_ptr<float>(), rows, cols);

    return output;
}
"""

logsumexp_cpp_source = """
torch::Tensor logsumexp_cuda(torch::Tensor input);
"""

# 编译CUDA代码
logsumexp = load_inline(
    name="logsumexp_final",
    cpp_sources=logsumexp_cpp_source,
    cuda_sources=logsumexp_source,
    functions=["logsumexp_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.logsumexp = logsumexp

    def forward(self, x):
        return self.logsumexp.logsumexp_cuda(x)
