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 post_forward(torch::Tensor y,
                           torch::Tensor lin_bias,
                           torch::Tensor gn_weight,
                           torch::Tensor gn_bias,
                           double eps,
                           double negative_slope,
                           int64_t num_groups);
"""

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

__inline__ __device__ float warp_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    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;
}

// Optimized path for the benchmark shape: group size is 16.  One warp computes
// one group, giving contiguous/coalesced accesses within each group.
__global__ void groupnorm_warp_kernel(float* __restrict__ y,
                                      const float* __restrict__ lin_bias,
                                      const float* __restrict__ gamma,
                                      const float* __restrict__ beta,
                                      int N, int C, int G, int group_size,
                                      float eps, float neg_slope) {
    int n = blockIdx.x;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int warps_per_block = blockDim.x >> 5;

    int row_base = n * C;
    float inv_m = 1.0f / (float)group_size;

    for (int g = warp; g < G; g += warps_per_block) {
        int c0 = g * group_size;
        float v = 0.0f;
        if (lane < group_size) {
            int c = c0 + lane;
            v = y[row_base + c] + lin_bias[c];
        }

        float s = warp_sum(v);
        float mean = __shfl_sync(0xffffffffu, s, 0) * inv_m;

        float d = (lane < group_size) ? (v - mean) : 0.0f;
        float ss = warp_sum(d * d);
        float var = __shfl_sync(0xffffffffu, ss, 0) * inv_m;
        float inv_std = rsqrtf(var + eps);

        if (lane < group_size) {
            int c = c0 + lane;
            float z = d * inv_std;
            z = z * gamma[c] + beta[c];
            z = (z >= 0.0f) ? z : z * neg_slope;
            y[row_base + c] = z + z;
        }
    }
}

// Generic fallback for unusual constructor arguments with large groups.
__global__ void groupnorm_serial_kernel(float* __restrict__ y,
                                        const float* __restrict__ lin_bias,
                                        const float* __restrict__ gamma,
                                        const float* __restrict__ beta,
                                        int N, int C, int G, int group_size,
                                        float eps, float neg_slope) {
    int n = blockIdx.x;
    int tid = threadIdx.x;
    int row_base = n * C;
    float inv_m = 1.0f / (float)group_size;

    for (int g = tid; g < G; g += blockDim.x) {
        int c0 = g * group_size;
        float sum = 0.0f;
        for (int i = 0; i < group_size; ++i) {
            int c = c0 + i;
            sum += y[row_base + c] + lin_bias[c];
        }
        float mean = sum * inv_m;
        float ss = 0.0f;
        for (int i = 0; i < group_size; ++i) {
            int c = c0 + i;
            float d = (y[row_base + c] + lin_bias[c]) - mean;
            ss += d * d;
        }
        float inv_std = rsqrtf(ss * inv_m + eps);
        for (int i = 0; i < group_size; ++i) {
            int c = c0 + i;
            float z = ((y[row_base + c] + lin_bias[c]) - mean) * inv_std;
            z = z * gamma[c] + beta[c];
            z = (z >= 0.0f) ? z : z * neg_slope;
            y[row_base + c] = z + z;
        }
    }
}

torch::Tensor post_forward(torch::Tensor y,
                           torch::Tensor lin_bias,
                           torch::Tensor gn_weight,
                           torch::Tensor gn_bias,
                           double eps,
                           double negative_slope,
                           int64_t num_groups) {
    TORCH_CHECK(y.is_cuda(), "y must be a CUDA tensor");
    TORCH_CHECK(y.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(y.is_contiguous(), "y must be contiguous");

    int N = (int)y.size(0);
    int C = (int)y.size(1);
    int G = (int)num_groups;
    int group_size = C / G;

    float* yp = y.data_ptr<float>();
    const float* bp = lin_bias.data_ptr<float>();
    const float* gp = gn_weight.data_ptr<float>();
    const float* betap = gn_bias.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int threads = 256;
    if (group_size <= 32) {
        groupnorm_warp_kernel<<<N, threads, 0, stream>>>(yp, bp, gp, betap, N, C, G, group_size,
                                                        (float)eps, (float)negative_slope);
    } else {
        groupnorm_serial_kernel<<<N, threads, 0, stream>>>(yp, bp, gp, betap, N, C, G, group_size,
                                                          (float)eps, (float)negative_slope);
    }
    return y;
}
"""

_ext = load_inline(
    name="modelnew_linear_gn_lrelu_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["post_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Matrix multiplication followed by fused GroupNorm + LeakyReLU + elementwise sum.
    """
    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.leaky_relu = nn.LeakyReLU(negative_slope=negative_slope)

    def forward(self, x):
        # GEMM only; the linear bias is consumed inside the fused CUDA post-kernel.
        y = torch.mm(x, self.fc.weight.t())
        return _ext.post_forward(y, self.fc.bias, self.gn.weight, self.gn.bias,
                                 float(self.gn.eps), float(self.leaky_relu.negative_slope),
                                 int(self.gn.num_groups))