```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 (corrected and improved).
#
# For each input row of length N the kernel 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 a zero slab along dim, slice [:-1] on dim 0,
# cumsum along dim) when dim is the last dimension, while fusing everything
# into a single kernel with no concatenated intermediate tensor.
#
# Design notes:
#   * Each warp owns one row -> no __syncthreads anywhere; the cross-tile
#     carry stays in a register and all scans are warp shuffles.
#   * Per 256-element tile each lane loads two float4: "a" (elements
#     4*lane..4*lane+3) and "b" (elements 128+4*lane..131+4*lane). Two
#     SEPARATE warp scans (one over the a-sums, one over the b-sums) keep the
#     element ordering correct: the b-part base is carry + totalA + exclB.
#     (Scanning the combined a+b sums is wrong because the b elements lie
#     after every a element in row order.)
#   * Per-lane results are staged in shared memory with a conflict-free
#     padded layout (slot(e) = e + e/32, STAGE = 264) so the global stores
#     are fully coalesced despite each lane owning strided elements.
#   * One-tile software prefetch keeps global loads in flight, and streaming
#     cache hints (__ldcs/__stcs) are used since the data is touched once.
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 264  // 256 floats + 8 floats of padding (1 per 32 elements)

__device__ __forceinline__ double warp_incl_scan_d(double v, const int lane) {
    #pragma unroll
    for (int o = 1; o < 32; o <<= 1) {
        double y = __shfl_up_sync(0xffffffffu, v, o);
        if (lane >= o) v += y;
    }
    return v;
}

__global__ void __launch_bounds__(NT, 6)
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 long P  = n4 >> 6;         // full 64-float4 (256-element) tiles
    const int  r4 = (int)(n4 & 63);  // leftover float4 count (< 64)

    // Padded staging slots for this lane's two 4-element groups.
    const int s0 = 4 * lane + (lane >> 3);
    const int s1 = 132 + s0;

    double carry = 0.0;

    float4 na = make_float4(0.f, 0.f, 0.f, 0.f);
    float4 nb = make_float4(0.f, 0.f, 0.f, 0.f);
    if (P > 0) {
        na = __ldcs(in_row + lane);
        nb = __ldcs(in_row + 32 + lane);
    }

    for (long t = 0; t < P; ++t) {
        const float4 a = na;
        const float4 b = nb;
        const long nt = t + 1;
        if (nt < P) {
            na = __ldcs(in_row + (nt << 6) + lane);
            nb = __ldcs(in_row + (nt << 6) + 32 + lane);
        }

        // Scan the a-group (elements 4*lane..4*lane+3 of this tile).
        const double sa = (double)a.x + (double)a.y + (double)a.z + (double)a.w;
        const double inclA = warp_incl_scan_d(sa, lane);
        const double totalA = __shfl_sync(0xffffffffu, inclA, 31);

        double run = carry + (inclA - sa);
        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;

        // Scan the b-group (elements 128+4*lane..128+4*lane+3 of this tile).
        // All a elements precede all b elements, hence the totalA offset.
        const double sb = (double)b.x + (double)b.y + (double)b.z + (double)b.w;
        const double inclB = warp_incl_scan_d(sb, lane);
        const double totalB = __shfl_sync(0xffffffffu, inclB, 31);

        run = carry + totalA + (inclB - sb);
        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 += totalA + totalB;
        __syncwarp(0xffffffffu);

        // Coalesced writeback of the whole 256-element tile from staging.
        const long ob = t << 8;
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            __stcs(out_row + ob + i * 32 + lane, st[i * 33 + lane]);
        }
        __syncwarp(0xffffffffu);
    }

    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 sa = (double)a.x + (double)a.y + (double)a.z + (double)a.w;
        const double inclA = warp_incl_scan_d(sa, lane);
        const double totalA = __shfl_sync(0xffffffffu, inclA, 31);

        double run = carry + (inclA - sa);
        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 double sb = (double)b.x + (double)b.y + (double)b.z + (double)b.w;
        const double inclB = warp_incl_scan_d(sb, lane);
        const double totalB = __shfl_sync(0xffffffffu, inclB, 31);

        run = carry + totalA + (inclB - sb);
        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 += totalA + totalB;
        __syncwarp(0xffffffffu);

        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(0xffffffffu);
    }

    // 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_row_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)
```