```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

# Warp-per-row exclusive cumsum.
#
# Each 32-lane warp independently scans one input row of length N and writes
# a row of length N+1: out[0] = 0, out[k] = sum(x[0..k-1]), out[N] = row
# total. This matches the reference (cat zero-slab + slice + cumsum) when
# dim is the last dimension, while:
#   * using no block-level synchronization at all (warp scans only,
#     carry kept in registers),
#   * keeping a depth-2 software prefetch pipeline so every warp always has
#     global loads in flight,
#   * staging the four per-lane results in shared memory with a conflict-free
#     padded layout (slot = e + e/32) so global stores stay fully coalesced
#     despite the odd (N+1) output row pitch.
exclusive_cumsum_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define NT 256
#define NWARP (NT / 32)
#define STAGE 132  // 128 floats + 1 float of padding per 32 elements

__device__ __forceinline__ void warp_scan_excl(double ts, int lane,
                                               double& excl, double& total) {
    double incl = ts;
    #pragma unroll
    for (int o = 1; o < 32; o <<= 1) {
        double y = __shfl_up_sync(0xffffffffu, incl, o);
        if (lane >= o) incl += y;
    }
    excl = incl - ts;
    total = __shfl_sync(0xffffffffu, incl, 31);
}

__global__ void __launch_bounds__(NT)
exclusive_cumsum_warp_rows_kernel(const float* __restrict__ x,
                                  float* __restrict__ out,
                                  const long n4,   // N / 4
                                  const long N,
                                  const long rows) {
    __shared__ float stage[NWARP][STAGE];

    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const long row = (long)blockIdx.x * NWARP + warp;
    if (row >= rows) return;

    const float4* __restrict__ in_row = reinterpret_cast<const float4*>(x + row * N);
    float* __restrict__ out_row = out + row * (N + 1);
    float* __restrict__ st = stage[warp];

    double carry = 0.0;

    const long P = n4 >> 6;              // full pair iterations (64 float4 each)
    const long rem = n4 - (P << 6);      // leftover float4 count (< 64)

    float4 pa0 = make_float4(0.f, 0.f, 0.f, 0.f);
    float4 pb0 = make_float4(0.f, 0.f, 0.f, 0.f);
    float4 pa1 = make_float4(0.f, 0.f, 0.f, 0.f);
    float4 pb1 = make_float4(0.f, 0.f, 0.f, 0.f);
    if (P > 0) {
        pa0 = __ldg(in_row + lane);
        pb0 = __ldg(in_row + 32 + lane);
        if (P > 1) {
            pa1 = __ldg(in_row + 64 + lane);
            pb1 = __ldg(in_row + 96 + lane);
        }
    }

    for (long t = 0; t < P; ++t) {
        const float4 a = pa0, b = pb0;
        pa0 = pa1; pb0 = pb1;
        const long nt = t + 2;
        if (nt < P) {
            pa1 = __ldg(in_row + (nt << 6) + lane);
            pb1 = __ldg(in_row + (nt << 6) + 32 + lane);
        }

        const double ts = ((double)a.x + (double)a.y + (double)a.z + (double)a.w)
                        + ((double)b.x + (double)b.y + (double)b.z + (double)b.w);
        double excl, total;
        warp_scan_excl(ts, lane, excl, total);

        const double base = carry + excl;
        double run = base;
        st[lane] = (float)run;            run += (double)a.x;
        st[32 + lane + 1] = (float)run;   run += (double)a.y;
        st[64 + lane + 2] = (float)run;   run += (double)a.z;
        st[96 + lane + 3] = (float)run;   run += (double)a.w;
        st[128 + lane + 4] = (float)run;  run += (double)b.x;
        st[160 + lane + 5] = (float)run;  run += (double)b.y;
        st[192 + lane + 6] = (float)run;  run += (double)b.z;
        st[224 + lane + 7] = (float)run;

        carry += total;
        __syncwarp();

        const long ob = t << 8;  // t * 256 output elements
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            out_row[ob + i * 32 + lane] = st[i * 33 + lane];
        }
        __syncwarp();
    }

    // Remainder: up to 63 float4s, processed in chunks of 32.
    for (int c = 0; c < 2; ++c) {
        const long chunk_start = (P << 6) + c * 32;
        const long left = n4 - chunk_start;
        if (left <= 0) break;
        const int cnt = left < 32 ? (int)left : 32;

        float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
        if (lane < cnt) v = __ldg(in_row + chunk_start + lane);

        const double ts = (double)v.x + (double)v.y + (double)v.z + (double)v.w;
        double excl, total;
        warp_scan_excl(ts, lane, excl, total);

        const double base = carry + excl;
        double run = base;
        st[lane] = (float)run;            run += (double)v.x;
        st[32 + lane + 1] = (float)run;   run += (double)v.y;
        st[64 + lane + 2] = (float)run;   run += (double)v.z;
        st[96 + lane + 3] = (float)run;

        carry += total;
        __syncwarp();

        const long ob = chunk_start * 4;
        const int lim = cnt * 4;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int e = i * 32 + lane;
            if (e < lim) out_row[ob + e] = st[i * 33 + lane];
        }
        __syncwarp();
    }

    // Append the row total at position N.
    if (lane == 0) out_row[N] = (float)carry;
}

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);
    const int blocks = (int)((out_rows + NWARP - 1) / NWARP);
    auto stream = at::cuda::getCurrentCUDAStream();
    exclusive_cumsum_warp_rows_kernel<<<blocks, NT, 0, stream>>>(
        x.data_ptr<float>(), out.data_ptr<float>(), n4, (long)N, (long)out_rows);
    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_warp_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 warp-per-row 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)
```