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

# C++ 源代码
cpp_source = """
#include <torch/extension.h>

torch::Tensor cumsum_cuda_forward(const torch::Tensor& input, int64_t dim);
"""

# CUDA 源代码
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

// 固定配置
#define BLOCK_SIZE 1024
// 支持的最大长度 2048 (每个线程处理 2 个元素)
#define MAX_SEQ_LEN 2048

template <typename T>
__global__ void cumsum_hillis_steele_kernel(
    T* __restrict__ output,
    const T* __restrict__ input,
    const int n,
    const int num_rows)
{
    int row_idx = blockIdx.x;
    if (row_idx >= num_rows) return;

    int tid = threadIdx.x;
    
    using AccT = double;

    // 双缓冲防止读写竞争
    __shared__ AccT sdata[2][MAX_SEQ_LEN];

    const T* row_in = input + row_idx * n;
    T* row_out = output + row_idx * n;

    // 每个线程负责加载和处理 2 个元素
    int idx1 = tid;
    int idx2 = tid + BLOCK_SIZE;

    // 1. Load Global -> Shared (Buffer 0)
    if (idx1 < n) sdata[0][idx1] = static_cast<AccT>(row_in[idx1]);
    else sdata[0][idx1] = 0.0;

    if (idx2 < n) sdata[0][idx2] = static_cast<AccT>(row_in[idx2]);
    else sdata[0][idx2] = 0.0;

    __syncthreads();

    // 2. Hillis-Steele Inclusive Scan
    // 每次迭代 stride 翻倍，读 in_buf，写 out_buf
    int in_buf = 0;
    int out_buf = 1;

    for (int stride = 1; stride < MAX_SEQ_LEN; stride *= 2) {
        
        // 处理第一个元素 idx1
        if (idx1 >= stride) {
            sdata[out_buf][idx1] = sdata[in_buf][idx1] + sdata[in_buf][idx1 - stride];
        } else {
            sdata[out_buf][idx1] = sdata[in_buf][idx1];
        }

        // 处理第二个元素 idx2
        if (idx2 >= stride) {
            sdata[out_buf][idx2] = sdata[in_buf][idx2] + sdata[in_buf][idx2 - stride];
        } else {
            sdata[out_buf][idx2] = sdata[in_buf][idx2];
        }

        __syncthreads();

        // 交换缓冲区
        in_buf = 1 - in_buf;
        out_buf = 1 - out_buf;
    }

    // 3. Write Back
    if (idx1 < n) row_out[idx1] = static_cast<T>(sdata[in_buf][idx1]);
    if (idx2 < n) row_out[idx2] = static_cast<T>(sdata[in_buf][idx2]);
}

torch::Tensor cumsum_cuda_forward(const torch::Tensor& input, int64_t dim) {
    TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");
    
    int ndim = input.dim();
    if (dim < 0) dim += ndim;
    TORCH_CHECK(dim == ndim - 1, "Optimized only for the last dimension");

    int n = input.size(dim);
    int num_rows = input.numel() / n;

    // 确保 dim 长度不超过 Kernel 支持范围
    TORCH_CHECK(n <= MAX_SEQ_LEN, "Seq len exceeds limit (2048)");

    auto output = torch::empty_like(input);

    dim3 grid(num_rows);
    dim3 block(BLOCK_SIZE);

    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "cumsum_hillis_steele", ([&] {
        cumsum_hillis_steele_kernel<scalar_t><<<grid, block>>>(
            output.data_ptr<scalar_t>(),
            input.data_ptr<scalar_t>(),
            n,
            num_rows
        );
    }));

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, dim):
        super(ModelNew, self).__init__()
        self.dim = dim
        self.op = load_inline(
            name='cumsum',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['cumsum_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=['-O3','--use_fast_math']
        )

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