```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 kernels.
#
# Strategy:
#  - The normalized dimension here is very large (64*256*256 = 4M elements per
#    row) with only batch_size rows, so the classic "one block per row" design
#    severely underutilizes the GPU. We instead split every row into many
#    chunks and use a two-kernel approach:
#       1) A multi-block partial reduction computes per-chunk sum & sum-of-
#          squares with float4 vectorized loads (blockIdx.x = row,
#          blockIdx.y = chunk).
#       2) A normalization kernel in which every block first cheaply reduces
#          the tiny per-row partials (fused finalize, no extra kernel launch)
#          and then normalizes its chunk with vectorized float4 loads/stores.
#  - The grid is ordered so that all rows sharing the same chunk are scheduled
#    together, giving L2 reuse of the (large) gamma/beta tensors across rows.
_layer_norm_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define LN_BLOCK 512

// Computes per-chunk partial sums and partial sums of squares.
// grid = (rows, chunks_per_row), block = LN_BLOCK
__global__ void ln_partial_stats_kernel(const float4* __restrict__ x,
                                        float* __restrict__ psum,
                                        float* __restrict__ psumsq,
                                        int n_vec,
                                        int chunk_vec) {
    const int row = blockIdx.x;
    const int chunk = blockIdx.y;
    const int start = chunk * chunk_vec;
    const int end = min(start + chunk_vec, n_vec);
    const float4* xr = x + (size_t)row * (size_t)n_vec;

    float s = 0.0f;
    float q = 0.0f;
    #pragma unroll 4
    for (int i = start + threadIdx.x; i < end; i += LN_BLOCK) {
        float4 v = xr[i];
        s += (v.x + v.y) + (v.z + v.w);
        q += fmaf(v.x, v.x, fmaf(v.y, v.y, fmaf(v.z, v.z, v.w * v.w)));
    }

    // Warp reduction.
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        s += __shfl_down_sync(0xffffffffu, s, off);
        q += __shfl_down_sync(0xffffffffu, q, off);
    }
    __shared__ float sh_s[LN_BLOCK / 32];
    __shared__ float sh_q[LN_BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) {
        sh_s[wid] = s;
        sh_q[wid] = q;
    }
    __syncthreads();
    if (wid == 0) {
        const int nw = LN_BLOCK / 32;
        s = (lane < nw) ? sh_s[lane] : 0.0f;
        q = (lane < nw) ? sh_q[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            s += __shfl_down_sync(0xffffffffu, s, off);
            q += __shfl_down_sync(0xffffffffu, q, off);
        }
        if (lane == 0) {
            const int idx = row * gridDim.y + chunk;
            psum[idx] = s;
            psumsq[idx] = q;
        }
    }
}

// Fused finalize + normalization. Each block first reduces the small set of
// partials for its row to obtain mean/rstd, then normalizes its chunk.
// grid = (rows, chunks_per_row), block = LN_BLOCK
__global__ void ln_normalize_kernel(const float4* __restrict__ x,
                                    const float4* __restrict__ w,
                                    const float4* __restrict__ b,
                                    const float* __restrict__ psum,
                                    const float* __restrict__ psumsq,
                                    float4* __restrict__ y,
                                    int n_vec,
                                    int chunk_vec,
                                    float inv_n,
                                    float eps) {
    const int row = blockIdx.x;
    const int chunk = blockIdx.y;
    const int C = gridDim.y;

    // Finalize stats for this row from the partials (cheap: C <= few hundred).
    const float* pr = psum + (size_t)row * C;
    const float* qr = psumsq + (size_t)row * C;
    float s = 0.0f;
    float q = 0.0f;
    for (int i = threadIdx.x; i < C; i += LN_BLOCK) {
        s += pr[i];
        q += qr[i];
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        s += __shfl_down_sync(0xffffffffu, s, off);
        q += __shfl_down_sync(0xffffffffu, q, off);
    }
    __shared__ float sh_s[LN_BLOCK / 32];
    __shared__ float sh_q[LN_BLOCK / 32];
    __shared__ float sh_mean;
    __shared__ float sh_rstd;
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) {
        sh_s[wid] = s;
        sh_q[wid] = q;
    }
    __syncthreads();
    if (wid == 0) {
        const int nw = LN_BLOCK / 32;
        s = (lane < nw) ? sh_s[lane] : 0.0f;
        q = (lane < nw) ? sh_q[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            s += __shfl_down_sync(0xffffffffu, s, off);
            q += __shfl_down_sync(0xffffffffu, q, off);
        }
        if (lane == 0) {
            const float m = s * inv_n;
            const float var = fmaf(-m, m, q * inv_n);
            sh_mean = m;
            sh_rstd = rsqrtf(var + eps);
        }
    }
    __syncthreads();
    const float m = sh_mean;
    const float r = sh_rstd;

    // Normalize this block's chunk with vectorized accesses.
    const size_t base = (size_t)row * (size_t)n_vec;
    const float4* xr = x + base;
    float4* yr = y + base;
    const int start = chunk * chunk_vec;
    const int end = min(start + chunk_vec, n_vec);
    #pragma unroll 4
    for (int i = start + threadIdx.x; i < end; i += LN_BLOCK) {
        float4 v = xr[i];
        float4 g = w[i];
        float4 bb = b[i];
        float4 o;
        o.x = fmaf((v.x - m) * r, g.x, bb.x);
        o.y = fmaf((v.y - m) * r, g.y, bb.y);
        o.z = fmaf((v.z - m) * r, g.z, bb.z);
        o.w = fmaf((v.w - m) * r, g.w, bb.w);
        yr[i] = o;
    }
}

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

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

    int64_t n = 1;
    for (int64_t i = xc.dim() - normalized_ndim; i < xc.dim(); ++i) {
        n *= xc.size(i);
    }
    const int64_t rows = xc.numel() / n;
    TORCH_CHECK(n % 4 == 0, "normalized size must be divisible by 4");
    const int n_vec = (int)(n >> 2);

    auto out = at::empty_like(xc);

    // Choose chunks per row so that rows * chunks saturates the GPU.
    const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int target_blocks = sm_count * 8;
    int64_t C = (target_blocks + rows - 1) / rows;
    if (C < 1) C = 1;
    int64_t chunk_vec = ((int64_t)n_vec + C - 1) / C;
    C = ((int64_t)n_vec + chunk_vec - 1) / chunk_vec;
    if (C < 1) C = 1;
    TORCH_CHECK(C <= 65535, "too many chunks");

    auto partial = at::empty({2, rows, C}, xc.options());
    float* psum = partial.data_ptr<float>();
    float* psumsq = psum + rows * C;

    const dim3 grid((unsigned int)rows, (unsigned int)C);
    auto stream = at::cuda::getCurrentCUDAStream();

    ln_partial_stats_kernel<<<grid, LN_BLOCK, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        psum, psumsq, n_vec, (int)chunk_vec);

    const float inv_n = 1.0f / (float)n;
    ln_normalize_kernel<<<grid, LN_BLOCK, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        reinterpret_cast<const float4*>(wc.data_ptr<float>()),
        reinterpret_cast<const float4*>(bc.data_ptr<float>()),
        psum, psumsq,
        reinterpret_cast<float4*>(out.data_ptr<float>()),
        n_vec, (int)chunk_vec, inv_n, (float)eps);

    return out;
}
"""

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

# Compile the fused LayerNorm CUDA extension.
_fused_layer_norm = load_inline(
    name="fused_layer_norm_ext",
    cpp_sources=_layer_norm_cpp_src,
    cuda_sources=_layer_norm_cuda_src,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    LayerNorm model accelerated with a custom fused CUDA kernel.
    """

    def __init__(self, normalized_shape: tuple):
        super(ModelNew, self).__init__()
        # Keep an nn.LayerNorm submodule so parameter names/state_dict match
        # the reference model exactly (weight, bias, eps).
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)
        if isinstance(normalized_shape, int):
            ns = (normalized_shape,)
        else:
            ns = tuple(normalized_shape)
        self._normalized_shape = ns
        self._normalized_ndim = len(ns)
        n = 1
        for s in ns:
            n *= int(s)
        self._normalized_numel = n
        # The vectorized kernel requires the normalized size divisible by 4.
        self._fast_ok = (n % 4 == 0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (
            self._fast_ok
            and x.is_cuda
            and x.dtype == torch.float32
            and self.ln.elementwise_affine
        ):
            return _fused_layer_norm.layernorm_forward(
                x,
                self.ln.weight,
                self.ln.bias,
                self._normalized_ndim,
                float(self.ln.eps),
            )
        # Safe fallback for any configuration the custom kernel does not cover.
        return F.layer_norm(
            x,
            self._normalized_shape,
            self.ln.weight if self.ln.elementwise_affine else None,
            self.ln.bias if self.ln.elementwise_affine else None,
            float(self.ln.eps),
        )


batch_size = 16
features = 64
dim1 = 256
dim2 = 256


def get_inputs():
    x = torch.rand(batch_size, features, dim1, dim2)
    return [x]


def get_init_inputs():
    return [(features, dim1, dim2)]
```