```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.
# Strategy: two-kernel approach.
#   Kernel 1: multi-block-per-row partial sum / sum-of-squares reduction
#             (vectorized float4 loads) into a small workspace.
#   Kernel 2: each block cheaply re-reduces the per-row partials to get
#             mean/rstd, then normalizes its chunk (fused affine transform).
_layernorm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <vector>
#include <cstdint>

#define LN_BLOCK_SIZE 256

// Reduces (sum, sumsq) across the whole block; result returned to all threads.
__device__ __forceinline__ void block_reduce_sum2(float &sum, float &sumsq,
                                                  float *sh_sum, float *sh_sumsq) {
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum   += __shfl_down_sync(0xffffffffu, sum, offset);
        sumsq += __shfl_down_sync(0xffffffffu, sumsq, offset);
    }
    if (lane == 0) {
        sh_sum[wid]   = sum;
        sh_sumsq[wid] = sumsq;
    }
    __syncthreads();
    if (wid == 0) {
        int num_warps = (blockDim.x + 31) >> 5;
        sum   = (lane < num_warps) ? sh_sum[lane]   : 0.0f;
        sumsq = (lane < num_warps) ? sh_sumsq[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum   += __shfl_down_sync(0xffffffffu, sum, offset);
            sumsq += __shfl_down_sync(0xffffffffu, sumsq, offset);
        }
        if (lane == 0) {
            sh_sum[0]   = sum;
            sh_sumsq[0] = sumsq;
        }
    }
    __syncthreads();
    sum   = sh_sum[0];
    sumsq = sh_sumsq[0];
}

__global__ void ln_partial_stats_kernel(
    const float *__restrict__ x,
    float *__restrict__ partial_sum,
    float *__restrict__ partial_sumsq,
    int64_t N,
    int blocks_per_row)
{
    int64_t row = blockIdx.x;
    int blk     = blockIdx.y;
    const float *row_x = x + row * N;

    float sum = 0.0f, sumsq = 0.0f;
    int64_t stride = (int64_t)blocks_per_row * blockDim.x;
    int64_t start  = (int64_t)blk * blockDim.x + threadIdx.x;

    if ((N & 3) == 0) {
        int64_t n4 = N >> 2;
        const float4 *row_x4 = reinterpret_cast<const float4 *>(row_x);
        for (int64_t i = start; i < n4; i += stride) {
            float4 v = row_x4[i];
            sum   += (v.x + v.y) + (v.z + v.w);
            sumsq += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
        }
    } else {
        for (int64_t i = start; i < N; i += stride) {
            float v = row_x[i];
            sum   += v;
            sumsq += v * v;
        }
    }

    __shared__ float sh_sum[32];
    __shared__ float sh_sumsq[32];
    block_reduce_sum2(sum, sumsq, sh_sum, sh_sumsq);

    if (threadIdx.x == 0) {
        int64_t idx = row * blocks_per_row + blk;
        partial_sum[idx]   = sum;
        partial_sumsq[idx] = sumsq;
    }
}

__global__ void ln_normalize_kernel(
    const float *__restrict__ x,
    const float *__restrict__ partial_sum,
    const float *__restrict__ partial_sumsq,
    const float *__restrict__ weight,
    const float *__restrict__ bias,
    float *__restrict__ out,
    int64_t N,
    int blocks_per_row,
    float eps)
{
    int64_t row = blockIdx.x;
    int blk     = blockIdx.y;

    // Re-reduce the small per-row partials (cheap, redundant across blocks).
    float sum = 0.0f, sumsq = 0.0f;
    int64_t base = row * blocks_per_row;
    for (int i = threadIdx.x; i < blocks_per_row; i += blockDim.x) {
        sum   += partial_sum[base + i];
        sumsq += partial_sumsq[base + i];
    }
    __shared__ float sh_sum[32];
    __shared__ float sh_sumsq[32];
    block_reduce_sum2(sum, sumsq, sh_sum, sh_sumsq);

    float mean = sum / (float)N;
    float var  = sumsq / (float)N - mean * mean;
    var = fmaxf(var, 0.0f);
    float rstd = rsqrtf(var + eps);

    const float *row_x   = x + row * N;
    float *row_out       = out + row * N;
    int64_t stride = (int64_t)blocks_per_row * blockDim.x;
    int64_t start  = (int64_t)blk * blockDim.x + threadIdx.x;

    if ((N & 3) == 0) {
        int64_t n4 = N >> 2;
        const float4 *row_x4 = reinterpret_cast<const float4 *>(row_x);
        const float4 *w4     = reinterpret_cast<const float4 *>(weight);
        const float4 *b4     = reinterpret_cast<const float4 *>(bias);
        float4 *out4         = reinterpret_cast<float4 *>(row_out);
        for (int64_t i = start; i < n4; i += stride) {
            float4 v = row_x4[i];
            float4 w = w4[i];
            float4 b = b4[i];
            float4 o;
            o.x = (v.x - mean) * rstd * w.x + b.x;
            o.y = (v.y - mean) * rstd * w.y + b.y;
            o.z = (v.z - mean) * rstd * w.z + b.z;
            o.w = (v.w - mean) * rstd * w.w + b.w;
            out4[i] = o;
        }
    } else {
        for (int64_t i = start; i < N; i += stride) {
            row_out[i] = (row_x[i] - mean) * rstd * weight[i] + bias[i];
        }
    }
}

torch::Tensor layernorm_forward(torch::Tensor x,
                                torch::Tensor weight,
                                torch::Tensor bias,
                                std::vector<int64_t> normalized_shape,
                                double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    const c10::cuda::CUDAGuard device_guard(x.device());

    auto x_c = x.contiguous();
    auto w_c = weight.contiguous();
    auto b_c = bias.contiguous();

    int64_t N = 1;
    for (auto s : normalized_shape) N *= s;
    TORCH_CHECK(N > 0, "normalized_shape must be non-empty");
    int64_t M = x_c.numel() / N;

    auto out = torch::empty_like(x_c);
    if (M == 0) return out;

    int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    int64_t target_blocks = (int64_t)sms * 16;
    int64_t max_bpr = (N + LN_BLOCK_SIZE - 1) / LN_BLOCK_SIZE;
    int64_t bpr = (target_blocks + M - 1) / M;
    if (bpr < 1) bpr = 1;
    if (bpr > max_bpr) bpr = max_bpr;
    if (bpr > 65535) bpr = 65535;  // grid.y limit
    int blocks_per_row = (int)bpr;

    auto partial = torch::empty({M, 2 * (int64_t)blocks_per_row}, x_c.options());
    float *partial_ptr   = partial.data_ptr<float>();
    float *partial_sum   = partial_ptr;
    float *partial_sumsq = partial_ptr + M * (int64_t)blocks_per_row;

    dim3 grid((unsigned int)M, (unsigned int)blocks_per_row);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    ln_partial_stats_kernel<<<grid, LN_BLOCK_SIZE, 0, stream>>>(
        x_c.data_ptr<float>(), partial_sum, partial_sumsq, N, blocks_per_row);

    ln_normalize_kernel<<<grid, LN_BLOCK_SIZE, 0, stream>>>(
        x_c.data_ptr<float>(), partial_sum, partial_sumsq,
        w_c.data_ptr<float>(), b_c.data_ptr<float>(),
        out.data_ptr<float>(), N, blocks_per_row, (float)eps);

    return out;
}
"""

_layernorm_cpp_source = (
    "torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, std::vector<int64_t> normalized_shape, double eps);"
)

_layernorm_ext = load_inline(
    name="fused_layernorm_ext_v1",
    cpp_sources=_layernorm_cpp_source,
    cuda_sources=_layernorm_cuda_source,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    LayerNorm model backed by a custom fused CUDA kernel.
    """
    def __init__(self, normalized_shape) -> None:
        super().__init__()
        if isinstance(normalized_shape, int):
            normalized_shape = (normalized_shape,)
        self.normalized_shape = tuple(normalized_shape)
        self.ln = nn.LayerNorm(normalized_shape=self.normalized_shape)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32:
            return _layernorm_ext.layernorm_forward(
                x,
                self.ln.weight,
                self.ln.bias,
                list(self.normalized_shape),
                float(self.ln.eps),
            )
        return F.layer_norm(
            x, self.normalized_shape, self.ln.weight, self.ln.bias, self.ln.eps
        )
```