```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Custom CUDA kernel: for each input row of length N it produces a row of
# length N+1 containing the exclusive cumulative sum with the row total
# appended (i.e. out[0]=0, out[k]=sum(x[0..k-1]), out[N]=sum(x[0..N-1])),
# processing only the first `out_rows` rows. This reproduces exactly the
# semantics of the reference (cat-with-zero + slice + cumsum) for the case
# dim == last dimension, while fusing everything into a single kernel and
# avoiding the huge intermediate concatenated tensor.
exclusive_cumsum_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define NT 1024
#define NWARP (NT / 32)

__device__ __forceinline__ double warp_inclusive_scan_d(double v) {
    #pragma unroll
    for (int offset = 1; offset < 32; offset <<= 1) {
        double y = __shfl_up_sync(0xffffffffu, v, offset);
        if ((threadIdx.x & 31) >= offset) v += y;
    }
    return v;
}

// One thread block per row. Each thread handles one float4 per tile.
__global__ void __launch_bounds__(NT)
exclusive_cumsum_rows_kernel(const float* __restrict__ x,
                             float* __restrict__ out,
                             const long n4,   // N / 4
                             const long N) {  // row length of input
    __shared__ double warp_sums[NWARP];
    __shared__ float stage[NT * 4];

    const long row = blockIdx.x;
    const float4* __restrict__ in_row = reinterpret_cast<const float4*>(x + row * N);
    float* __restrict__ out_row = out + row * (N + 1);

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    const long num_tiles = (n4 + NT - 1) / NT;
    double carry = 0.0;

    for (long t = 0; t < num_tiles; ++t) {
        const long idx = t * NT + tid;
        const bool valid = idx < n4;

        float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
        if (valid) v = __ldg(in_row + idx);

        const double sum = (double)v.x + (double)v.y + (double)v.z + (double)v.w;

        // Block-wide exclusive scan of the per-thread sums (in double precision).
        double incl = warp_inclusive_scan_d(sum);
        if (lane == 31) warp_sums[warp] = incl;
        __syncthreads();
        if (warp == 0) {
            double w = (lane < NWARP) ? warp_sums[lane] : 0.0;
            w = warp_inclusive_scan_d(w);
            if (lane < NWARP) warp_sums[lane] = w;
        }
        __syncthreads();

        const double warp_prefix = (warp == 0) ? 0.0 : warp_sums[warp - 1];
        const double tile_total  = warp_sums[NWARP - 1];
        const double base        = carry + warp_prefix + (incl - sum);

        // Stage the four results so the global stores below are fully coalesced.
        stage[tid * 4 + 0] = (float)(base);
        stage[tid * 4 + 1] = (float)(base + (double)v.x);
        stage[tid * 4 + 2] = (float)(base + (double)v.x + (double)v.y);
        stage[tid * 4 + 3] = (float)(base + (double)v.x + (double)v.y + (double)v.z);

        // The very last thread appends the row total at position N.
        if (idx == n4 - 1) {
            out_row[N] = (float)(base + sum);
        }

        carry += tile_total;
        __syncthreads();

        const long tile_base = t * (long)(NT * 4);
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const long k = tile_base + j * NT + tid;
            if (k < N) out_row[k] = stage[j * NT + tid];
        }
        __syncthreads();
    }
}

torch::Tensor exclusive_cumsum_rows(torch::Tensor x, int64_t out_rows) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    const int64_t N = x.size(-1);
    TORCH_CHECK(N > 0 && (N % 4) == 0, "last dim must be a positive multiple of 4");

    auto out = at::empty({out_rows, N + 1}, x.options());
    if (out_rows <= 0) return out;

    const long n4 = (long)(N / 4);
    auto stream = at::cuda::getCurrentCUDAStream();
    exclusive_cumsum_rows_kernel<<<(int)out_rows, NT, 0, stream>>>(
        x.data_ptr<float>(), out.data_ptr<float>(), n4, (long)N);
    return out;
}
"""

exclusive_cumsum_cpp_source = (
    "torch::Tensor exclusive_cumsum_rows(torch::Tensor x, int64_t out_rows);"
)

exclusive_cumsum_ext = load_inline(
    name="exclusive_cumsum_rows_ext",
    cpp_sources=exclusive_cumsum_cpp_source,
    cuda_sources=exclusive_cumsum_cuda_source,
    functions=["exclusive_cumsum_rows"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized exclusive cumulative sum.

    For a CUDA float32 contiguous tensor with `dim` equal to the last
    dimension (the benchmark case), a single fused CUDA kernel computes the
    per-row exclusive cumsum (with the row total appended, last slice of
    dim-0 dropped) directly, reproducing the reference semantics without
    materializing the concatenated intermediate tensor. All other cases fall
    back to the exact reference implementation.
    """

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

    def forward(self, x):
        d = self.dim if self.dim >= 0 else self.dim + x.dim()
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.is_contiguous()
            and x.dim() >= 2
            and d == x.dim() - 1
            and x.size(0) >= 1
            and x.size(-1) > 0
            and (x.size(-1) % 4) == 0
        ):
            n = x.size(-1)
            total_rows = x.numel() // n
            # Number of last-dim rows removed by the reference's `[:-1]`
            # slice on dimension 0.
            tail_rows = x.numel() // (x.size(0) * n)
            out_rows = total_rows - tail_rows
            out = exclusive_cumsum_ext.exclusive_cumsum_rows(
                x.view(total_rows, n), out_rows
            )
            out_shape = list(x.shape)
            out_shape[0] -= 1
            out_shape[-1] = n + 1
            return out.view(out_shape)

        # Exact reference fallback for all other configurations.
        exclusive_cumsum = torch.cat(
            (torch.zeros_like(x.select(self.dim, 0).unsqueeze(self.dim)), x),
            dim=self.dim,
        )[:-1]
        return torch.cumsum(exclusive_cumsum, dim=self.dim)
```