import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a GEMM, GroupNorm, Swish, Multiply, and Swish operations.
    """
    def __init__(self, in_features, out_features, num_groups, multiply_weight_shape):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.multiply_weight = nn.Parameter(torch.randn(multiply_weight_shape)) 

    def forward(self, x):
        # (batch_size, in_features) -> (batch_size, out_features)
        x = self.gemm(x)
        # (batch_size, out_features) -> (batch_size, out_features)
        x = self.group_norm(x)
        # (batch_size, out_features) -> (batch_size, out_features)
        x = x * torch.sigmoid(x)
        # (batch_size, out_features) -> (batch_size, out_features)
        x = x * self.multiply_weight
        # (batch_size, out_features) -> (batch_size, out_features)
        x = x * torch.sigmoid(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
num_groups = 256
multiply_weight_shape = (out_features,)

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

def get_init_inputs():
    return [in_features, out_features, num_groups, multiply_weight_shape]

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

cpp_source = r'''
#include <torch/extension.h>

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                 torch::Tensor gn_weight, torch::Tensor gn_bias,
                                 torch::Tensor mult_weight, int64_t num_groups, double eps);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                            torch::Tensor gn_weight, torch::Tensor gn_bias,
                            torch::Tensor mult_weight, int64_t num_groups, double eps) {
    return fused_forward_cuda(x, weight, bias, gn_weight, gn_bias, mult_weight, num_groups, eps);
}
'''

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

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

template<int WPB>
__global__ void gn_swish_8192x256_kernel(float* __restrict__ y,
                                         const float* __restrict__ gamma,
                                         const float* __restrict__ beta,
                                         const float* __restrict__ mw,
                                         int total_units, float eps) {
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int unit = blockIdx.x * WPB + warp;
    if (unit >= total_units) return;

    int idx = (unit << 5) + lane;        // row-major output, group size is 32
    int c = ((unit & 255) << 5) + lane;  // channel within the row (256 groups)

    float v = y[idx];
    float mean = warp_sum_bcast(v) * 0.03125f;
    float d = v - mean;
    float invstd = rsqrtf(warp_sum_bcast(d * d) * 0.03125f + eps);

    float z = d * invstd;
    z = z * gamma[c] + beta[c];
    z = z / (1.0f + expf(-z));
    z *= mw[c];
    z = z / (1.0f + expf(-z));
    y[idx] = z;
}

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                 torch::Tensor gn_weight, torch::Tensor gn_bias,
                                 torch::Tensor mult_weight, int64_t num_groups, double eps) {
    auto y = at::linear(x, weight, bias);

    int N = (int)y.size(0);
    int total = N * (int)num_groups;
    const int WPB = 8;
    int blocks = (total + WPB - 1) / WPB;
    gn_swish_8192x256_kernel<WPB><<<blocks, WPB * 32, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), gn_weight.data_ptr<float>(), gn_bias.data_ptr<float>(),
        mult_weight.data_ptr<float>(), total, (float)eps);
    return y;
}
'''

_ext = load_inline(name='fused_gemm_gn_swish_ext_submit',
                   cpp_sources=cpp_source,
                   cuda_sources=cuda_source,
                   functions=['fused_forward'],
                   extra_cflags=['-O3'],
                   extra_cuda_cflags=['-O3', '--use_fast_math'],
                   verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, num_groups, multiply_weight_shape):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.multiply_weight = nn.Parameter(torch.randn(multiply_weight_shape))
        self.num_groups = num_groups

    def forward(self, x):
        return _ext.fused_forward(x, self.gemm.weight, self.gemm.bias,
                                  self.group_norm.weight, self.group_norm.bias,
                                  self.multiply_weight, self.num_groups, self.group_norm.eps)