```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 CUDA LayerNorm:
#  - Splits each row (normalized_shape flattened) across MANY blocks so the
#    whole GPU is used even when the batch (number of rows) is small.
#  - Kernel 1 computes per-block partial sums / sums of squares (vectorized
#    float4 loads, double precision accumulation for accuracy).
#  - Kernel 2 reduces the (few) partials per row inside every block and then
#    normalizes + applies the affine transform in a fused, vectorized pass.
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cstdint>

#define LN_THREADS 256

__device__ __forceinline__ void block_reduce_sum_sq(double& s, double& q) {
    __shared__ double sh_s[LN_THREADS];
    __shared__ double sh_q[LN_THREADS];
    sh_s[threadIdx.x] = s;
    sh_q[threadIdx.x] = q;
    __syncthreads();
    #pragma unroll
    for (int stride = LN_THREADS / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            sh_s[threadIdx.x] += sh_s[threadIdx.x + stride];
            sh_q[threadIdx.x] += sh_q[threadIdx.x + stride];
        }
        __syncthreads();
    }
    s = sh_s[0];
    q = sh_q[0];
}

// ------------------------- vectorized (N % 4 == 0) path -------------------------

__global__ void ln_stats_vec4_kernel(
    const float* __restrict__ x,
    double* __restrict__ part_sum,
    double* __restrict__ part_sumsq,
    const int n4,
    const int chunk4)
{
    const int row = blockIdx.x;
    const int bid = blockIdx.y;
    const float4* x4 = reinterpret_cast<const float4*>(x) + (int64_t)row * n4;
    const int start = bid * chunk4;
    int end = start + chunk4;
    if (end > n4) end = n4;

    double s = 0.0, q = 0.0;
    for (int i = start + threadIdx.x; i < end; i += LN_THREADS) {
        float4 v = x4[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;
    }
    block_reduce_sum_sq(s, q);
    if (threadIdx.x == 0) {
        const int idx = row * gridDim.y + bid;
        part_sum[idx] = s;
        part_sumsq[idx] = q;
    }
}

__global__ void ln_apply_vec4_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    const double* __restrict__ part_sum,
    const double* __restrict__ part_sumsq,
    const int n4,
    const int chunk4,
    const double inv_n,
    const double eps)
{
    const int row = blockIdx.x;
    const int bid = blockIdx.y;
    const int bpr = gridDim.y;

    double s = 0.0, q = 0.0;
    for (int i = threadIdx.x; i < bpr; i += LN_THREADS) {
        s += part_sum[row * bpr + i];
        q += part_sumsq[row * bpr + i];
    }
    block_reduce_sum_sq(s, q);

    const double mean = s * inv_n;
    double var = q * inv_n - mean * mean;
    if (var < 0.0) var = 0.0;
    const float mean_f = (float)mean;
    const float rstd_f = (float)(1.0 / sqrt(var + eps));

    const float4* x4 = reinterpret_cast<const float4*>(x) + (int64_t)row * n4;
    const float4* w4 = reinterpret_cast<const float4*>(w);
    const float4* b4 = reinterpret_cast<const float4*>(b);
    float4* y4 = reinterpret_cast<float4*>(y) + (int64_t)row * n4;

    const int start = bid * chunk4;
    int end = start + chunk4;
    if (end > n4) end = n4;
    for (int i = start + threadIdx.x; i < end; i += LN_THREADS) {
        float4 xv = x4[i];
        float4 wv = w4[i];
        float4 bv = b4[i];
        float4 ov;
        ov.x = (xv.x - mean_f) * rstd_f * wv.x + bv.x;
        ov.y = (xv.y - mean_f) * rstd_f * wv.y + bv.y;
        ov.z = (xv.z - mean_f) * rstd_f * wv.z + bv.z;
        ov.w = (xv.w - mean_f) * rstd_f * wv.w + bv.w;
        y4[i] = ov;
    }
}

// ------------------------- scalar (generic N) path -------------------------

__global__ void ln_stats_scalar_kernel(
    const float* __restrict__ x,
    double* __restrict__ part_sum,
    double* __restrict__ part_sumsq,
    const int n,
    const int chunk)
{
    const int row = blockIdx.x;
    const int bid = blockIdx.y;
    const float* xr = x + (int64_t)row * n;
    const int start = bid * chunk;
    int end = start + chunk;
    if (end > n) end = n;

    double s = 0.0, q = 0.0;
    for (int i = start + threadIdx.x; i < end; i += LN_THREADS) {
        double v = (double)xr[i];
        s += v;
        q += v * v;
    }
    block_reduce_sum_sq(s, q);
    if (threadIdx.x == 0) {
        const int idx = row * gridDim.y + bid;
        part_sum[idx] = s;
        part_sumsq[idx] = q;
    }
}

__global__ void ln_apply_scalar_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    const double* __restrict__ part_sum,
    const double* __restrict__ part_sumsq,
    const int n,
    const int chunk,
    const double inv_n,
    const double eps)
{
    const int row = blockIdx.x;
    const int bid = blockIdx.y;
    const int bpr = gridDim.y;

    double s = 0.0, q = 0.0;
    for (int i = threadIdx.x; i < bpr; i += LN_THREADS) {
        s += part_sum[row * bpr + i];
        q += part_sumsq[row * bpr + i];
    }
    block_reduce_sum_sq(s, q);

    const double mean = s * inv_n;
    double var = q * inv_n - mean * mean;
    if (var < 0.0) var = 0.0;
    const float mean_f = (float)mean;
    const float rstd_f = (float)(1.0 / sqrt(var + eps));

    const float* xr = x + (int64_t)row * n;
    float* yr = y + (int64_t)row * n;
    const int start = bid * chunk;
    int end = start + chunk;
    if (end > n) end = n;
    for (int i = start + threadIdx.x; i < end; i += LN_THREADS) {
        yr[i] = (xr[i] - mean_f) * rstd_f * w[i] + b[i];
    }
}

// ------------------------- host wrapper -------------------------

torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double eps) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(weight.scalar_type() == at::kFloat && bias.scalar_type() == at::kFloat,
                "weight/bias must be float32");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();
    auto y = torch::empty_like(xc);

    const int64_t N = wc.numel();
    if (N == 0) return y;
    const int64_t rows = xc.numel() / N;
    if (rows == 0) return y;
    TORCH_CHECK(N <= INT32_MAX, "normalized dimension too large");

    const int threads = LN_THREADS;
    const int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    // Aim for ~16 blocks per SM spread across the available rows.
    const int64_t want = ((int64_t)sms * 16 + rows - 1) / rows;

    auto stream = at::cuda::getCurrentCUDAStream();
    auto opts = torch::TensorOptions().dtype(torch::kFloat64).device(xc.device());

    if ((N & 3) == 0) {
        const int n4 = (int)(N >> 2);
        const int64_t maxb = (n4 + threads - 1) / threads;
        const int bpr = (int)std::min<int64_t>(std::max<int64_t>(want, 1), maxb);
        const int chunk4 = (n4 + bpr - 1) / bpr;

        auto partials = torch::empty({2, rows * (int64_t)bpr}, opts);
        double* psum = partials.data_ptr<double>();
        double* psq = psum + rows * (int64_t)bpr;

        dim3 grid((unsigned)rows, (unsigned)bpr);
        ln_stats_vec4_kernel<<<grid, threads, 0, stream>>>(
            xc.data_ptr<float>(), psum, psq, n4, chunk4);
        ln_apply_vec4_kernel<<<grid, threads, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
            y.data_ptr<float>(), psum, psq, n4, chunk4, 1.0 / (double)N, eps);
    } else {
        const int n = (int)N;
        const int64_t maxb = (n + threads - 1) / threads;
        const int bpr = (int)std::min<int64_t>(std::max<int64_t>(want, 1), maxb);
        const int chunk = (n + bpr - 1) / bpr;

        auto partials = torch::empty({2, rows * (int64_t)bpr}, opts);
        double* psum = partials.data_ptr<double>();
        double* psq = psum + rows * (int64_t)bpr;

        dim3 grid((unsigned)rows, (unsigned)bpr);
        ln_stats_scalar_kernel<<<grid, threads, 0, stream>>>(
            xc.data_ptr<float>(), psum, psq, n, chunk);
        ln_apply_scalar_kernel<<<grid, threads, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
            y.data_ptr<float>(), psum, psq, n, chunk, 1.0 / (double)N, eps);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

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

_ln_ext = load_inline(
    name="fused_layernorm_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    LayerNorm accelerated with a custom multi-block CUDA kernel.
    """
    def __init__(self, normalized_shape: tuple):
        super().__init__()
        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:
            return _ln_ext.layernorm_forward(
                x, self.ln.weight, self.ln.bias, float(self.ln.eps)
            )
        # Safe 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
        )
```