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 <cuda.h>
#include <cuda_runtime.h>

__inline__ __device__ float warp_reduce_sum(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}

template<int BLOCK>
__global__ void l2norm_rows_kernel(const float* __restrict__ x,
                                   float* __restrict__ y,
                                   int rows, int cols) {
    int row = blockIdx.x;
    if (row >= rows) return;
    int tid = threadIdx.x;
    long long base = (long long)row * (long long)cols;

    float ss = 0.0f;
    for (int c = tid; c < cols; c += BLOCK) {
        float v = x[base + c];
        ss = fmaf(v, v, ss);
    }

    ss = warp_reduce_sum(ss);

    __shared__ float smem[32];
    if ((tid & 31) == 0) smem[tid >> 5] = ss;
    __syncthreads();

    float total = 0.0f;
    if (tid < 32) {
        total = (tid < (BLOCK >> 5)) ? smem[tid] : 0.0f;
        total = warp_reduce_sum(total);
        if (tid == 0) smem[0] = rsqrtf(total);
    }
    __syncthreads();

    float inv = smem[0];
    for (int c = tid; c < cols; c += BLOCK) {
        y[base + c] = x[base + 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, "only float32 supported");
    TORCH_CHECK(x.dim() == 2, "expected a 2D tensor");
    auto xc = x.contiguous();
    auto y = torch::empty_like(xc);
    int rows = (int)xc.size(0);
    int cols = (int)xc.size(1);
    constexpr int BLOCK = 256;
    dim3 grid(rows);
    dim3 block(BLOCK);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    l2norm_rows_kernel<BLOCK><<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), y.data_ptr<float>(), rows, cols);
    return y;
}
'''

_ext = load_inline(
    name='l2norm_dim1_ext_v1',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['l2norm_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    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)