```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 kernels.
# The given model normalizes over (features, dim1, dim2), which means the input
# has very few rows (M = batch_size = 16) and extremely long rows
# (N = 64 * 256 * 256 = 4,194,304). PyTorch's native kernel uses only one
# thread block per row, which severely under-utilizes the GPU for this shape.
# These kernels instead use many thread blocks per row:
#   1) layerNormStatsKernel   - multi-block partial sum / sum-of-squares (double
#                               precision accumulation, float4 vectorized loads)
#   2) layerNormFinalizeKernel- tiny per-block reduction of partials -> mean/rstd
#   3) layerNormApplyKernel   - y = (x - mean) * rstd * weight + bias
#                               (float4 vectorized, fused affine transform)
_layernorm_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

constexpr int kThreads = 256;

__device__ __forceinline__ void blockReduceSum2(double& v0, double& v1,
                                                double* s0, double* s1) {
    const unsigned int mask = 0xffffffffu;
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
#pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v0 += __shfl_down_sync(mask, v0, offset);
        v1 += __shfl_down_sync(mask, v1, offset);
    }
    if (lane == 0) {
        s0[wid] = v0;
        s1[wid] = v1;
    }
    __syncthreads();
    const int nwarps = (blockDim.x + 31) >> 5;
    v0 = (threadIdx.x < nwarps) ? s0[threadIdx.x] : 0.0;
    v1 = (threadIdx.x < nwarps) ? s1[threadIdx.x] : 0.0;
    if (wid == 0) {
#pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            v0 += __shfl_down_sync(mask, v0, offset);
            v1 += __shfl_down_sync(mask, v1, offset);
        }
    }
    // results are valid in thread 0
}

__global__ __launch_bounds__(256) void layerNormStatsKernel(
        const float* __restrict__ x,
        double* __restrict__ psum,
        double* __restrict__ psumsq,
        long N) {
    const int row = blockIdx.y;
    const float* xr = x + (long)row * N;
    double s = 0.0;
    double q = 0.0;
    const long n4 = N & ~3L;
    const long stride = (long)gridDim.x * blockDim.x * 4;
    for (long i = ((long)blockIdx.x * blockDim.x + threadIdx.x) * 4; i < n4;
         i += stride) {
        const float4 v = *reinterpret_cast<const float4*>(xr + i);
        s += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
        q += (double)v.x * v.x + (double)v.y * v.y +
             (double)v.z * v.z + (double)v.w * v.w;
    }
    // scalar tail (only used when N is not divisible by 4)
    for (long i = n4 + (long)blockIdx.x * blockDim.x + threadIdx.x; i < N;
         i += (long)gridDim.x * blockDim.x) {
        const double v = (double)xr[i];
        s += v;
        q += v * v;
    }
    __shared__ double s0[32];
    __shared__ double s1[32];
    blockReduceSum2(s, q, s0, s1);
    if (threadIdx.x == 0) {
        const long idx = (long)row * gridDim.x + blockIdx.x;
        psum[idx]   = s;
        psumsq[idx] = q;
    }
}

__global__ __launch_bounds__(256) void layerNormFinalizeKernel(
        const double* __restrict__ psum,
        const double* __restrict__ psumsq,
        float* __restrict__ mean,
        float* __restrict__ rstd,
        int npart, long N, float eps) {
    const int row = blockIdx.x;
    double s = 0.0;
    double q = 0.0;
    for (int i = threadIdx.x; i < npart; i += blockDim.x) {
        s += psum[(long)row * npart + i];
        q += psumsq[(long)row * npart + i];
    }
    __shared__ double s0[32];
    __shared__ double s1[32];
    blockReduceSum2(s, q, s0, s1);
    if (threadIdx.x == 0) {
        const double mu = s / (double)N;
        double var = q / (double)N - mu * mu;
        var = var > 0.0 ? var : 0.0;
        mean[row] = (float)mu;
        rstd[row] = (float)rsqrt(var + (double)eps);
    }
}

__global__ __launch_bounds__(256) void layerNormApplyKernel(
        const float* __restrict__ x,
        const float* __restrict__ w,
        const float* __restrict__ b,
        const float* __restrict__ mean,
        const float* __restrict__ rstd,
        float* __restrict__ out,
        long N) {
    const int row = blockIdx.y;
    const float m = mean[row];
    const float r = rstd[row];
    const float* xr  = x + (long)row * N;
    float* outr      = out + (long)row * N;
    const long i = ((long)blockIdx.x * blockDim.x + threadIdx.x) * 4;
    if (i + 3 < N) {
        const float4 xv = *reinterpret_cast<const float4*>(xr + i);
        const float4 wv = *reinterpret_cast<const float4*>(w + i);
        const float4 bv = *reinterpret_cast<const float4*>(b + i);
        float4 ov;
        ov.x = (xv.x - m) * r * wv.x + bv.x;
        ov.y = (xv.y - m) * r * wv.y + bv.y;
        ov.z = (xv.z - m) * r * wv.z + bv.z;
        ov.w = (xv.w - m) * r * wv.w + bv.w;
        *reinterpret_cast<float4*>(outr + i) = ov;
    } else if (i < N) {
        for (long j = i; j < N; ++j) {
            outr[j] = (xr[j] - m) * r * w[j] + b[j];
        }
    }
}

}  // namespace

torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight,
                                torch::Tensor bias, int64_t N, double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();
    const long n = (long)N;
    const long m = xc.numel() / n;
    TORCH_CHECK(m >= 1 && m <= 65535, "invalid row count");

    auto out = torch::empty_like(xc);
    const auto stream = at::cuda::getCurrentCUDAStream();

    const int threads = kThreads;
    const long elemsPerBlock = (long)threads * 4 * 8;
    int npart = (int)((n + elemsPerBlock - 1) / elemsPerBlock);
    if (npart < 1) npart = 1;
    if (npart > 1024) npart = 1024;

    auto dblOpts = torch::TensorOptions().dtype(torch::kFloat64).device(xc.device());
    auto fltOpts = torch::TensorOptions().dtype(torch::kFloat32).device(xc.device());
    auto psum   = torch::empty({m, (long)npart}, dblOpts);
    auto psumsq = torch::empty({m, (long)npart}, dblOpts);
    auto mean   = torch::empty({m}, fltOpts);
    auto rstd   = torch::empty({m}, fltOpts);

    dim3 statsGrid((unsigned int)npart, (unsigned int)m);
    layerNormStatsKernel<<<statsGrid, threads, 0, stream>>>(
        xc.data_ptr<float>(), psum.data_ptr<double>(),
        psumsq.data_ptr<double>(), n);

    layerNormFinalizeKernel<<<(unsigned int)m, threads, 0, stream>>>(
        psum.data_ptr<double>(), psumsq.data_ptr<double>(),
        mean.data_ptr<float>(), rstd.data_ptr<float>(), npart, n, (float)eps);

    const long chunks = (n + 3) / 4;
    const int applyBlocks = (int)((chunks + threads - 1) / threads);
    dim3 applyGrid((unsigned int)applyBlocks, (unsigned int)m);
    layerNormApplyKernel<<<applyGrid, threads, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
        mean.data_ptr<float>(), rstd.data_ptr<float>(),
        out.data_ptr<float>(), n);

    return out;
}
"""

_layernorm_cpp_src = (
    "torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, int64_t N, double eps);"
)

_layernorm_ext = load_inline(
    name="custom_layernorm_ext_v1",
    cpp_sources=_layernorm_cpp_src,
    cuda_sources=_layernorm_cuda_src,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    LayerNorm model accelerated with custom CUDA kernels.
    Keeps an nn.LayerNorm submodule so weights/bias/eps semantics (and
    state_dict compatibility) are identical to the reference Model.
    """
    def __init__(self, normalized_shape: tuple):
        super(ModelNew, self).__init__()
        self.normalized_shape = tuple(normalized_shape)
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and self.ln.elementwise_affine
            and self.ln.weight.is_cuda
        ):
            n = 1
            for d in self.normalized_shape:
                n *= d
            rows = x.numel() // n if n > 0 else 0
            if n > 0 and x.numel() % n == 0 and 0 < rows <= 65535:
                return _layernorm_ext.layernorm_forward(
                    x, self.ln.weight, self.ln.bias, n, float(self.ln.eps)
                )
        # Fallback (CPU tensors / non-float32 / no affine): exact reference op.
        return F.layer_norm(
            x, self.normalized_shape, self.ln.weight, self.ln.bias, self.ln.eps
        )
```