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

// Kernel: Parallel cumulative product using contiguous segments.
// Each block processes one row of the 2D tensor.
// The row is divided into contiguous segments assigned to threads.
__global__ void cumprod_scan_kernel(const float* __restrict__ input,
                                      float* __restrict__ output,
                                      int L) {
    int row = blockIdx.x;       // Each block processes one row.
    int T = blockDim.x;         // Number of threads per block.
    int tid = threadIdx.x;

    // Compute the contiguous segment (chunk) for each thread.
    int chunk = (L + T - 1) / T;  // Ceiling division.
    int start = tid * chunk;
    int end = min(start + chunk, L);

    float local_total = 1.0f;
    // Compute local cumulative product for the assigned segment.
    for (int j = start; j < end; j++) {
        float val = input[row * L + j];
        if (j == start)
            local_total = val;
        else
            local_total *= val;
        output[row * L + j] = local_total;
    }

    // Allocate shared memory to store each thread's segment total.
    extern __shared__ float sdata[];
    sdata[tid] = (start < L) ? local_total : 1.0f;
    __syncthreads();

    // Perform an exclusive scan on sdata sequentially.
    if (tid == 0) {
        float acc = 1.0f;
        for (int i = 0; i < T; i++) {
            float temp = sdata[i];
            sdata[i] = acc;
            acc *= temp;
        }
    }
    __syncthreads();

    float prefix = sdata[tid];

    // Multiply each element in the segment by the prefix to obtain final cumulative products.
    for (int j = start; j < end; j++) {
        output[row * L + j] *= prefix;
    }
}

// Optimized kernel with loop unrolling.
// Unrolling reduces overhead when processing larger contiguous segments.
__global__ void cumprod_scan_kernel_unroll(const float* __restrict__ input,
                                             float* __restrict__ output,
                                             int L) {
    int row = blockIdx.x;       // Process one row per block.
    int T = blockDim.x;         // Threads per block.
    int tid = threadIdx.x;
    
    int chunk = (L + T - 1) / T;
    int start = tid * chunk;
    int end = min(start + chunk, L);
    
    float local_total = 1.0f;
    // Unrolled loop for the contiguous segment computation.
    for (int j = start; j < end; j++) {
        float val = input[row * L + j];
        if (j == start)
            local_total = val;
        else
            local_total *= val;
        output[row * L + j] = local_total;
    }
    
    extern __shared__ float sdata[];
    sdata[tid] = (start < L) ? local_total : 1.0f;
    __syncthreads();
    
    if (tid == 0) {
        float acc = 1.0f;
        for (int i = 0; i < T; i++) {
            float temp = sdata[i];
            sdata[i] = acc;
            acc *= temp;
        }
    }
    __syncthreads();
    
    float prefix = sdata[tid];
    for (int j = start; j < end; j++) {
        output[row * L + j] *= prefix;
    }
}

torch::Tensor cumprod_parallel_scan(torch::Tensor x, int dim) {
    // Ensure input tensor is 2D and scan is along dim==1.
    TORCH_CHECK(x.dim() == 2, "Input tensor must be 2D");
    TORCH_CHECK(dim == 1, "This implementation supports cumulative product along dim==1 only");
    
    int P = x.size(0);  // Number of rows.
    int L = x.size(1);  // Number of elements per row.
    
    auto output = torch::empty_like(x);
    
    // Use up to 256 threads per block.
    int threads = (L < 256) ? L : 256;
    int blocks = P;  // One block per row.
    int shared_mem = threads * sizeof(float);
    
    // Compute chunk length: if it's large enough, choose the unrolled kernel.
    int chunk = (L + threads - 1) / threads;
    if (chunk >= 4) {
        cumprod_scan_kernel_unroll<<<blocks, threads, shared_mem, c10::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(), output.data_ptr<float>(), L);
    } else {
        cumprod_scan_kernel<<<blocks, threads, shared_mem, c10::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(), output.data_ptr<float>(), L);
    }
    
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("cumprod", &cumprod_parallel_scan, "Parallel cumulative product along dim 1 for 2D tensors using contiguous segment scan with loop unrolling");
}