import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs L2 normalization.
    """
    def __init__(self):
        """
        Initializes the L2Norm layer.

        Args:
            dim (int): Dimension along which to normalize.
        """
        super(Model, self).__init__()

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

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

        Returns:
            torch.Tensor: Output tensor with L2 normalization applied, same shape as input.
        """
        return x / torch.norm(x, p=2, dim=1, keepdim=True)

batch_size = 32768
# choose dim so total <2^31
dim = 65535

def get_inputs():
    x = torch.rand(batch_size, dim)
    return [x]

def get_init_inputs():
    return []

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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor l2norm_forward(torch::Tensor x);
"""

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

__inline__ __device__ float warp_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(0xffffffff, v, offset);
    return v;
}

template<int BLOCK>
__global__ void l2norm_inplace_kernel(float* __restrict__ x, long long rows, long long cols) {
    int row = (int)blockIdx.x;
    float* __restrict__ xp = x + ((long long)row) * cols;
    int tid = threadIdx.x;
    float sum = 0.0f;
    long long c = tid;
    const long long step = (long long)BLOCK * 4;
    for (; c + (long long)3 * BLOCK < cols; c += step) {
        float v0 = xp[c];
        float v1 = xp[c + BLOCK];
        float v2 = xp[c + (long long)2 * BLOCK];
        float v3 = xp[c + (long long)3 * BLOCK];
        sum = fmaf(v0, v0, sum);
        sum = fmaf(v1, v1, sum);
        sum = fmaf(v2, v2, sum);
        sum = fmaf(v3, v3, sum);
    }
    for (; c < cols; c += BLOCK) {
        float v = xp[c];
        sum = fmaf(v, v, sum);
    }
    __shared__ float smem[32];
    int lane = tid & 31;
    int wid = tid >> 5;
    sum = warp_sum(sum);
    if (lane == 0) smem[wid] = sum;
    __syncthreads();
    float total = (tid < (BLOCK >> 5)) ? smem[lane] : 0.0f;
    if (wid == 0) total = warp_sum(total);
    if (tid == 0) smem[0] = 1.0f / sqrtf(total);
    __syncthreads();
    float inv = smem[0];
    c = tid;
    for (; c + (long long)3 * BLOCK < cols; c += step) {
        xp[c] *= inv;
        xp[c + BLOCK] *= inv;
        xp[c + (long long)2 * BLOCK] *= inv;
        xp[c + (long long)3 * BLOCK] *= inv;
    }
    for (; c < cols; c += BLOCK) xp[c] *= inv;
}

torch::Tensor l2norm_forward(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(x.dim() == 2, "x must be 2D");
    const auto rows = x.size(0);
    const auto cols = x.size(1);
    if (rows == 0 || cols == 0) return x;
    c10::cuda::CUDAGuard device_guard(x.device());
    const int BLOCK = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    l2norm_inplace_kernel<BLOCK><<<dim3(rows), dim3(BLOCK), 0, stream>>>(x.data_ptr<float>(), (long long)rows, (long long)cols);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return x;
}
"""

_ext = load_inline(name="l2norm_ext_inplace_unroll_v1", cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=["l2norm_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.l2norm_forward(x)