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_source = r'''
#include <torch/extension.h>

torch::Tensor reverse_cumsum_cuda(torch::Tensor x, int64_t dim);
'''

cuda_source = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

__inline__ __device__ float warp_inclusive_scan(float v) {
    unsigned mask = 0xffffffffu;
    float n;
    n = __shfl_up_sync(mask, v, 1);  if ((threadIdx.x & 31) >= 1)  v += n;
    n = __shfl_up_sync(mask, v, 2);  if ((threadIdx.x & 31) >= 2)  v += n;
    n = __shfl_up_sync(mask, v, 4);  if ((threadIdx.x & 31) >= 4)  v += n;
    n = __shfl_up_sync(mask, v, 8);  if ((threadIdx.x & 31) >= 8)  v += n;
    n = __shfl_up_sync(mask, v, 16); if ((threadIdx.x & 31) >= 16) v += n;
    return v;
}

__global__ __launch_bounds__(1024, 1)
void reverse_cumsum_warpchunks_kernel(const float* __restrict__ x, float* __restrict__ y) {
    __shared__ float chunk_sums[32];
    __shared__ float chunk_offsets[32];

    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;  // one warp per 1024-element chunk

    const float* xr = x + ((size_t)row << 15);  // row * 32768
    float* yr = y + ((size_t)row << 15);

    int chunk_end = (warp << 10) + 1023;
    float carry = 0.0f;
    float outv[32];

    // Each warp scans one 1024-element chunk from right to left.  For each
    // 32-element group, lanes access consecutive addresses (in reverse order),
    // then a warp inclusive scan gives the partial reverse cumulative sums.
    #pragma unroll
    for (int g = 0; g < 32; ++g) {
        int idx = chunk_end - ((g << 5) + lane);
        float v = xr[idx];
        float scan = warp_inclusive_scan(v);
        outv[g] = scan + carry;
        float group_sum = __shfl_sync(0xffffffffu, scan, 31);
        carry += group_sum;
    }

    if (lane == 0) chunk_sums[warp] = carry;
    __syncthreads();

    // Scan the 32 chunk sums from right to left to get the amount that must be
    // added to every element in each chunk.
    if (warp == 0) {
        float v = chunk_sums[31 - lane];
        float scan = warp_inclusive_scan(v);
        chunk_offsets[31 - lane] = scan - v;  // exclusive suffix over chunks
    }
    __syncthreads();

    float off = chunk_offsets[warp];
    #pragma unroll
    for (int g = 0; g < 32; ++g) {
        int idx = chunk_end - ((g << 5) + lane);
        yr[idx] = outv[g] + off;
    }
}

torch::Tensor reverse_cumsum_cuda(torch::Tensor x, int64_t dim) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.dim() == 2, "only 2D supported");
    TORCH_CHECK(dim == 1 || dim == -1, "only dim=1 supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(x.size(1) == 32768, "only N=32768 supported");

    auto y = torch::empty_like(x);
    int rows = (int)x.size(0);
    reverse_cumsum_warpchunks_kernel<<<rows, 1024, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), y.data_ptr<float>());
    return y;
}
'''

_ext = load_inline(
    name='reverse_cumsum_ext_final',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['reverse_cumsum_cuda'],
    extra_cuda_cflags=['-O3'],
    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_cuda(x, self.dim)