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

# Fused LayerNorm CUDA implementation.
#
# Workload shape: a few rows (16) x a huge normalized dim (4M floats, 16MB/row,
# 256MB total).  Exact LayerNorm needs global row stats before any output, so
# the classic lower bound on DRAM traffic is 2 reads of x + 1 write of the
# output.  This revision adds a *single-pass-DRAM* path:
#
#   ln_coop_vec4 is launched cooperatively (guaranteed co-residency).  Each
#   block owns a contiguous chunk of columns, sized so one row of chunks fits
#   in shared memory chip-wide (~38KB/block, 4 blocks/SM).  For every row:
#     1. block loads its chunk of x once from DRAM into shared memory while
#        accumulating sum/sumsq partials,
#     2. a lightweight atomic grid barrier,
#     3. every block redundantly reduces the few hundred partials -> mean/rstd
#        (no second barrier needed),
#     4. the block normalizes its shared-memory chunk and writes it out.
#
#   x is therefore read from DRAM exactly once.  Weights/bias are re-read per
#   row but stay resident in L2 (x loads and output stores use evict-first
#   hints), so DRAM traffic is read(x) + read(w,b once) + write(out) instead
#   of the previous revision's 2*read(x) + write(out) + read(w,b) -> ~1.4x
#   less traffic.  Partials are double-buffered by row parity so only one
#   barrier per row is required.
#
# A tuned 3-kernel fallback (partial sums -> stats -> normalize) handles
# arbitrary shapes, alignments, dtypes, and GPUs without cooperative launch.

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);
        if (lane == 0) { sh_sum[0] = sum; sh_sumsq[0] = sumsq; }
    }
    __syncthreads();
    sum = sh_sum[0];
    sumsq = sh_sumsq[0];
    __syncthreads();
}

__device__ __forceinline__ float ln_f4_sum(const float4 v) {
    return (v.x + v.y) + (v.z + v.w);
}
__device__ __forceinline__ float ln_f4_sq(const float4 v) {
    return fmaf(v.x, v.x, fmaf(v.y, v.y, fmaf(v.z, v.z, v.w * v.w)));
}

// Sense-safe monotonic grid barrier.  bar[0] = arrival ticket counter,
// bar[1] = generation.  Requires all blocks to be co-resident (cooperative
// launch).  Tickets are never reset, so there is no reset race.
__device__ __forceinline__ void ln_grid_barrier(unsigned int* __restrict__ bar,
                                                int nblocks) {
    __syncthreads();
    if (threadIdx.x == 0) {
        volatile unsigned int* genp = (volatile unsigned int*)&bar[1];
        const unsigned int gen = *genp;
        __threadfence();
        const unsigned int ticket = atomicAdd(&bar[0], 1u);
        if ((ticket + 1u) % (unsigned int)nblocks == 0u) {
            __threadfence();
            atomicAdd(&bar[1], 1u);
        } else {
            while (*genp == gen) {
#if __CUDA_ARCH__ >= 700
                __nanosleep(32);
#endif
            }
            __threadfence();
        }
    }
    __syncthreads();
}

// ---------------- Single-DRAM-pass cooperative kernel ---------------------
__global__ void __launch_bounds__(LN_THREADS, 4)
ln_coop_vec4(const float4* __restrict__ x, float4* __restrict__ out,
             const float4* __restrict__ w, const float4* __restrict__ b,
             float* __restrict__ partials, unsigned int* __restrict__ bar,
             int N4, int rows, int chunk, float nInv, float eps) {
    extern __shared__ float4 sdata[];  // chunk float4
    __shared__ float sh_sum[LN_WARPS], sh_sumsq[LN_WARPS];
    __shared__ float s_mean, s_rstd;

    const int tid     = threadIdx.x;
    const int nt      = blockDim.x;
    const int bid     = blockIdx.x;
    const int nblocks = gridDim.x;
    const int c0      = bid * chunk;
    const int cnt     = min(chunk, N4 - c0);
    const int pstride = 2 * nblocks;  // partials per parity buffer

    for (int r = 0; r < rows; ++r) {
        const size_t rowoff = (size_t)r * N4;
        const float4* xr = x + rowoff + c0;

        // Phase 1: load chunk -> shared, accumulate partials.
        float s0 = 0.f, q0 = 0.f, s1 = 0.f, q1 = 0.f;
        float s2 = 0.f, q2 = 0.f, s3 = 0.f, q3 = 0.f;
        int i = tid;
        for (; i + 3 * nt < cnt; i += 4 * nt) {
            const float4 v0 = __ldcs(xr + i);
            const float4 v1 = __ldcs(xr + i + nt);
            const float4 v2 = __ldcs(xr + i + 2 * nt);
            const float4 v3 = __ldcs(xr + i + 3 * nt);
            sdata[i]          = v0;
            sdata[i + nt]     = v1;
            sdata[i + 2 * nt] = v2;
            sdata[i + 3 * nt] = v3;
            s0 += ln_f4_sum(v0); q0 += ln_f4_sq(v0);
            s1 += ln_f4_sum(v1); q1 += ln_f4_sq(v1);
            s2 += ln_f4_sum(v2); q2 += ln_f4_sq(v2);
            s3 += ln_f4_sum(v3); q3 += ln_f4_sq(v3);
        }
        for (; i < cnt; i += nt) {
            const float4 v = __ldcs(xr + i);
            sdata[i] = v;
            s0 += ln_f4_sum(v); q0 += ln_f4_sq(v);
        }
        float sum = (s0 + s1) + (s2 + s3);
        float ssq = (q0 + q1) + (q2 + q3);
        ln_block_reduce(sum, ssq, sh_sum, sh_sumsq);

        float* pbuf = partials + (r & 1) * pstride;
        if (tid == 0) {
            pbuf[bid]            = sum;
            pbuf[nblocks + bid]  = ssq;
        }
        ln_grid_barrier(bar, nblocks);

        // Phase 2: every block redundantly reduces the partials.
        sum = 0.f; ssq = 0.f;
        for (int j = tid; j < nblocks; j += nt) {
            sum += pbuf[j];
            ssq += pbuf[nblocks + j];
        }
        ln_block_reduce(sum, ssq, sh_sum, sh_sumsq);
        if (tid == 0) {
            const float m = sum * nInv;
            float var = ssq * nInv - m * m;
            var = var < 0.f ? 0.f : var;
            s_mean = m;
            s_rstd = rsqrtf(var + eps);
        }
        __syncthreads();
        const float m = s_mean;
        const float t = s_rstd;

        // Phase 3: normalize the shared chunk and write it out.
        float4* orow = out + rowoff + c0;
        const float4* wc = w + c0;
        const float4* bc = b + c0;
        for (i = tid; i + nt < cnt; i += 2 * nt) {
            const float4 v0 = sdata[i];
            const float4 v1 = sdata[i + nt];
            const float4 w0 = wc[i];
            const float4 w1 = wc[i + nt];
            const float4 b0 = bc[i];
            const float4 b1 = bc[i + nt];
            float4 o0, o1;
            o0.x = (v0.x - m) * t * w0.x + b0.x;
            o0.y = (v0.y - m) * t * w0.y + b0.y;
            o0.z = (v0.z - m) * t * w0.z + b0.z;
            o0.w = (v0.w - m) * t * w0.w + b0.w;
            o1.x = (v1.x - m) * t * w1.x + b1.x;
            o1.y = (v1.y - m) * t * w1.y + b1.y;
            o1.z = (v1.z - m) * t * w1.z + b1.z;
            o1.w = (v1.w - m) * t * w1.w + b1.w;
            __stcs(orow + i, o0);
            __stcs(orow + i + nt, o1);
        }
        for (; i < cnt; i += nt) {
            const float4 v = sdata[i];
            const float4 wv = wc[i];
            const float4 bv = bc[i];
            float4 o;
            o.x = (v.x - m) * t * wv.x + bv.x;
            o.y = (v.y - m) * t * wv.y + bv.y;
            o.z = (v.z - m) * t * wv.z + bv.z;
            o.w = (v.w - m) * t * wv.w + bv.w;
            __stcs(orow + i, o);
        }
    }
}

// ---------------- Fallback pass 1: partial sums ----------------
__global__ void ln_partial_sums_vec4(const float* __restrict__ x,
                                     float* __restrict__ ps,
                                     float* __restrict__ pq,
                                     int N, int splits) {
    const int row = blockIdx.y;
    const int n4 = N >> 2;
    const float4* xv = reinterpret_cast<const float4*>(x + (long long)row * N);
    const int gstride = gridDim.x * blockDim.x;
    const int tid0 = blockIdx.x * blockDim.x + threadIdx.x;

    float s[4], q[4];
    #pragma unroll
    for (int k = 0; k < 4; ++k) { s[k] = 0.f; q[k] = 0.f; }
    int i = tid0;
    for (; i + 3 * gstride < n4; i += 4 * gstride) {
        #pragma unroll
        for (int k = 0; k < 4; ++k) {
            const float4 v = __ldcs(xv + i + k * gstride);
            s[k] += ln_f4_sum(v);
            q[k] += ln_f4_sq(v);
        }
    }
    for (; i < n4; i += gstride) {
        const float4 v = __ldcs(xv + i);
        s[0] += ln_f4_sum(v);
        q[0] += ln_f4_sq(v);
    }
    float sum = (s[0] + s[1]) + (s[2] + s[3]);
    float ssq = (q[0] + q[1]) + (q[2] + q[3]);
    __shared__ float sh_sum[LN_WARPS], sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, ssq, sh_sum, sh_sumsq);
    if (threadIdx.x == 0) {
        ps[(long long)row * splits + blockIdx.x] = sum;
        pq[(long long)row * splits + blockIdx.x] = ssq;
    }
}

__global__ void ln_partial_sums_scalar(const float* __restrict__ x,
                                       float* __restrict__ ps,
                                       float* __restrict__ pq,
                                       int N, int splits) {
    const int row = blockIdx.y;
    const float* row_ptr = x + (long long)row * N;
    const int gstride = gridDim.x * blockDim.x;
    float s0 = 0.f, q0 = 0.f, s1 = 0.f, q1 = 0.f;
    float s2 = 0.f, q2 = 0.f, s3 = 0.f, q3 = 0.f;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    for (; i + 3 * gstride < N; i += 4 * gstride) {
        const float v0 = row_ptr[i];
        const float v1 = row_ptr[i + gstride];
        const float v2 = row_ptr[i + 2 * gstride];
        const float v3 = row_ptr[i + 3 * gstride];
        s0 += v0; q0 += v0 * v0;
        s1 += v1; q1 += v1 * v1;
        s2 += v2; q2 += v2 * v2;
        s3 += v3; q3 += v3 * v3;
    }
    for (; i < N; i += gstride) {
        const float v = row_ptr[i];
        s0 += v; q0 += v * v;
    }
    float sum = (s0 + s1) + (s2 + s3);
    float ssq = (q0 + q1) + (q2 + q3);
    __shared__ float sh_sum[LN_WARPS], sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, ssq, sh_sum, sh_sumsq);
    if (threadIdx.x == 0) {
        ps[(long long)row * splits + blockIdx.x] = sum;
        pq[(long long)row * splits + blockIdx.x] = ssq;
    }
}

// ---------------- Fallback pass 2: final per-row statistics ----------------
__global__ void ln_final_stats(const float* __restrict__ ps,
                               const float* __restrict__ pq,
                               float* __restrict__ mean,
                               float* __restrict__ rstd,
                               int splits, int N, float eps) {
    const int row = blockIdx.x;
    float sum = 0.f, ssq = 0.f;
    const long long base = (long long)row * splits;
    for (int i = threadIdx.x; i < splits; i += blockDim.x) {
        sum += ps[base + i];
        ssq += pq[base + i];
    }
    __shared__ float sh_sum[LN_WARPS], sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, ssq, sh_sum, sh_sumsq);
    if (threadIdx.x == 0) {
        const float m = sum / (float)N;
        float var = ssq / (float)N - m * m;
        var = var < 0.f ? 0.f : var;
        mean[row] = m;
        rstd[row] = rsqrtf(var + eps);
    }
}

// ---------------- Fallback pass 3: normalize ----------------
__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;
    }
}

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

    // ---------- Fast path: single-DRAM-pass cooperative kernel ----------
    if (vec4 && rows <= 4096 && N >= (1 << 18)) {
        static int coop_ok = -1;
        static int num_sms = 0;
        if (coop_ok < 0) {
            int dev = 0, attr = 0, sms = 0;
            cudaGetDevice(&dev);
            cudaDeviceGetAttribute(&attr, cudaDevAttrCooperativeLaunch, dev);
            cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
            coop_ok = attr ? 1 : 0;
            num_sms = sms;
        }
        if (coop_ok) {
            const int64_t N4 = N >> 2;
            int grid = 0, chunk = 0;
            size_t shbytes = 0;
            for (int t = 4; t >= 1 && grid == 0; --t) {
                const int g = num_sms * t;
                const int c = (int)((N4 + g - 1) / g);
                const size_t sh = (size_t)c * sizeof(float4) + 64;
                if (sh > 160 * 1024) continue;
                if (sh > 48 * 1024) {
                    cudaFuncSetAttribute((const void*)ln_coop_vec4,
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         (int)sh);
                }
                int occ = 0;
                cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                    &occ, (const void*)ln_coop_vec4, LN_THREADS, sh);
                if (occ >= t) { grid = g; chunk = c; shbytes = sh; }
            }
            if (grid > 0 && (int64_t)grid * chunk >= N4) {
                auto scratch = at::empty({4 * grid + 2},
                                         xc.options().dtype(at::kInt));
                unsigned int* bar = reinterpret_cast<unsigned int*>(
                    scratch.data_ptr<int>());
                float* partials = reinterpret_cast<float*>(bar + 2);
                cudaMemsetAsync(bar, 0, 2 * sizeof(unsigned int), stream);

                const float nInv = 1.0f / (float)N;
                int N4i = (int)N4, rowsi = (int)rows, chunki = chunk;
                float nInvf = nInv, epsf = (float)eps;
                const float4* x4 = reinterpret_cast<const float4*>(xp);
                float4* o4 = reinterpret_cast<float4*>(op);
                const float4* w4 = reinterpret_cast<const float4*>(wp);
                const float4* b4 = reinterpret_cast<const float4*>(bp);
                void* args[] = {(void*)&x4, (void*)&o4, (void*)&w4, (void*)&b4,
                                (void*)&partials, (void*)&bar, (void*)&N4i,
                                (void*)&rowsi, (void*)&chunki, (void*)&nInvf,
                                (void*)&epsf};
                const cudaError_t err = cudaLaunchCooperativeKernel(
                    (const void*)ln_coop_vec4, dim3(grid), dim3(LN_THREADS),
                    args, shbytes, stream.stream());
                TORCH_CHECK(err == cudaSuccess,
                            "cooperative layernorm launch failed: ",
                            cudaGetErrorString(err));
                TORCH_CHECK(cudaGetLastError() == cudaSuccess,
                            "cooperative layernorm kernel failed");
                return out;
            }
        }
    }

    // ---------- Fallback: 3-kernel path ----------
    int splits;
    if (vec4) {
        const int64_t n4 = N >> 2;
        splits = (int)std::max<int64_t>(
            32, std::min<int64_t>(2048, n4 / (LN_THREADS * 8)));
    } else {
        splits = (int)std::max<int64_t>(
            32, std::min<int64_t>(2048, N / (LN_THREADS * 4)));
    }

    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* ps = partials.data_ptr<float>();
    float* pq = ps + rows * (int64_t)splits;
    float* mp = mean.data_ptr<float>();
    float* rp = rstd.data_ptr<float>();

    const int64_t maxY = 65535;
    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, ps + r0 * splits, pq + r0 * splits, (int)N, splits);
        } else {
            ln_partial_sums_scalar<<<grid1, LN_THREADS, 0, stream>>>(
                xp + r0 * N, ps + r0 * splits, pq + r0 * splits, (int)N, splits);
        }
    }

    ln_final_stats<<<(int)rows, LN_THREADS, 0, stream>>>(
        ps, pq, mp, rp, splits, (int)N, (float)eps);

    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_v4",
    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
        )
```