```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 multi-block-per-row LayerNorm CUDA implementation.
#
# The benchmark shape normalizes over (64, 256, 256) = 4M elements per row with
# only 16 rows. A standard "one block per row" LayerNorm would only launch 16
# blocks and severely underutilize the GPU. Here each row is processed by many
# blocks:
#   Kernel 1: every block computes a partial sum / sum-of-squares of its slice
#             of the row, atomically accumulates them in a double workspace, and
#             the last block to finish finalizes mean/rstd for the row.
#   Kernel 2: normalizes the row with the affine transform using float4 I/O.

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

#define LN_THREADS 256

__device__ __forceinline__ float ln_warp_sum(float v) {
#pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        v += __shfl_down_sync(0xffffffffu, v, off);
    }
    return v;
}

// Pass 1: partial reduction + finalize (performed by the last block of each row).
__global__ void __launch_bounds__(LN_THREADS)
ln_partial_kernel(const float* __restrict__ x,
                  double* __restrict__ g_sum,
                  double* __restrict__ g_sumsq,
                  unsigned int* __restrict__ counter,
                  float* __restrict__ mean,
                  float* __restrict__ rstd,
                  const int vec_size,      // row_size / 4
                  const int row_size,
                  const int blocks_per_row,
                  const float eps) {
    const int row = blockIdx.x / blocks_per_row;
    const int blk = blockIdx.x - row * blocks_per_row;

    const float4* row_ptr = reinterpret_cast<const float4*>(x) + (size_t)row * vec_size;

    float s = 0.f, sq = 0.f;
    const int stride = blocks_per_row * LN_THREADS;
    for (int i = blk * LN_THREADS + threadIdx.x; i < vec_size; i += stride) {
        const float4 v = row_ptr[i];
        s  += (v.x + v.y) + (v.z + v.w);
        sq += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
    }

    s  = ln_warp_sum(s);
    sq = ln_warp_sum(sq);

    __shared__ float sh_sum[LN_THREADS / 32];
    __shared__ float sh_sq[LN_THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        sh_sum[wid] = s;
        sh_sq[wid]  = sq;
    }
    __syncthreads();

    if (wid == 0) {
        const int nwarp = LN_THREADS / 32;
        s  = (lane < nwarp) ? sh_sum[lane] : 0.f;
        sq = (lane < nwarp) ? sh_sq[lane]  : 0.f;
        s  = ln_warp_sum(s);
        sq = ln_warp_sum(sq);
        if (lane == 0) {
            atomicAdd(g_sum + row, (double)s);
            atomicAdd(g_sumsq + row, (double)sq);
            __threadfence();
            const unsigned int prev = atomicAdd(counter + row, 1u);
            if (prev == (unsigned int)(blocks_per_row - 1)) {  // last block of this row
                const double tot = atomicAdd(g_sum + row, 0.0);
                const double tsq = atomicAdd(g_sumsq + row, 0.0);
                const double m = tot / (double)row_size;
                double var = tsq / (double)row_size - m * m;
                if (var < 0.0) var = 0.0;
                mean[row] = (float)m;
                rstd[row] = (float)rsqrt(var + (double)eps);
            }
        }
    }
}

// Pass 2: normalize + affine transform, vectorized with float4.
__global__ void __launch_bounds__(LN_THREADS)
ln_normalize_kernel(const float* __restrict__ x,
                    float* __restrict__ out,
                    const float* __restrict__ w,
                    const float* __restrict__ b,
                    const float* __restrict__ mean,
                    const float* __restrict__ rstd,
                    const int vec_size,
                    const int blocks_per_row,
                    const bool affine) {
    const int row = blockIdx.x / blocks_per_row;
    const int blk = blockIdx.x - row * blocks_per_row;
    const float m = mean[row];
    const float r = rstd[row];

    const size_t row_off = (size_t)row * vec_size;
    const float4* in_row  = reinterpret_cast<const float4*>(x) + row_off;
    float4* out_row       = reinterpret_cast<float4*>(out) + row_off;
    const float4* w4 = reinterpret_cast<const float4*>(w);
    const float4* b4 = reinterpret_cast<const float4*>(b);

    const int stride = blocks_per_row * LN_THREADS;
    if (affine) {
        for (int i = blk * LN_THREADS + threadIdx.x; i < vec_size; i += stride) {
            const float4 v  = in_row[i];
            const float4 wv = w4[i];
            const float4 bv = b4[i];
            float4 o;
            o.x = (v.x - m) * r * wv.x + bv.x;
            o.y = (v.y - m) * r * wv.y + bv.y;
            o.z = (v.z - m) * r * wv.z + bv.z;
            o.w = (v.w - m) * r * wv.w + bv.w;
            out_row[i] = o;
        }
    } else {
        for (int i = blk * LN_THREADS + threadIdx.x; i < vec_size; i += stride) {
            const float4 v = in_row[i];
            float4 o;
            o.x = (v.x - m) * r;
            o.y = (v.y - m) * r;
            o.z = (v.z - m) * r;
            o.w = (v.w - m) * r;
            out_row[i] = o;
        }
    }
}

torch::Tensor layernorm_forward(torch::Tensor x,
                                torch::Tensor weight,
                                torch::Tensor bias,
                                int64_t normalized_size,
                                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(normalized_size > 0, "normalized_size must be positive");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();
    const bool affine = wc.numel() > 0;

    const int64_t row_size = normalized_size;
    const int64_t num_rows = xc.numel() / row_size;
    TORCH_CHECK(num_rows > 0 && num_rows * row_size == xc.numel(), "input shape mismatch");
    if (affine) {
        TORCH_CHECK(wc.numel() == row_size && bc.numel() == row_size,
                    "weight/bias must match normalized_shape");
    }

    // Rare fallback for shapes the vectorized path doesn't cover.
    if ((row_size & 3) != 0 || row_size > INT32_MAX || num_rows > INT32_MAX) {
        auto x2 = xc.reshape({num_rows, row_size});
        c10::optional<at::Tensor> wopt = c10::nullopt;
        c10::optional<at::Tensor> bopt = c10::nullopt;
        if (affine) {
            wopt = wc.reshape({row_size});
            bopt = bc.reshape({row_size});
        }
        auto y2 = at::layer_norm(x2, {row_size}, wopt, bopt, eps, false);
        return y2.reshape(xc.sizes());
    }

    // Guarantee 16-byte alignment for float4 accesses.
    if ((reinterpret_cast<uintptr_t>(xc.data_ptr()) & 0xF) != 0) xc = xc.clone();
    if (affine && (reinterpret_cast<uintptr_t>(wc.data_ptr()) & 0xF) != 0) wc = wc.clone();
    if (affine && (reinterpret_cast<uintptr_t>(bc.data_ptr()) & 0xF) != 0) bc = bc.clone();

    auto out = at::empty_like(xc);

    const float* xptr = xc.data_ptr<float>();
    const float* wptr = affine ? wc.data_ptr<float>() : nullptr;
    const float* bptr = affine ? bc.data_ptr<float>() : nullptr;
    float* optr = out.data_ptr<float>();

    // Workspace layout: [sum(double) | sumsq(double) | counter(uint) | mean(float) | rstd(float)]
    auto ws = at::zeros({num_rows * 28},
                        at::TensorOptions().dtype(at::kByte).device(xc.device()));
    uint8_t* wsp = ws.data_ptr<uint8_t>();
    double* g_sum        = reinterpret_cast<double*>(wsp);
    double* g_sumsq      = g_sum + num_rows;
    unsigned int* counter = reinterpret_cast<unsigned int*>(g_sumsq + num_rows);
    float* mean          = reinterpret_cast<float*>(counter + num_rows);
    float* rstd          = mean + num_rows;

    const int64_t vec_size = row_size >> 2;
    const int64_t by_work = std::max<int64_t>(1, (vec_size + LN_THREADS * 16 - 1) / (LN_THREADS * 16));
    const int64_t by_rows = std::max<int64_t>(1, 4096 / num_rows);
    const int blocks_per_row = (int)std::min(by_work, by_rows);

    const int grid = blocks_per_row * (int)num_rows;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    ln_partial_kernel<<<grid, LN_THREADS, 0, stream>>>(
        xptr, g_sum, g_sumsq, counter, mean, rstd,
        (int)vec_size, (int)row_size, blocks_per_row, (float)eps);

    ln_normalize_kernel<<<grid, LN_THREADS, 0, stream>>>(
        xptr, optr, wptr, bptr, mean, rstd,
        (int)vec_size, blocks_per_row, affine);

    return out;
}
"""

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

_ln_ext = load_inline(
    name="fast_multi_block_layernorm",
    cpp_sources=_layernorm_cpp_src,
    cuda_sources=_layernorm_cuda_src,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    LayerNorm model backed by a custom multi-block-per-row CUDA kernel.
    """
    def __init__(self, normalized_shape: tuple):
        super().__init__()
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)
        ns = 1
        for d in self.ln.normalized_shape:
            ns *= int(d)
        self.normalized_size = ns

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.ln.weight
        b = self.ln.bias
        if (not x.is_cuda) or x.dtype != torch.float32 or \
           (w is not None and (w.dtype != torch.float32 or b is None or b.dtype != torch.float32)):
            return self.ln(x)
        if w is None:
            z = torch.empty(0, device=x.device, dtype=torch.float32)
            return _ln_ext.layernorm_forward(
                x, z, z, self.normalized_size, float(self.ln.eps))
        return _ln_ext.layernorm_forward(
            x, w, b, self.normalized_size, float(self.ln.eps))
```