import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a matrix multiplication, group normalization, leaky ReLU activation, and element-wise sum.
    """
    def __init__(self, input_size, hidden_size, num_groups, eps=1e-5, negative_slope=0.01):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, hidden_size)
        self.gn = nn.GroupNorm(num_groups=num_groups, num_channels=hidden_size, eps=eps)
        self.leaky_relu = nn.LeakyReLU(negative_slope=negative_slope)

    def forward(self, x):
        """
        Performs the forward pass of the model.

        Args:
            x: Input tensor of shape (batch_size, input_size).

        Returns:
            Output tensor of shape (batch_size, hidden_size).
        """
        x = self.fc(x)
        x = self.gn(x)
        x = self.leaky_relu(x)
        x = x + x
        return x


batch_size = 1024
input_size = 8192
hidden_size = 8192
num_groups = 512

def get_inputs():
    return [torch.rand(batch_size, input_size)]

def get_init_inputs():
    return [input_size, hidden_size, num_groups]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src=r'''
#include <torch/extension.h>
torch::Tensor fused_gn_lrelu_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps, double neg_slope, int64_t num_groups);
'''
_cuda_src=r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
__inline__ __device__ float half_reduce_sum(float v, unsigned mask){
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}
__inline__ __device__ float warp_reduce_sum(float v){
    v += __shfl_down_sync(0xffffffffu, v, 16);
    v += __shfl_down_sync(0xffffffffu, v, 8);
    v += __shfl_down_sync(0xffffffffu, v, 4);
    v += __shfl_down_sync(0xffffffffu, v, 2);
    v += __shfl_down_sync(0xffffffffu, v, 1);
    return v;
}
__global__ void gn16_halfwarp_inplace_kernel(float* __restrict__ x,const float* __restrict__ gamma,const float* __restrict__ beta,int N,int C,int G,float eps,float neg_slope){
    int tid=threadIdx.x; int lane=tid&31; int warp=tid>>5; int half=lane>>4; int sub=lane&15;
    int groups_per_block=(blockDim.x>>5)*2; int idx=blockIdx.x*groups_per_block + warp*2 + half; if(idx>=N*G) return;
    int n=idx/G; int g=idx-n*G; int base=n*C + g*16;
    unsigned mask = half ? 0xffff0000u : 0x0000ffffu;
    float v=x[base+sub];
    float sum=half_reduce_sum(v,mask); sum=__shfl_sync(mask,sum,half*16); float mean=sum*0.0625f;
    float d=v-mean;
    float ss=half_reduce_sum(d*d,mask); ss=__shfl_sync(mask,ss,half*16); float inv=rsqrtf(ss*0.0625f+eps);
    int c=g*16+sub; float y=d*inv*gamma[c]+beta[c]; y=(y>=0.0f)?y:y*neg_slope; x[base+sub]=y+y;
}
__global__ void gn_generic_inplace_kernel(float* __restrict__ x,const float* __restrict__ gamma,const float* __restrict__ beta,int N,int C,int G,float eps,float neg_slope){
    int tid=threadIdx.x,lane=tid&31,warp=tid>>5; int wpb=blockDim.x>>5; int idx=blockIdx.x*wpb+warp; if(idx>=N*G) return; int n=idx/G; int g=idx-n*G; int gs=C/G; int base=n*C+g*gs;
    float sum=0.f; for(int i=lane;i<gs;i+=32) sum+=x[base+i]; sum=warp_reduce_sum(sum); sum=__shfl_sync(0xffffffffu,sum,0); float mean=sum/(float)gs;
    float ss=0.f; for(int i=lane;i<gs;i+=32){float d=x[base+i]-mean; ss+=d*d;} ss=warp_reduce_sum(ss); ss=__shfl_sync(0xffffffffu,ss,0); float inv=rsqrtf(ss/(float)gs+eps);
    for(int i=lane;i<gs;i+=32){int c=g*gs+i; float y=(x[base+i]-mean)*inv*gamma[c]+beta[c]; y=(y>=0.f)?y:y*neg_slope; x[base+i]=y+y;}
}
torch::Tensor fused_gn_lrelu_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps, double neg_slope, int64_t num_groups){
    int N=x.size(0), C=x.size(1), G=num_groups; int gs=C/G; cudaStream_t stream=at::cuda::getCurrentCUDAStream(); const int threads=256;
    if(gs==16){ int groups_per_block=(threads/32)*2; int blocks=(N*G+groups_per_block-1)/groups_per_block; gn16_halfwarp_inplace_kernel<<<blocks,threads,0,stream>>>(x.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),N,C,G,(float)eps,(float)neg_slope); }
    else { int warps=threads/32; int blocks=(N*G+warps-1)/warps; gn_generic_inplace_kernel<<<blocks,threads,0,stream>>>(x.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),N,C,G,(float)eps,(float)neg_slope); }
    return x;
}
'''
_ext=load_inline(name='fused_gn_lrelu2_ext_final',cpp_sources=[_cpp_src],cuda_sources=[_cuda_src],functions=['fused_gn_lrelu_forward'],verbose=False,extra_cuda_cflags=['-O3','--use_fast_math'],extra_cflags=['-O3'])
class ModelNew(nn.Module):
    def __init__(self,input_size,hidden_size,num_groups,eps=1e-5,negative_slope=0.01):
        super(ModelNew,self).__init__(); self.fc=nn.Linear(input_size,hidden_size); self.gn=nn.GroupNorm(num_groups=num_groups,num_channels=hidden_size,eps=eps); self.negative_slope=negative_slope; self.num_groups=num_groups; self.eps=eps
    def forward(self,x):
        y=torch.nn.functional.linear(x,self.fc.weight,self.fc.bias)
        return _ext.fused_gn_lrelu_forward(y,self.gn.weight,self.gn.bias,self.eps,self.negative_slope,self.num_groups)