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

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

_layernorm_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <algorithm>

namespace {

constexpr int kBlockSize = 256;
constexpr int64_t kMaxBlocksPerRow = 512;

// Computes per-row sum and sum-of-squares. Multiple blocks cooperate on each
// row (blockIdx.y == row) using float4 vectorized loads, warp-shuffle
// reductions, and double-precision accumulators + atomics for accuracy.
__global__ void reduce_sums_kernel(const float4* __restrict__ x4,
                                   double* __restrict__ sum,
                                   double* __restrict__ sumsq,
                                   int64_t D4) {
    const int row = blockIdx.y;
    const float4* row_ptr = x4 + static_cast<int64_t>(row) * D4;

    float ls = 0.0f;
    float lsq = 0.0f;
    const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
    for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
         i < D4; i += stride) {
        const float4 v = row_ptr[i];
        ls += v.x + v.y + v.z + v.w;
        lsq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }

    double ds = static_cast<double>(ls);
    double dsq = static_cast<double>(lsq);
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        ds += __shfl_down_sync(0xffffffffu, ds, offset);
        dsq += __shfl_down_sync(0xffffffffu, dsq, offset);
    }

    __shared__ double s_sum[32];
    __shared__ double s_sumsq[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) {
        s_sum[warp] = ds;
        s_sumsq[warp] = dsq;
    }
    __syncthreads();

    if (warp == 0) {
        const int num_warps = (blockDim.x + 31) >> 5;
        ds = (lane < num_warps) ? s_sum[lane] : 0.0;
        dsq = (lane < num_warps) ? s_sumsq[lane] : 0.0;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            ds += __shfl_down_sync(0xffffffffu, ds, offset);
            dsq += __shfl_down_sync(0xffffffffu, dsq, offset);
        }
        if (lane == 0) {
            atomicAdd(&sum[row], ds);
            atomicAdd(&sumsq[row], dsq);
        }
    }
}

// Applies y = (x - mean) * rstd * weight + bias with float4 vectorization.
// mean/rstd are derived from the partial sums directly (fused, no extra pass).
__global__ void normalize_kernel(const float4* __restrict__ x4,
                                 float4* __restrict__ out4,
                                 const float4* __restrict__ w4,
                                 const float4* __restrict__ b4,
                                 const double* __restrict__ sum,
                                 const double* __restrict__ sumsq,
                                 double inv_d, double eps,
                                 int64_t D4) {
    const int row = blockIdx.y;
    const double m = sum[row] * inv_d;
    double var = sumsq[row] * inv_d - m * m;
    var = var < 0.0 ? 0.0 : var;
    const float mf = static_cast<float>(m);
    const float rf = static_cast<float>(rsqrt(var + eps));

    const float4* xrow = x4 + static_cast<int64_t>(row) * D4;
    float4* orow = out4 + static_cast<int64_t>(row) * D4;
    const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
    for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
         i < D4; i += stride) {
        const float4 v = xrow[i];
        const float4 w = w4[i];
        const float4 b = b4[i];
        float4 o;
        o.x = (v.x - mf) * rf * w.x + b.x;
        o.y = (v.y - mf) * rf * w.y + b.y;
        o.z = (v.z - mf) * rf * w.z + b.z;
        o.w = (v.w - mf) * rf * w.w + b.w;
        orow[i] = o;
    }
}

}  // namespace

torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight,
                                torch::Tensor bias, double eps, int64_t D) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(D > 0 && D % 4 == 0, "D must be positive and divisible by 4");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();

    const int64_t rows = xc.numel() / D;
    const int64_t D4 = D / 4;
    TORCH_CHECK(rows >= 1 && rows <= 65535, "invalid number of rows");

    auto out = at::empty_like(xc);
    auto dopts = at::TensorOptions().dtype(at::kDouble).device(xc.device());
    auto sum = at::zeros({rows}, dopts);
    auto sumsq = at::zeros({rows}, dopts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int64_t needed = (D4 + kBlockSize - 1) / kBlockSize;
    const int blocks_per_row =
        static_cast<int>(std::min<int64_t>(needed, kMaxBlocksPerRow));
    dim3 grid(blocks_per_row, static_cast<unsigned int>(rows));

    reduce_sums_kernel<<<grid, kBlockSize, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        sum.data_ptr<double>(),
        sumsq.data_ptr<double>(),
        D4);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    normalize_kernel<<<grid, kBlockSize, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        reinterpret_cast<float4*>(out.data_ptr<float>()),
        reinterpret_cast<const float4*>(wc.data_ptr<float>()),
        reinterpret_cast<const float4*>(bc.data_ptr<float>()),
        sum.data_ptr<double>(),
        sumsq.data_ptr<double>(),
        1.0 / static_cast<double>(D),
        eps,
        D4);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

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

_layernorm_op = load_inline(
    name="fused_layernorm_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 with a custom multi-block-per-row CUDA kernel.
    """

    def __init__(self, normalized_shape: tuple):
        super(ModelNew, self).__init__()
        self.normalized_shape = tuple(normalized_shape)
        self.ln = nn.LayerNorm(normalized_shape=self.normalized_shape)
        self._D = int(math.prod(self.normalized_shape))
        self._eps = float(self.ln.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rows = x.numel() // self._D if self._D > 0 else 0
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and self._D % 4 == 0
            and x.numel() % self._D == 0
            and 0 < rows <= 65535
        ):
            return _layernorm_op.layernorm_forward(
                x, self.ln.weight, self.ln.bias, self._eps, self._D
            )
        return F.layer_norm(
            x, self.normalized_shape, self.ln.weight, self.ln.bias, self._eps
        )
```