import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a reverse cumulative sum operation along a specified dimension.

    Parameters:
        dim (int): The dimension along which to perform the reverse cumulative sum.
    """

    def __init__(self, dim):
        super(Model, self).__init__()
        self.dim = dim

    def forward(self, x):
        return torch.cumsum(x.flip(self.dim), dim=self.dim).flip(self.dim)

batch_size = 32768
input_shape = (32768,)
dim = 1

def get_inputs():
    return [torch.rand(batch_size, *input_shape)]

def get_init_inputs():
    return [dim]

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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor reverse_cumsum_forward(torch::Tensor x, int64_t dim);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cub/cub.cuh>
#include <cstdint>

// One CTA computes one row.  The row length in the benchmark is 32768 = 32*1024.
// We process the row from right to left in coalesced 1024-element tiles.  For each
// tile, scanning the reversed tile gives exactly the suffix sums within that tile;
// a scalar carry holds the sum of all tiles to its right.
template<int BLOCK_THREADS>
__global__ void reverse_cumsum_lastdim_kernel(const float* __restrict__ x,
                                              float* __restrict__ y,
                                              int64_t rows,
                                              int64_t cols) {
    const int64_t row = static_cast<int64_t>(blockIdx.x);
    if (row >= rows) return;

    const int tid = threadIdx.x;
    const int64_t base = row * cols;

    using BlockScan = cub::BlockScan<float, BLOCK_THREADS>;
    __shared__ typename BlockScan::TempStorage scan_storage;

    float carry = 0.0f;

    for (int64_t tile_end = cols; tile_end > 0; tile_end -= BLOCK_THREADS) {
        int64_t tile_start = tile_end - BLOCK_THREADS;
        if (tile_start < 0) tile_start = 0;
        const int count = static_cast<int>(tile_end - tile_start);

        // Reversed indexing: tid 0 corresponds to the rightmost element of the tile.
        const int64_t col = tile_end - 1 - tid;
        float v = 0.0f;
        if (tid < count) {
            v = x[base + col];
        }

        float scanned = 0.0f;
        float tile_sum = 0.0f;
        BlockScan(scan_storage).InclusiveSum(v, scanned, tile_sum);
        __syncthreads();  // scan_storage may be reused next loop iteration

        if (tid < count) {
            y[base + col] = scanned + carry;
        }
        carry += tile_sum;
    }
}

torch::Tensor reverse_cumsum_forward(torch::Tensor x, int64_t dim) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(x.dim() >= 1, "x must have at least one dimension");

    const int64_t ndim = x.dim();
    if (dim < 0) dim += ndim;
    TORCH_CHECK(dim == ndim - 1, "this optimized kernel supports cumsum over the last dimension");

    const int64_t cols = x.size(ndim - 1);
    const int64_t rows = x.numel() / cols;

    auto y = torch::empty_like(x);
    if (x.numel() == 0) return y;

    constexpr int BLOCK = 1024;
    dim3 block(BLOCK);
    dim3 grid(rows);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    reverse_cumsum_lastdim_kernel<BLOCK><<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), rows, cols);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

_ext = load_inline(
    name="reverse_cumsum_lastdim_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["reverse_cumsum_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)

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

    def forward(self, x):
        return _ext.reverse_cumsum_forward(x, self.dim)