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_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta, int64_t groups, double eps, double negative_slope);
'''

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

__device__ __forceinline__ float finish_val(float x, float mean, float invstd, float gamma, float beta, float neg) {
    float v = (x - mean) * invstd;
    v = v * gamma + beta;
    v = (v >= 0.0f) ? v : v * neg;
    return v + v;
}

__global__ void gn16_act2_kernel(float* __restrict__ y,
                                 const float* __restrict__ gamma,
                                 const float* __restrict__ beta,
                                 int total_groups, float eps, float negative_slope) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    if (gid >= total_groups) return;
    int base = gid << 4;
    int wbase = (gid & 511) << 4;

    const float4* py = reinterpret_cast<const float4*>(y + base);
    float4 a0 = py[0];
    float4 a1 = py[1];
    float4 a2 = py[2];
    float4 a3 = py[3];

    float sum = a0.x+a0.y+a0.z+a0.w + a1.x+a1.y+a1.z+a1.w + a2.x+a2.y+a2.z+a2.w + a3.x+a3.y+a3.z+a3.w;
    float mean = sum * 0.0625f;
    float d0=a0.x-mean, d1=a0.y-mean, d2=a0.z-mean, d3=a0.w-mean;
    float d4=a1.x-mean, d5=a1.y-mean, d6=a1.z-mean, d7=a1.w-mean;
    float d8=a2.x-mean, d9=a2.y-mean, d10=a2.z-mean, d11=a2.w-mean;
    float d12=a3.x-mean, d13=a3.y-mean, d14=a3.z-mean, d15=a3.w-mean;
    float ss = d0*d0+d1*d1+d2*d2+d3*d3 + d4*d4+d5*d5+d6*d6+d7*d7 + d8*d8+d9*d9+d10*d10+d11*d11 + d12*d12+d13*d13+d14*d14+d15*d15;
    float invstd = rsqrtf(ss * 0.0625f + eps);

    const float4* pg = reinterpret_cast<const float4*>(gamma + wbase);
    const float4* pb = reinterpret_cast<const float4*>(beta + wbase);
    float4 g0=pg[0], g1=pg[1], g2=pg[2], g3=pg[3];
    float4 b0=pb[0], b1=pb[1], b2=pb[2], b3=pb[3];

    float4* po = reinterpret_cast<float4*>(y + base);
    po[0] = make_float4(finish_val(a0.x, mean, invstd, g0.x, b0.x, negative_slope),
                        finish_val(a0.y, mean, invstd, g0.y, b0.y, negative_slope),
                        finish_val(a0.z, mean, invstd, g0.z, b0.z, negative_slope),
                        finish_val(a0.w, mean, invstd, g0.w, b0.w, negative_slope));
    po[1] = make_float4(finish_val(a1.x, mean, invstd, g1.x, b1.x, negative_slope),
                        finish_val(a1.y, mean, invstd, g1.y, b1.y, negative_slope),
                        finish_val(a1.z, mean, invstd, g1.z, b1.z, negative_slope),
                        finish_val(a1.w, mean, invstd, g1.w, b1.w, negative_slope));
    po[2] = make_float4(finish_val(a2.x, mean, invstd, g2.x, b2.x, negative_slope),
                        finish_val(a2.y, mean, invstd, g2.y, b2.y, negative_slope),
                        finish_val(a2.z, mean, invstd, g2.z, b2.z, negative_slope),
                        finish_val(a2.w, mean, invstd, g2.w, b2.w, negative_slope));
    po[3] = make_float4(finish_val(a3.x, mean, invstd, g3.x, b3.x, negative_slope),
                        finish_val(a3.y, mean, invstd, g3.y, b3.y, negative_slope),
                        finish_val(a3.z, mean, invstd, g3.z, b3.z, negative_slope),
                        finish_val(a3.w, mean, invstd, g3.w, b3.w, negative_slope));
}

torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta, int64_t groups, double eps, double negative_slope) {
    auto y = at::linear(input, weight, bias);
    int64_t total_groups = y.size(0) * groups;
    const int threads = 256;
    int blocks = (int)((total_groups + threads - 1) / threads);
    gn16_act2_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
        (int)total_groups, (float)eps, (float)negative_slope);
    return y;
}
'''

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

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):
        return _ext.fused_forward(x, self.fc.weight, self.fc.bias, self.gn.weight, self.gn.bias, self.num_groups, self.eps, self.negative_slope)