```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
from torch.utils.cpp_extension import load_inline

# Custom CUDA implementation of LayerNorm.
#
# The normalized shape here is (64, 256, 256), i.e. each "row" that has to be
# normalized contains 64*256*256 = 4,194,304 elements and there are only 16
# rows. A single-block-per-row kernel (like the standard PyTorch one) severely
# underutilizes the GPU in this regime, so we use:
#   1) A multi-block-per-row partial reduction kernel (sum / sum of squares)
#      with float4 vectorized loads and warp-shuffle reductions.
#   2) A vectorized normalization kernel where each block redundantly combines
#      the few partials (cheap, held in shared memory) to get mean / rstd and
#      then applies y = (x - mean) * rstd * weight + bias.
# Scalar fallback kernels are provided for row sizes not divisible by 4.

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

#define THREADS 256
#define MAX_PARTS 128

// Reduce (sum, sumsq) accumulated in a thread across the whole block.
// Result is written to *out_s / *out_q by a single thread.
__device__ __forceinline__ void block_reduce_ss(float s, float q,
                                                float* out_s, float* out_q) {
    for (int o = 16; o > 0; o >>= 1) {
        s += __shfl_down_sync(0xffffffffu, s, o);
        q += __shfl_down_sync(0xffffffffu, q, o);
    }
    __shared__ float ws[THREADS / 32];
    __shared__ float wq[THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) { ws[wid] = s; wq[wid] = q; }
    __syncthreads();
    if (wid == 0) {
        s = (lane < THREADS / 32) ? ws[lane] : 0.f;
        q = (lane < THREADS / 32) ? wq[lane] : 0.f;
        for (int o = 16; o > 0; o >>= 1) {
            s += __shfl_down_sync(0xffffffffu, s, o);
            q += __shfl_down_sync(0xffffffffu, q, o);
        }
        if (lane == 0) { *out_s = s; *out_q = q; }
    }
}

// Combine the per-block partials of one row into mean / rstd.
__device__ __forceinline__ void finalize_stats(const float* __restrict__ psum,
                                               const float* __restrict__ psumsq,
                                               int row, int parts,
                                               long long row_size, float eps,
                                               float& mean, float& rstd) {
    __shared__ float sp[MAX_PARTS];
    __shared__ float sqp[MAX_PARTS];
    if (threadIdx.x < parts) {
        sp[threadIdx.x]  = psum[(long long)row * parts + threadIdx.x];
        sqp[threadIdx.x] = psumsq[(long long)row * parts + threadIdx.x];
    }
    __syncthreads();
    float s = 0.f, q = 0.f;
    for (int i = 0; i < parts; ++i) { s += sp[i]; q += sqp[i]; }
    mean = s / (float)row_size;
    float var = q / (float)row_size - mean * mean;
    var = fmaxf(var, 0.f);
    rstd = rsqrtf(var + eps);
}

// -------------------- vectorized kernels (row_size % 4 == 0) --------------------

__global__ void ln_partial_vec_kernel(const float* __restrict__ x,
                                      float* __restrict__ psum,
                                      float* __restrict__ psumsq,
                                      long long row_size) {
    const int row = blockIdx.y;
    const long long n4 = row_size >> 2;
    const float4* xr = reinterpret_cast<const float4*>(x + (long long)row * row_size);
    float s = 0.f, q = 0.f;
    const long long stride = (long long)gridDim.x * THREADS;
    for (long long i = (long long)blockIdx.x * THREADS + threadIdx.x; i < n4; i += stride) {
        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)));
    }
    const int parts = gridDim.x;
    block_reduce_ss(s, q,
                    &psum[(long long)row * parts + blockIdx.x],
                    &psumsq[(long long)row * parts + blockIdx.x]);
}

__global__ void ln_normalize_vec_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        const float* __restrict__ b,
                                        const float* __restrict__ psum,
                                        const float* __restrict__ psumsq,
                                        float* __restrict__ y,
                                        long long row_size,
                                        int parts,
                                        float eps) {
    const int row = blockIdx.y;
    float mean, rstd;
    finalize_stats(psum, psumsq, row, parts, row_size, eps, mean, rstd);

    const long long n4 = row_size >> 2;
    const long long base = (long long)row * row_size;
    const float4* xr = reinterpret_cast<const float4*>(x + base);
    const float4* wr = reinterpret_cast<const float4*>(w);
    const float4* br = reinterpret_cast<const float4*>(b);
    float4* yr = reinterpret_cast<float4*>(y + base);
    const long long stride = (long long)gridDim.x * THREADS;
    for (long long i = (long long)blockIdx.x * THREADS + threadIdx.x; i < n4; i += stride) {
        float4 xv = xr[i];
        float4 wv = wr[i];
        float4 bv = br[i];
        float4 o;
        o.x = (xv.x - mean) * rstd * wv.x + bv.x;
        o.y = (xv.y - mean) * rstd * wv.y + bv.y;
        o.z = (xv.z - mean) * rstd * wv.z + bv.z;
        o.w = (xv.w - mean) * rstd * wv.w + bv.w;
        yr[i] = o;
    }
}

// -------------------- scalar fallback kernels --------------------

__global__ void ln_partial_scalar_kernel(const float* __restrict__ x,
                                         float* __restrict__ psum,
                                         float* __restrict__ psumsq,
                                         long long row_size) {
    const int row = blockIdx.y;
    const float* xr = x + (long long)row * row_size;
    float s = 0.f, q = 0.f;
    const long long stride = (long long)gridDim.x * THREADS;
    for (long long i = (long long)blockIdx.x * THREADS + threadIdx.x; i < row_size; i += stride) {
        float v = xr[i];
        s += v;
        q += v * v;
    }
    const int parts = gridDim.x;
    block_reduce_ss(s, q,
                    &psum[(long long)row * parts + blockIdx.x],
                    &psumsq[(long long)row * parts + blockIdx.x]);
}

__global__ void ln_normalize_scalar_kernel(const float* __restrict__ x,
                                           const float* __restrict__ w,
                                           const float* __restrict__ b,
                                           const float* __restrict__ psum,
                                           const float* __restrict__ psumsq,
                                           float* __restrict__ y,
                                           long long row_size,
                                           int parts,
                                           float eps) {
    const int row = blockIdx.y;
    float mean, rstd;
    finalize_stats(psum, psumsq, row, parts, row_size, eps, mean, rstd);

    const long long base = (long long)row * row_size;
    const long long stride = (long long)gridDim.x * THREADS;
    for (long long i = (long long)blockIdx.x * THREADS + threadIdx.x; i < row_size; i += stride) {
        y[base + i] = (x[base + i] - mean) * rstd * 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(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be float32");

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

    const long long row_size = (long long)wc.numel();
    TORCH_CHECK(row_size > 0, "normalized shape must be non-empty");
    const long long total = (long long)xc.numel();
    TORCH_CHECK(total % row_size == 0, "input shape does not match normalized shape");
    const long long num_rows = total / row_size;
    TORCH_CHECK(num_rows <= 65535, "too many rows for this kernel");

    auto y = torch::empty_like(xc);
    if (total == 0) return y;

    const bool vec = (row_size % 4 == 0);
    const long long units = vec ? (row_size / 4) : row_size;

    long long p = (units + THREADS - 1) / THREADS;
    if (p < 1) p = 1;
    if (p > MAX_PARTS) p = MAX_PARTS;
    const int parts = (int)p;

    auto wsbuf = torch::empty({2 * num_rows * parts}, xc.options());
    float* psum = wsbuf.data_ptr<float>();
    float* psumsq = psum + num_rows * parts;

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

    dim3 gridR((unsigned)parts, (unsigned)num_rows);
    long long nb = (units + THREADS - 1) / THREADS;
    if (nb < 1) nb = 1;
    if (nb > 4096) nb = 4096;
    dim3 gridN((unsigned)nb, (unsigned)num_rows);

    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    const float* bp = bc.data_ptr<float>();
    float* yp = y.data_ptr<float>();

    if (vec) {
        ln_partial_vec_kernel<<<gridR, THREADS, 0, stream>>>(xp, psum, psumsq, row_size);
        ln_normalize_vec_kernel<<<gridN, THREADS, 0, stream>>>(xp, wp, bp, psum, psumsq,
                                                               yp, row_size, parts, (float)eps);
    } else {
        ln_partial_scalar_kernel<<<gridR, THREADS, 0, stream>>>(xp, psum, psumsq, row_size);
        ln_normalize_scalar_kernel<<<gridN, THREADS, 0, stream>>>(xp, wp, bp, psum, psumsq,
                                                                  yp, row_size, parts, (float)eps);
    }
    return y;
}
"""

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

# Compile the inline CUDA LayerNorm kernels
_layernorm_ext = load_inline(
    name="fast_layernorm_multi_block",
    cpp_sources=_layernorm_cpp_src,
    cuda_sources=_layernorm_cuda_src,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    LayerNorm model accelerated with a custom multi-block-per-row CUDA kernel.
    """
    def __init__(self, normalized_shape: tuple):
        super(ModelNew, self).__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 and self.ln.elementwise_affine
                and self.ln.weight.dtype == torch.float32
                and self.ln.bias.dtype == torch.float32):
            return _layernorm_ext.layernorm_forward(
                x, self.ln.weight, self.ln.bias, float(self.ln.eps)
            )
        # Fallback for any configuration the custom kernel does not cover.
        return self.ln(x)


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)]
```