```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Custom fused LayerNorm CUDA implementation.
#
# Strategy (tuned for the target case: few very long rows):
#   1) partial-sums kernel: each block reduces a chunk of a row -> (sum, sumsq)
#      using vectorized float4 loads and warp/block reductions.
#   2) tiny kernel reduces the partials -> per-row mean and rstd.
#   3) normalize kernel: out = (x - mean) * rstd * weight + bias.
#
# Key optimization vs. a naive per-(row, chunk) grid: weight/bias are shared by
# every row, so re-reading them once per row can add hundreds of MB of DRAM
# traffic (L2 gets thrashed by the streaming x/out data).  The normalize kernel
# therefore assigns each block a fixed column chunk and loops over ALL rows
# inside the kernel: weight/bias are loaded exactly once into registers, and
# x/out are accessed as pure streaming loads/stores (__ldcs/__stcs).  The row
# loop runs in reverse so the rows most recently read by the reduction pass
# (and thus possibly still resident in L2) are consumed first.

layer_norm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cstdint>

#define LN_THREADS 256
#define LN_WARPS (LN_THREADS / 32)

__device__ __forceinline__ void ln_warp_reduce(float &sum, float &sumsq) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum   += __shfl_down_sync(0xffffffffu, sum, off);
        sumsq += __shfl_down_sync(0xffffffffu, sumsq, off);
    }
}

__device__ __forceinline__ void ln_block_reduce(float &sum, float &sumsq,
                                                float *sh_sum, float *sh_sumsq) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    ln_warp_reduce(sum, sumsq);
    if (lane == 0) {
        sh_sum[warp]   = sum;
        sh_sumsq[warp] = sumsq;
    }
    __syncthreads();
    if (warp == 0) {
        const int nwarp = (blockDim.x + 31) >> 5;
        sum   = (lane < nwarp) ? sh_sum[lane]   : 0.0f;
        sumsq = (lane < nwarp) ? sh_sumsq[lane] : 0.0f;
        ln_warp_reduce(sum, sumsq);
    }
}

// ---------------- Kernel 1: partial sums (vectorized) ----------------
__global__ void ln_partial_sums_vec4(const float* __restrict__ x,
                                     float* __restrict__ partials,
                                     int N, int chunk, int splits) {
    const int row   = blockIdx.y;
    const int split = blockIdx.x;
    const float* row_ptr = x + (long long)row * N;

    const int start = split * chunk;
    int end = start + chunk;
    if (end > N) end = N;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float q0 = 0.f, q1 = 0.f, q2 = 0.f, q3 = 0.f;
    const int stride = blockDim.x * 4;
    #pragma unroll 4
    for (int i = start + threadIdx.x * 4; i < end; i += stride) {
        const float4 v = reinterpret_cast<const float4*>(row_ptr)[i >> 2];
        s0 += v.x; s1 += v.y; s2 += v.z; s3 += v.w;
        q0 += v.x * v.x; q1 += v.y * v.y; q2 += v.z * v.z; q3 += v.w * v.w;
    }
    float sum   = (s0 + s1) + (s2 + s3);
    float sumsq = (q0 + q1) + (q2 + q3);

    __shared__ float sh_sum[LN_WARPS];
    __shared__ float sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, sumsq, sh_sum, sh_sumsq);

    if (threadIdx.x == 0) {
        float* p = partials + ((long long)row * splits + split) * 2;
        p[0] = sum;
        p[1] = sumsq;
    }
}

// ---------------- Kernel 1: partial sums (scalar fallback) ----------------
__global__ void ln_partial_sums_scalar(const float* __restrict__ x,
                                       float* __restrict__ partials,
                                       int N, int chunk, int splits) {
    const int row   = blockIdx.y;
    const int split = blockIdx.x;
    const float* row_ptr = x + (long long)row * N;

    const int start = split * chunk;
    int end = start + chunk;
    if (end > N) end = N;

    float sum = 0.f, sumsq = 0.f;
    #pragma unroll 4
    for (int i = start + threadIdx.x; i < end; i += blockDim.x) {
        const float v = row_ptr[i];
        sum   += v;
        sumsq += v * v;
    }

    __shared__ float sh_sum[LN_WARPS];
    __shared__ float sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, sumsq, sh_sum, sh_sumsq);

    if (threadIdx.x == 0) {
        float* p = partials + ((long long)row * splits + split) * 2;
        p[0] = sum;
        p[1] = sumsq;
    }
}

// ---------------- Kernel 2: final per-row statistics ----------------
__global__ void ln_final_stats(const float* __restrict__ partials,
                               float* __restrict__ mean,
                               float* __restrict__ rstd,
                               int splits, int N, float eps) {
    const int row = blockIdx.x;
    float sum = 0.f, sumsq = 0.f;
    for (int i = threadIdx.x; i < splits; i += blockDim.x) {
        const float* p = partials + ((long long)row * splits + i) * 2;
        sum   += p[0];
        sumsq += p[1];
    }

    __shared__ float sh_sum[LN_WARPS];
    __shared__ float sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, sumsq, sh_sum, sh_sumsq);

    if (threadIdx.x == 0) {
        const float m = sum / (float)N;
        float var = sumsq / (float)N - m * m;
        var = var < 0.f ? 0.f : var;  // guard against tiny negative roundoff
        mean[row] = m;
        rstd[row] = rsqrtf(var + eps);
    }
}

// ------------- Kernel 3a: normalize, block owns a column chunk and
// loops over all rows (weight/bias loaded once into registers). ---------
__global__ void ln_apply_rows_vec4(const float* __restrict__ x,
                                   float* __restrict__ out,
                                   const float* __restrict__ mean,
                                   const float* __restrict__ rstd,
                                   const float* __restrict__ weight,
                                   const float* __restrict__ bias,
                                   int N, int rows) {
    const int n4 = N >> 2;
    const int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= n4) return;

    const float4 w = reinterpret_cast<const float4*>(weight)[c];
    const float4 b = reinterpret_cast<const float4*>(bias)[c];
    const float4* xv = reinterpret_cast<const float4*>(x) + c;
    float4* ov = reinterpret_cast<float4*>(out) + c;

    int r = rows;
    for (; r >= 4; r -= 4) {
        const long long o0 = (long long)(r - 1) * n4;
        const long long o1 = (long long)(r - 2) * n4;
        const long long o2 = (long long)(r - 3) * n4;
        const long long o3 = (long long)(r - 4) * n4;
        const float4 v0 = __ldcs(xv + o0);
        const float4 v1 = __ldcs(xv + o1);
        const float4 v2 = __ldcs(xv + o2);
        const float4 v3 = __ldcs(xv + o3);
        const float m0 = mean[r - 1], t0 = rstd[r - 1];
        const float m1 = mean[r - 2], t1 = rstd[r - 2];
        const float m2 = mean[r - 3], t2 = rstd[r - 3];
        const float m3 = mean[r - 4], t3 = rstd[r - 4];
        float4 u0, u1, u2, u3;
        u0.x = (v0.x - m0) * t0 * w.x + b.x;
        u0.y = (v0.y - m0) * t0 * w.y + b.y;
        u0.z = (v0.z - m0) * t0 * w.z + b.z;
        u0.w = (v0.w - m0) * t0 * w.w + b.w;
        u1.x = (v1.x - m1) * t1 * w.x + b.x;
        u1.y = (v1.y - m1) * t1 * w.y + b.y;
        u1.z = (v1.z - m1) * t1 * w.z + b.z;
        u1.w = (v1.w - m1) * t1 * w.w + b.w;
        u2.x = (v2.x - m2) * t2 * w.x + b.x;
        u2.y = (v2.y - m2) * t2 * w.y + b.y;
        u2.z = (v2.z - m2) * t2 * w.z + b.z;
        u2.w = (v2.w - m2) * t2 * w.w + b.w;
        u3.x = (v3.x - m3) * t3 * w.x + b.x;
        u3.y = (v3.y - m3) * t3 * w.y + b.y;
        u3.z = (v3.z - m3) * t3 * w.z + b.z;
        u3.w = (v3.w - m3) * t3 * w.w + b.w;
        __stcs(ov + o0, u0);
        __stcs(ov + o1, u1);
        __stcs(ov + o2, u2);
        __stcs(ov + o3, u3);
    }
    for (; r > 0; --r) {
        const long long o = (long long)(r - 1) * n4;
        const float4 v = __ldcs(xv + o);
        const float m = mean[r - 1], t = rstd[r - 1];
        float4 u;
        u.x = (v.x - m) * t * w.x + b.x;
        u.y = (v.y - m) * t * w.y + b.y;
        u.z = (v.z - m) * t * w.z + b.z;
        u.w = (v.w - m) * t * w.w + b.w;
        __stcs(ov + o, u);
    }
}

__global__ void ln_apply_rows_scalar(const float* __restrict__ x,
                                     float* __restrict__ out,
                                     const float* __restrict__ mean,
                                     const float* __restrict__ rstd,
                                     const float* __restrict__ weight,
                                     const float* __restrict__ bias,
                                     int N, int rows) {
    const int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= N) return;

    const float w = weight[c];
    const float b = bias[c];
    const float* xp = x + c;
    float* op = out + c;

    int r = rows;
    for (; r >= 4; r -= 4) {
        const long long o0 = (long long)(r - 1) * N;
        const long long o1 = (long long)(r - 2) * N;
        const long long o2 = (long long)(r - 3) * N;
        const long long o3 = (long long)(r - 4) * N;
        const float v0 = __ldcs(xp + o0);
        const float v1 = __ldcs(xp + o1);
        const float v2 = __ldcs(xp + o2);
        const float v3 = __ldcs(xp + o3);
        __stcs(op + o0, (v0 - mean[r - 1]) * rstd[r - 1] * w + b);
        __stcs(op + o1, (v1 - mean[r - 2]) * rstd[r - 2] * w + b);
        __stcs(op + o2, (v2 - mean[r - 3]) * rstd[r - 3] * w + b);
        __stcs(op + o3, (v3 - mean[r - 4]) * rstd[r - 4] * w + b);
    }
    for (; r > 0; --r) {
        const long long o = (long long)(r - 1) * N;
        __stcs(op + o, (__ldcs(xp + o) - mean[r - 1]) * rstd[r - 1] * w + b);
    }
}

// ------------- Kernel 3b: fallback 2D normalize (vectorized) -----------
__global__ void ln_normalize_vec4(const float* __restrict__ x,
                                  float* __restrict__ out,
                                  const float* __restrict__ mean,
                                  const float* __restrict__ rstd,
                                  const float* __restrict__ weight,
                                  const float* __restrict__ bias,
                                  int N) {
    const int row = blockIdx.y;
    const float m = mean[row];
    const float r = rstd[row];

    const float4* xv = reinterpret_cast<const float4*>(x + (long long)row * N);
    const float4* wv = reinterpret_cast<const float4*>(weight);
    const float4* bv = reinterpret_cast<const float4*>(bias);
    float4* ov = reinterpret_cast<float4*>(out + (long long)row * N);

    const int n4 = N >> 2;
    const int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n4; i += stride) {
        const float4 v = xv[i];
        const float4 w = wv[i];
        const float4 b = bv[i];
        float4 o;
        o.x = (v.x - m) * r * w.x + b.x;
        o.y = (v.y - m) * r * w.y + b.y;
        o.z = (v.z - m) * r * w.z + b.z;
        o.w = (v.w - m) * r * w.w + b.w;
        ov[i] = o;
    }
}

// ------------- Kernel 3b: fallback 2D normalize (scalar) ---------------
__global__ void ln_normalize_scalar(const float* __restrict__ x,
                                    float* __restrict__ out,
                                    const float* __restrict__ mean,
                                    const float* __restrict__ rstd,
                                    const float* __restrict__ weight,
                                    const float* __restrict__ bias,
                                    int N) {
    const int row = blockIdx.y;
    const float m = mean[row];
    const float r = rstd[row];
    const float* row_x = x + (long long)row * N;
    float* row_o = out + (long long)row * N;

    const int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; i += stride) {
        row_o[i] = (row_x[i] - m) * r * weight[i] + bias[i];
    }
}

// ---------------- Host wrapper ----------------
torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight,
                                torch::Tensor bias, double eps) {
    const c10::cuda::CUDAGuard device_guard(x.device());
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
                "All tensors must be CUDA tensors");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();
    TORCH_CHECK(xc.scalar_type() == at::kFloat &&
                wc.scalar_type() == at::kFloat &&
                bc.scalar_type() == at::kFloat,
                "Only float32 is supported by the custom kernel");

    const int64_t N = wc.numel();
    const int64_t rows = (N > 0) ? (xc.numel() / N) : 0;
    auto out = at::empty_like(xc);
    if (rows == 0 || N == 0) return out;
    TORCH_CHECK(N < (int64_t)1 << 31, "Normalized dim too large");
    TORCH_CHECK(rows < (int64_t)1 << 31, "Too many rows");

    auto stream = at::cuda::getCurrentCUDAStream();
    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    const float* bp = bc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    const bool vec4 = (N % 4 == 0) &&
        ((reinterpret_cast<uintptr_t>(xp) & 15) == 0) &&
        ((reinterpret_cast<uintptr_t>(op) & 15) == 0) &&
        ((reinterpret_cast<uintptr_t>(wp) & 15) == 0) &&
        ((reinterpret_cast<uintptr_t>(bp) & 15) == 0);

    int splits;
    int64_t chunk;
    if (vec4) {
        splits = (int)std::min<int64_t>(128, std::max<int64_t>(1, N / (LN_THREADS * 4)));
        chunk  = (((N + splits - 1) / splits) + 3) & ~3LL;  // keep 16B alignment
        splits = (int)((N + chunk - 1) / chunk);
    } else {
        splits = (int)std::min<int64_t>(128, std::max<int64_t>(1, N / LN_THREADS));
        chunk  = (N + splits - 1) / splits;
        splits = (int)((N + chunk - 1) / chunk);
    }

    auto partials = at::empty({rows * (int64_t)splits * 2}, xc.options());
    auto mean     = at::empty({rows}, xc.options());
    auto rstd     = at::empty({rows}, xc.options());
    float* pp = partials.data_ptr<float>();
    float* mp = mean.data_ptr<float>();
    float* rp = rstd.data_ptr<float>();

    const int64_t maxY = 65535;
    // Pass 1: partial reductions (chunk rows to respect the grid-y limit).
    for (int64_t r0 = 0; r0 < rows; r0 += maxY) {
        const int rc = (int)std::min<int64_t>(maxY, rows - r0);
        dim3 grid1(splits, rc);
        if (vec4) {
            ln_partial_sums_vec4<<<grid1, LN_THREADS, 0, stream>>>(
                xp + r0 * N, pp + r0 * splits * 2, (int)N, (int)chunk, splits);
        } else {
            ln_partial_sums_scalar<<<grid1, LN_THREADS, 0, stream>>>(
                xp + r0 * N, pp + r0 * splits * 2, (int)N, (int)chunk, splits);
        }
    }

    // Pass 2: per-row mean / rstd.
    ln_final_stats<<<(int)rows, LN_THREADS, 0, stream>>>(
        pp, mp, rp, splits, (int)N, (float)eps);

    // Pass 3: apply normalization with affine transform.
    const int64_t units = vec4 ? (N >> 2) : N;
    const int64_t col_blocks = (units + LN_THREADS - 1) / LN_THREADS;

    if (col_blocks >= 400 && rows <= 32768) {
        // Preferred: each block owns a column chunk and loops over all rows,
        // so weight/bias are fetched exactly once.
        dim3 grid3((unsigned int)col_blocks);
        if (vec4) {
            ln_apply_rows_vec4<<<grid3, LN_THREADS, 0, stream>>>(
                xp, op, mp, rp, wp, bp, (int)N, (int)rows);
        } else {
            ln_apply_rows_scalar<<<grid3, LN_THREADS, 0, stream>>>(
                xp, op, mp, rp, wp, bp, (int)N, (int)rows);
        }
    } else {
        // Fallback for narrow normalized dims or extremely many rows.
        for (int64_t r0 = 0; r0 < rows; r0 += maxY) {
            const int rc = (int)std::min<int64_t>(maxY, rows - r0);
            if (vec4) {
                const int64_t n4 = N >> 2;
                const int xb = (int)std::min<int64_t>(4096, (n4 + LN_THREADS - 1) / LN_THREADS);
                dim3 grid3(xb, rc);
                ln_normalize_vec4<<<grid3, LN_THREADS, 0, stream>>>(
                    xp + r0 * N, op + r0 * N, mp + r0, rp + r0, wp, bp, (int)N);
            } else {
                const int xb = (int)std::min<int64_t>(4096, (N + LN_THREADS - 1) / LN_THREADS);
                dim3 grid3(xb, rc);
                ln_normalize_scalar<<<grid3, LN_THREADS, 0, stream>>>(
                    xp + r0 * N, op + r0 * N, mp + r0, rp + r0, wp, bp, (int)N);
            }
        }
    }

    return out;
}
"""

layer_norm_cpp_source = (
    "torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, double eps);"
)

# Compile the inline CUDA LayerNorm kernels
_ln_ext = load_inline(
    name="fused_layer_norm_ext_v2",
    cpp_sources=layer_norm_cpp_source,
    cuda_sources=layer_norm_cuda_source,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    LayerNorm model accelerated with custom fused CUDA kernels.
    """

    def __init__(self, normalized_shape: tuple):
        """
        Initializes the LayerNorm layer.

        Args:
            normalized_shape (tuple): Shape of the input tensor to be normalized.
        """
        super(ModelNew, self).__init__()
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Layer Normalization using the custom CUDA kernel.

        Args:
            x (torch.Tensor): Input tensor of shape (*, normalized_shape).

        Returns:
            torch.Tensor: Output tensor with Layer Normalization applied, same shape as input.
        """
        if x.is_cuda and x.dtype == torch.float32:
            return _ln_ext.layernorm_forward(
                x, self.ln.weight, self.ln.bias, float(self.ln.eps)
            )
        # Fallback for CPU tensors or non-float32 dtypes.
        return F.layer_norm(
            x, self.ln.normalized_shape, self.ln.weight, self.ln.bias, self.ln.eps
        )
```