import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs RMS Normalization.
    """
    def __init__(self, num_features: int, eps: float = 1e-5):
        """
        Initializes the RMSNorm layer.

        Args:
            num_features (int): Number of features in the input tensor.
            eps (float, optional): A small value added to the denominator to avoid division by zero. Defaults to 1e-5.
        """
        super(Model, self).__init__()
        self.num_features = num_features
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies RMS Normalization to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, num_features, *).

        Returns:
            torch.Tensor: Output tensor with RMS Normalization applied, same shape as input.
        """
        # Calculate the RMS along the feature dimension
        rms = torch.sqrt(torch.mean(x ** 2, dim=1, keepdim=True) + self.eps)

        # Normalize the input by dividing by the RMS
        return x / rms

batch_size = 112
features = 64
dim1 = 512
dim2 = 512

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

def get_init_inputs():
    return [features]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor rmsnorm_forward(torch::Tensor x, double eps);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdint>

__global__ void rmsnorm64_vec4_kernel(const float* __restrict__ x,
                                      float* __restrict__ y,
                                      int64_t vec_cols,
                                      int S4,
                                      int S,
                                      float eps) {
    int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;

    for (int64_t idx = tid; idx < vec_cols; idx += stride) {
        int n = (int)(idx / S4);
        int sv = (int)(idx - (int64_t)n * S4);
        int s = sv << 2;
        int64_t base = ((int64_t)n * 64LL * (int64_t)S) + (int64_t)s;

        float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f, sum3 = 0.0f;

        #pragma unroll
        for (int c = 0; c < 64; ++c) {
            const float4 v = *reinterpret_cast<const float4*>(x + base + (int64_t)c * S);
            sum0 = fmaf(v.x, v.x, sum0);
            sum1 = fmaf(v.y, v.y, sum1);
            sum2 = fmaf(v.z, v.z, sum2);
            sum3 = fmaf(v.w, v.w, sum3);
        }

        const float inv0 = rsqrtf(sum0 * 0.015625f + eps);
        const float inv1 = rsqrtf(sum1 * 0.015625f + eps);
        const float inv2 = rsqrtf(sum2 * 0.015625f + eps);
        const float inv3 = rsqrtf(sum3 * 0.015625f + eps);

        #pragma unroll
        for (int c = 0; c < 64; ++c) {
            const int64_t off = base + (int64_t)c * S;
            const float4 v = *reinterpret_cast<const float4*>(x + off);
            float4 o;
            o.x = v.x * inv0;
            o.y = v.y * inv1;
            o.z = v.z * inv2;
            o.w = v.w * inv3;
            *reinterpret_cast<float4*>(y + off) = o;
        }
    }
}

__global__ void rmsnorm_generic_vec4_kernel(const float* __restrict__ x,
                                            float* __restrict__ y,
                                            int64_t vec_cols,
                                            int C,
                                            int S4,
                                            int S,
                                            float invC,
                                            float eps) {
    int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;

    for (int64_t idx = tid; idx < vec_cols; idx += stride) {
        int n = (int)(idx / S4);
        int sv = (int)(idx - (int64_t)n * S4);
        int s = sv << 2;
        int64_t base = ((int64_t)n * (int64_t)C * (int64_t)S) + (int64_t)s;

        float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f, sum3 = 0.0f;
        for (int c = 0; c < C; ++c) {
            const float4 v = *reinterpret_cast<const float4*>(x + base + (int64_t)c * S);
            sum0 = fmaf(v.x, v.x, sum0);
            sum1 = fmaf(v.y, v.y, sum1);
            sum2 = fmaf(v.z, v.z, sum2);
            sum3 = fmaf(v.w, v.w, sum3);
        }

        const float inv0 = rsqrtf(sum0 * invC + eps);
        const float inv1 = rsqrtf(sum1 * invC + eps);
        const float inv2 = rsqrtf(sum2 * invC + eps);
        const float inv3 = rsqrtf(sum3 * invC + eps);

        for (int c = 0; c < C; ++c) {
            const int64_t off = base + (int64_t)c * S;
            const float4 v = *reinterpret_cast<const float4*>(x + off);
            float4 o;
            o.x = v.x * inv0;
            o.y = v.y * inv1;
            o.z = v.z * inv2;
            o.w = v.w * inv3;
            *reinterpret_cast<float4*>(y + off) = o;
        }
    }
}

__global__ void rmsnorm_generic_scalar_kernel(const float* __restrict__ x,
                                              float* __restrict__ y,
                                              int64_t cols,
                                              int C,
                                              int S,
                                              float invC,
                                              float eps) {
    int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;

    for (int64_t idx = tid; idx < cols; idx += stride) {
        int n = (int)(idx / S);
        int s = (int)(idx - (int64_t)n * S);
        int64_t base = ((int64_t)n * (int64_t)C * (int64_t)S) + (int64_t)s;
        float sum = 0.0f;
        for (int c = 0; c < C; ++c) {
            float v = x[base + (int64_t)c * S];
            sum = fmaf(v, v, sum);
        }
        float inv = rsqrtf(sum * invC + eps);
        for (int c = 0; c < C; ++c) {
            int64_t off = base + (int64_t)c * S;
            y[off] = x[off] * inv;
        }
    }
}

torch::Tensor rmsnorm_forward(torch::Tensor x, double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(x.dim() >= 2, "x must have at least 2 dimensions");

    auto y = torch::empty_like(x);
    const int64_t B64 = x.size(0);
    const int64_t C64 = x.size(1);
    const int64_t numel = x.numel();
    if (numel == 0) return y;
    const int64_t S64 = numel / (B64 * C64);

    TORCH_CHECK(C64 <= 2147483647LL && S64 <= 2147483647LL, "dimensions too large");
    const int C = (int)C64;
    const int S = (int)S64;

    const float* xp = x.data_ptr<float>();
    float* yp = y.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;

    if ((S & 3) == 0) {
        const int S4 = S >> 2;
        const int64_t vec_cols = B64 * (int64_t)S4;
        int blocks = (int)((vec_cols + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        if (C == 64) {
            rmsnorm64_vec4_kernel<<<blocks, threads, 0, stream>>>(xp, yp, vec_cols, S4, S, (float)eps);
        } else {
            rmsnorm_generic_vec4_kernel<<<blocks, threads, 0, stream>>>(xp, yp, vec_cols, C, S4, S, 1.0f / (float)C, (float)eps);
        }
    } else {
        const int64_t cols = B64 * S64;
        int blocks = (int)((cols + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        rmsnorm_generic_scalar_kernel<<<blocks, threads, 0, stream>>>(xp, yp, cols, C, S, 1.0f / (float)C, (float)eps);
    }
    return y;
}
'''

_ext = load_inline(
    name='rmsnorm_forward_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['rmsnorm_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, num_features: int, eps: float = 1e-5):
        super(ModelNew, self).__init__()
        self.num_features = num_features
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.rmsnorm_forward(x, self.eps)