```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, the
#     cross-tile carry is kept in registers),
#   * keeping a depth-2 software prefetch pipeline so every warp always has
#     global loads in flight,
#   * staging the 256 per-tile results in shared memory with a conflict-free
#     padded layout (slot(e) = e + e/32) so global stores stay fully coalesced
#     despite each lane owning 8 strided elements,
#   * using streaming cache hints (__ldcs/__stcs) since data is touched once.
exclusive_cumsum_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

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

__device__ __forceinline__ void warp_scan_excl_d(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, 8)
exclusive_cumsum_warp_row_kernel(const float* __restrict__ x,
                                 float* __restrict__ out,
                                 const long n4,   // N / 4
                                 const long N,    // input row length
                                 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];
    const int pad = lane >> 3;

    double carry = 0.0;
    const long P  = n4 >> 6;         // full 64-float4 (256-element) tiles
    const int  r4 = (int)(n4 & 63);  // 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 = __ldcs(in_row + lane);
        pb0 = __ldcs(in_row + 32 + lane);
        if (P > 1) {
            pa1 = __ldcs(in_row + 64 + lane);
            pb1 = __ldcs(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 = __ldcs(in_row + (nt << 6) + lane);
            pb1 = __ldcs(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_d(ts, lane, excl, total);

        // Lane l owns elements 4l..4l+3 (from a) and 128+4l..128+4l+3 (from b)
        // within this 256-element tile. Stage them at slot(e) = e + e/32 so
        // the readback below is fully coalesced and bank-conflict free.
        double run = carry + excl;
        const int s0 = 4 * lane + pad;
        st[s0 + 0] = (float)run;  run += (double)a.x;
        st[s0 + 1] = (float)run;  run += (double)a.y;
        st[s0 + 2] = (float)run;  run += (double)a.z;
        st[s0 + 3] = (float)run;  run += (double)a.w;
        const int s1 = 132 + s0;
        st[s1 + 0] = (float)run;  run += (double)b.x;
        st[s1 + 1] = (float)run;  run += (double)b.y;
        st[s1 + 2] = (float)run;  run += (double)b.z;
        st[s1 + 3] = (float)run;

        carry += total;
        __syncwarp();

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

    if (r4 > 0) {
        const long base4 = P << 6;
        float4 a = make_float4(0.f, 0.f, 0.f, 0.f);
        float4 b = make_float4(0.f, 0.f, 0.f, 0.f);
        if (lane < r4)      a = __ldcs(in_row + base4 + lane);
        if (32 + lane < r4) b = __ldcs(in_row + base4 + 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_d(ts, lane, excl, total);

        double run = carry + excl;
        const int s0 = 4 * lane + pad;
        st[s0 + 0] = (float)run;  run += (double)a.x;
        st[s0 + 1] = (float)run;  run += (double)a.y;
        st[s0 + 2] = (float)run;  run += (double)a.z;
        st[s0 + 3] = (float)run;  run += (double)a.w;
        const int s1 = 132 + s0;
        st[s1 + 0] = (float)run;  run += (double)b.x;
        st[s1 + 1] = (float)run;  run += (double)b.y;
        st[s1 + 2] = (float)run;  run += (double)b.z;
        st[s1 + 3] = (float)run;

        carry += total;
        __syncwarp();

        const long ob = P << 8;
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            const long e = ob + i * 32 + lane;
            if (e < N) __stcs(out_row + e, st[i * 33 + lane]);
        }
        __syncwarp();
    }

    // Append the row total at position N.
    if (lane == 0) __stcs(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 >> 2);
    const int blocks = (int)((out_rows + NWARP - 1) / NWARP);
    auto stream = at::cuda::getCurrentCUDAStream();
    exclusive_cumsum_warp_row_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_stage_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)
```