```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 normalized shape (64, 256, 256) means each "row" contains
# 64*256*256 = 4,194,304 elements and there are only 16 rows, which
# under-utilizes the GPU in the naive "one block per row" scheme.
# We therefore use:
#   1) A multi-block-per-row vectorized (float4) stats kernel that
#      accumulates sum / sum-of-squares with warp shuffles + shared memory,
#      and atomically accumulates per-row partials in double precision.
#   2) A vectorized apply kernel that reads the stats, computes
#      mean / rstd in double, and writes the normalized, affine-transformed
#      output in a single pass.

layernorm_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

constexpr int kStatsBlock = 512;
constexpr int kApplyBlock = 256;

// ---------------------------------------------------------------------------
// Kernel 1: per-row partial sums / sums of squares (vectorized float4).
// grid = (blocks_per_row, rows)
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(kStatsBlock) ln_stats_kernel(
        const float* __restrict__ x,
        double* __restrict__ stats,   // [rows, 2]: sum, sumsq
        long Nv,                      // N / 4 (float4 elems per row)
        int blocks_per_row) {
    const int row = blockIdx.y;
    const float4* xr = reinterpret_cast<const float4*>(x + (long)row * (Nv << 2));

    float ssum = 0.f;
    float ssq  = 0.f;
    const long stride = (long)blocks_per_row * kStatsBlock;
    for (long i = (long)blockIdx.x * kStatsBlock + threadIdx.x; i < Nv; i += stride) {
        const float4 v = xr[i];
        ssum += v.x + v.y + v.z + v.w;
        ssq  += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }

    __shared__ float ssum_sh[kStatsBlock / 32];
    __shared__ float ssq_sh[kStatsBlock / 32];

    const unsigned mask = 0xffffffffu;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        ssum += __shfl_down_sync(mask, ssum, off);
        ssq  += __shfl_down_sync(mask, ssq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        ssum_sh[wid] = ssum;
        ssq_sh[wid]  = ssq;
    }
    __syncthreads();

    if (wid == 0) {
        const int nw = kStatsBlock >> 5;
        ssum = (lane < nw) ? ssum_sh[lane] : 0.f;
        ssq  = (lane < nw) ? ssq_sh[lane]  : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            ssum += __shfl_down_sync(mask, ssum, off);
            ssq  += __shfl_down_sync(mask, ssq, off);
        }
        if (lane == 0) {
            atomicAdd(&stats[row * 2 + 0], (double)ssum);
            atomicAdd(&stats[row * 2 + 1], (double)ssq);
        }
    }
}

// ---------------------------------------------------------------------------
// Kernel 2: apply normalization + affine transform (vectorized float4).
// grid = (blocks_per_row, rows)
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(kApplyBlock) ln_apply_kernel(
        const float4* __restrict__ x,
        const float4* __restrict__ w,
        const float4* __restrict__ b,
        float4* __restrict__ out,
        const double* __restrict__ stats,
        long Nv,               // float4 elems per row
        double invN,
        float eps,
        int blocks_per_row) {
    const int row = blockIdx.y;
    const double sum = stats[row * 2 + 0];
    const double sq  = stats[row * 2 + 1];
    const double mean = sum * invN;
    double var = sq * invN - mean * mean;
    var = var < 0.0 ? 0.0 : var;
    const float rstd  = rsqrtf((float)(var + (double)eps));
    const float meanf = (float)mean;

    const float4* xr = x + (long)row * Nv;
    float4* outr     = out + (long)row * Nv;

    const long stride = (long)blocks_per_row * kApplyBlock;
    for (long i = (long)blockIdx.x * kApplyBlock + threadIdx.x; i < Nv; i += stride) {
        const float4 xv = xr[i];
        const float4 wv = w[i];
        const float4 bv = b[i];
        float4 o;
        o.x = (xv.x - meanf) * rstd * wv.x + bv.x;
        o.y = (xv.y - meanf) * rstd * wv.y + bv.y;
        o.z = (xv.z - meanf) * rstd * wv.z + bv.z;
        o.w = (xv.w - meanf) * rstd * wv.w + bv.w;
        outr[i] = o;
    }
}

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

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = b.contiguous();
    auto out = torch::empty_like(xc);

    const long N = wc.numel();
    const long total = xc.numel();
    const long rows_l = total / N;
    TORCH_CHECK(rows_l <= 65535, "too many rows");
    TORCH_CHECK((N & 3) == 0, "normalized size must be divisible by 4");
    const int rows = (int)rows_l;

    auto stats = torch::zeros({rows * 2}, x.options().dtype(at::kDouble));

    const long Nv = N >> 2;
    auto stream = at::cuda::getCurrentCUDAStream();

    // ---- stats ----
    int sbpr = (int)((Nv + (long)kStatsBlock * 8 - 1) / ((long)kStatsBlock * 8));
    sbpr = sbpr < 1 ? 1 : (sbpr > 256 ? 256 : sbpr);
    dim3 sgrid(sbpr, rows);
    ln_stats_kernel<<<sgrid, kStatsBlock, 0, stream>>>(
            xc.data_ptr<float>(), stats.data_ptr<double>(), Nv, sbpr);

    // ---- apply ----
    int abpr = (int)((Nv + (long)kApplyBlock * 8 - 1) / ((long)kApplyBlock * 8));
    abpr = abpr < 1 ? 1 : (abpr > 512 ? 512 : abpr);
    dim3 agrid(abpr, rows);
    ln_apply_kernel<<<agrid, kApplyBlock, 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>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            stats.data_ptr<double>(),
            Nv, 1.0 / (double)N, (float)eps, abpr);

    return out;
}
"""

layernorm_cpp_source = (
    "torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b, double eps);"
)

# Compile the inline CUDA kernels
_layernorm_ext = load_inline(
    name="fused_layernorm_ext",
    cpp_sources=layernorm_cpp_source,
    cuda_sources=layernorm_cuda_source,
    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__()
        self.normalized_shape = tuple(normalized_shape)
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)
        self.ext = _layernorm_ext

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        N = self.ln.weight.numel()
        if (x.is_cuda and x.dtype == torch.float32 and x.is_contiguous()
                and self.ln.elementwise_affine and (N % 4 == 0)
                and (x.numel() // N) <= 65535):
            return self.ext.layernorm_forward(
                x, self.ln.weight, self.ln.bias, float(self.ln.eps))
        return F.layer_norm(
            x, self.normalized_shape, self.ln.weight, self.ln.bias, self.ln.eps)


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


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


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