#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>

constexpr int MAX_CHUNK = 32;  // Maximum number of elements a thread processes in its chunk

// Kernel to perform an inclusive scan (cumulative sum) for each row.
// Each block processes one row (of length 'scan_length'), and each thread computes
// a local prefix sum over its assigned contiguous chunk. Intra-warp scan is performed
// using __shfl_up_sync with loop unrolling, and an inter-warp exclusive scan is computed using shared memory.
template <int CHUNK_SIZE>
__global__ void cumsum_kernel(const float* __restrict__ input,
                              float* __restrict__ output,
                              int scan_length) {
    // Each block handles one row.
    int row = blockIdx.x;
    int T = blockDim.x;    // Number of threads per block.
    int tid = threadIdx.x;
    
    // Compute per-thread chunk size (ceiling division).
    int chunk = (scan_length + T - 1) / T;
    int start = tid * chunk;
    int end = start + chunk;
    if (end > scan_length)
        end = scan_length;
    
    // Pointers to the current row.
    const float* in_row = input + row * scan_length;
    float* out_row = output + row * scan_length;
    
    // Compute local sequential prefix sum.
    float local_prefix[CHUNK_SIZE];  // CHUNK_SIZE must be >= chunk.
    float local_sum = 0.0f;
    int count = 0;
    #pragma unroll
    for (int i = start; i < end; i++, count++) {
        // Use __ldg for read-only caching.
        float val = __ldg(&in_row[i]);
        local_sum += val;
        local_prefix[count] = local_sum;
    }
    float local_total = (count > 0) ? local_prefix[count - 1] : 0.0f;
    
    // --- Intra-warp scan using warp-level intrinsics with loop unrolling ---
    int lane = tid & 31;  // Lane index within the warp.
    float warp_inclusive = local_total;
    #pragma unroll
    for (int offset = 1; offset < 32; offset *= 2) {
        float n = __shfl_up_sync(0xffffffff, warp_inclusive, offset);
        if (lane >= offset)
            warp_inclusive += n;
    }
    // Calculate intra-warp exclusive prefix.
    float intra_offset = (lane == 0) ? 0.0f : (warp_inclusive - local_total);
    
    // --- Inter-warp scan using shared memory ---
    extern __shared__ float sdata[]; // Shared memory: one float per warp.
    int warp_id = tid >> 5;  // Warp index.
    int num_warps = blockDim.x / 32;
    if (lane == 31) {
        sdata[warp_id] = warp_inclusive;
    }
    __syncthreads();
    
    // First warp computes an exclusive scan over warp sums.
    if (tid < num_warps) {
        float offset = 0.0f;
        #pragma unroll
        for (int i = 0; i < tid; i++) {
            offset += sdata[i];
        }
        sdata[tid] = offset;
    }
    __syncthreads();
    float warp_offset = (warp_id > 0) ? sdata[warp_id] : 0.0f;
    
    // Combine inter-warp and intra-warp offsets.
    float total_offset = warp_offset + intra_offset;
    
    // --- Write back the final cumulative sum ---
    #pragma unroll
    for (int i = start, j = 0; i < end; i++, j++) {
        out_row[i] = local_prefix[j] + total_offset;
    }
}

torch::Tensor cumsum(torch::Tensor x, int scan_length) {
    // Ensure we run on the correct CUDA device.
    c10::cuda::CUDAGuard device_guard(x.device());
    
    // Retrieve original shape.
    auto sizes = x.sizes();
    int total_elements = x.numel();
    // Compute the "outer" dimension as product of all dimensions except the scan dimension.
    int outer = total_elements / scan_length;
    
    // Flatten x to a 2D tensor (outer, scan_length).
    auto x_2d = x.view({outer, scan_length});
    
    // Allocate the output tensor.
    auto output = torch::empty({outer, scan_length}, x.options());
    
    // Launch parameters: use 256 threads per block (multiple of 32).
    const int threads = 256;
    dim3 grid(outer);
    dim3 block(threads);
    // Allocate shared memory: one float for each warp.
    size_t shared_mem_size = (threads / 32) * sizeof(float);
    
    constexpr int CHUNK_SIZE = MAX_CHUNK;  // Template parameter.
    cumsum_kernel<CHUNK_SIZE><<<grid, block, shared_mem_size, c10::cuda::getCurrentCUDAStream()>>>(
        x_2d.data_ptr<float>(),
        output.data_ptr<float>(),
        scan_length
    );
    
    // Reshape the output back to the original tensor shape.
    return output.view(sizes);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("cumsum", &cumsum, "Custom cumulative sum (inclusive scan) along the specified dimension (CUDA)");
}