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

# Fused exclusive-cumsum kernel.
#
# Fixes the illegal-memory-access bug in the previous attempt: the number of
# full 1024-float4 (4096-element) tiles per row is n4 >> 10 (each tile holds
# 2 * NT = 1024 float4 values), not n4 >> 9. The incorrect shift doubled the
# tile count and made the kernel read/write past the end of every row.
#
# Design (same as the intended v2, now index-correct):
#   * One thread block per row; each thread owns TWO float4 groups per tile
#     (tile = 4096 elements) -> 8 tiles per 32768-element row, so few
#     barriers per row.
#   * Register double-buffered prefetch of the next tile with streaming
#     cache hints (__ldcs/__stcs): the data is touched exactly once.
#   * Block-wide scan in double precision via warp shuffles + one shared
#     array of per-warp totals (both groups' warp scans are issued before a
#     single barrier to minimize sync count: 3 barriers per tile).
#   * Results are staged in shared memory so the global stores are fully
#     coalesced even though output rows are only 4-byte aligned
#     (row length N+1 is odd).
#   * __launch_bounds__(512, 2) gives two resident blocks per SM so barriers
#     and memory latency overlap across blocks without register spills.
exclusive_cumsum_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

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

__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;
}

// One thread block per row. Each thread handles two float4 groups per tile.
__global__ void __launch_bounds__(NT, 2)
exclusive_cumsum_rows_kernel(const float* __restrict__ x,
                             float* __restrict__ out,
                             const int n4,   // N / 4
                             const int N) {  // input row length
    __shared__ double wsum[2][NWARP];
    __shared__ float stage[NT * 8];

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

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

    const int g4 = n4 >> 10;       // full 1024-float4 (4096-element) tiles
    const int r4 = n4 & 1023;      // leftover float4 count (< 1024)

    double carry = 0.0;

    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 (g4 > 0) {
        a = __ldcs(in_row + tid);
        b = __ldcs(in_row + NT + tid);
    }

    for (int t = 0; t < g4; ++t) {
        const float4 va = a;
        const float4 vb = b;
        const int nt = t + 1;
        if (nt < g4) {
            a = __ldcs(in_row + (nt << 10) + tid);
            b = __ldcs(in_row + (nt << 10) + NT + tid);
        }

        const double sa = (double)va.x + (double)va.y + (double)va.z + (double)va.w;
        const double sb = (double)vb.x + (double)vb.y + (double)vb.z + (double)vb.w;
        const double inclA = warp_incl_scan_d(sa, lane);
        const double inclB = warp_incl_scan_d(sb, lane);
        if (lane == 31) {
            wsum[0][warp] = inclA;
            wsum[1][warp] = inclB;
        }
        __syncthreads();
        if (warp == 0) {
            double wa = (lane < NWARP) ? wsum[0][lane] : 0.0;
            double wb = (lane < NWARP) ? wsum[1][lane] : 0.0;
            wa = warp_incl_scan_d(wa, lane);
            wb = warp_incl_scan_d(wb, lane);
            if (lane < NWARP) {
                wsum[0][lane] = wa;
                wsum[1][lane] = wb;
            }
        }
        __syncthreads();

        const double preA   = (warp == 0) ? 0.0 : wsum[0][warp - 1];
        const double totalA = wsum[0][NWARP - 1];
        const double preB   = (warp == 0) ? 0.0 : wsum[1][warp - 1];
        const double totalB = wsum[1][NWARP - 1];

        double run = carry + preA + (inclA - sa);
        stage[tid * 4 + 0] = (float)run;  run += (double)va.x;
        stage[tid * 4 + 1] = (float)run;  run += (double)va.y;
        stage[tid * 4 + 2] = (float)run;  run += (double)va.z;
        stage[tid * 4 + 3] = (float)run;

        run = carry + totalA + preB + (inclB - sb);
        stage[NT * 4 + tid * 4 + 0] = (float)run;  run += (double)vb.x;
        stage[NT * 4 + tid * 4 + 1] = (float)run;  run += (double)vb.y;
        stage[NT * 4 + tid * 4 + 2] = (float)run;  run += (double)vb.z;
        stage[NT * 4 + tid * 4 + 3] = (float)run;

        carry += totalA + totalB;
        __syncthreads();  // staging complete

        const int ob = t << 12;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            __stcs(out_row + ob + j * NT + tid, stage[j * NT + tid]);
        }
        // No trailing barrier: the next tile's staging writes happen only
        // after its two scan barriers, by which time every thread has
        // finished the writeback above (program order).
    }

    if (r4 > 0) {
        const int base4 = g4 << 10;
        float4 va = make_float4(0.f, 0.f, 0.f, 0.f);
        float4 vb = make_float4(0.f, 0.f, 0.f, 0.f);
        if (tid < r4)          va = __ldcs(in_row + base4 + tid);
        if (NT + tid < r4)     vb = __ldcs(in_row + base4 + NT + tid);

        const double sa = (double)va.x + (double)va.y + (double)va.z + (double)va.w;
        const double sb = (double)vb.x + (double)vb.y + (double)vb.z + (double)vb.w;
        const double inclA = warp_incl_scan_d(sa, lane);
        const double inclB = warp_incl_scan_d(sb, lane);
        if (lane == 31) {
            wsum[0][warp] = inclA;
            wsum[1][warp] = inclB;
        }
        __syncthreads();
        if (warp == 0) {
            double wa = (lane < NWARP) ? wsum[0][lane] : 0.0;
            double wb = (lane < NWARP) ? wsum[1][lane] : 0.0;
            wa = warp_incl_scan_d(wa, lane);
            wb = warp_incl_scan_d(wb, lane);
            if (lane < NWARP) {
                wsum[0][lane] = wa;
                wsum[1][lane] = wb;
            }
        }
        __syncthreads();

        const double preA   = (warp == 0) ? 0.0 : wsum[0][warp - 1];
        const double totalA = wsum[0][NWARP - 1];
        const double preB   = (warp == 0) ? 0.0 : wsum[1][warp - 1];
        const double totalB = wsum[1][NWARP - 1];

        double run = carry + preA + (inclA - sa);
        stage[tid * 4 + 0] = (float)run;  run += (double)va.x;
        stage[tid * 4 + 1] = (float)run;  run += (double)va.y;
        stage[tid * 4 + 2] = (float)run;  run += (double)va.z;
        stage[tid * 4 + 3] = (float)run;

        run = carry + totalA + preB + (inclB - sb);
        stage[NT * 4 + tid * 4 + 0] = (float)run;  run += (double)vb.x;
        stage[NT * 4 + tid * 4 + 1] = (float)run;  run += (double)vb.y;
        stage[NT * 4 + tid * 4 + 2] = (float)run;  run += (double)vb.z;
        stage[NT * 4 + tid * 4 + 3] = (float)run;

        carry += totalA + totalB;
        __syncthreads();  // staging complete

        const int ob = g4 << 12;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            const int k = ob + j * NT + tid;
            if (k < N) __stcs(out_row + k, stage[j * NT + tid]);
        }
    }

    // Append the row total at position N.
    if (tid == 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");
    TORCH_CHECK(N < (1LL << 30), "last dim too large for this kernel");
    TORCH_CHECK(out_rows <= (int64_t)2147483647, "too many rows");

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

    auto stream = at::cuda::getCurrentCUDAStream();
    exclusive_cumsum_rows_kernel<<<(int)out_rows, NT, 0, stream>>>(
        x.data_ptr<float>(), out.data_ptr<float>(), (int)(N >> 2), (int)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_v3",
    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)
```