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

cpp_source = """
#include <torch/extension.h>

torch::Tensor logsumexp_cuda_forward(const torch::Tensor& input, int64_t dim, bool keepdim);
"""

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <limits>

#define BLOCK_SIZE 512

template <typename T>
__global__ void logsumexp_kernel(
    T* output,
    const T* input,
    const int n_along_dim,
    const int stride_between_slices,
    const int n_slices)
{{
    // 每个块处理一个 1D 切片
    const int slice_idx = blockIdx.x;
    if (slice_idx >= n_slices) return;

    __shared__ T sdata[BLOCK_SIZE];
    
    const int tid = threadIdx.x;
    const T* input_slice = input + slice_idx * stride_between_slices;

    // --- Pass 1: Find Max (`m`) in parallel ---
    T my_max = -std::numeric_limits<T>::infinity();
    for (int i = tid; i < n_along_dim; i += blockDim.x) {
        my_max = fmaxf(my_max, input_slice[i]);
    }
    sdata[tid] = my_max;
    __syncthreads();

    // Reduce max in shared memory
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
        }
        __syncthreads();
    }
    const T max_val = sdata[0];
    
    // --- Pass 2: Sum(exp(x - m)) in parallel ---
    T my_sum = 0.0f;
    for (int i = tid; i < n_along_dim; i += blockDim.x) {
        my_sum += expf(input_slice[i] - max_val);
    }
    sdata[tid] = my_sum;
    __syncthreads();

    // Reduce sum in shared memory
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // --- Finalize and Write ---
    if (tid == 0) {
        output[slice_idx] = max_val + logf(sdata[0]);
    }
}}

torch::Tensor logsumexp_cuda_forward(const torch::Tensor& input, int64_t dim, bool keepdim) {
    auto sizes = input.sizes();
    int ndim = input.dim();
    if (dim < 0) dim += ndim;

    TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");
    TORCH_CHECK(dim == ndim - 1, "This custom kernel only supports logsumexp on the last dimension.");

    const int n_along_dim = sizes[dim];
    const int64_t n_slices = input.numel() / n_along_dim;
    const int stride_between_slices = n_along_dim;
    
    // Calculate output shape
    auto out_sizes = sizes.vec();
    if (keepdim) {
        out_sizes[dim] = 1;
    } else {
        out_sizes.erase(out_sizes.begin() + dim);
    }
    auto output = torch::empty(out_sizes, input.options());
    
    const dim3 grid(n_slices);
    const dim3 block(BLOCK_SIZE);

    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "logsumexp_kernel", ([&] {{
        logsumexp_kernel<scalar_t><<<grid, block>>>(
            output.data_ptr<scalar_t>(),
            input.data_ptr<scalar_t>(),
            n_along_dim,
            stride_between_slices,
            n_slices
        );
    }}));

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, dim, keepdim=False):
        super(ModelNew, self).__init__()
        self.dim = dim
        self.keepdim = keepdim
        self.op = load_inline(
            name='logsumexp_op_v2',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['logsumexp_cuda_forward'],
            verbose=False
        )

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return self.op.logsumexp_cuda_forward(input_tensor, self.dim, self.keepdim)