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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

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

torch::Tensor post_forward_cuda(torch::Tensor x,
                                torch::Tensor gamma,
                                torch::Tensor beta,
                                torch::Tensor mul);

torch::Tensor post_forward(torch::Tensor x,
                           torch::Tensor gamma,
                           torch::Tensor beta,
                           torch::Tensor mul) {
    return post_forward_cuda(x, gamma, beta, mul);
}
'''

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

static __forceinline__ __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;
}

__global__ void post_gn_swish_kernel(float* __restrict__ x,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     const float* __restrict__ mul,
                                     int C, int G) {
    // Specialized for the benchmark shape: group size = C / G = 32.
    const int n = blockIdx.x;
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int nwarps = blockDim.x >> 5;
    const int row_base = n * C;

    constexpr float inv32 = 0.03125f;
    constexpr float eps = 1.0e-5f;

    for (int g = warp; g < G; g += nwarps) {
        const int c0 = g * 32;
        const int idx = row_base + c0 + lane;

        float v = x[idx];
        float s = warp_reduce_sum(v);
        float mean = __shfl_sync(0xffffffff, s, 0) * inv32;

        float d = v - mean;
        float ss = warp_reduce_sum(d * d);
        float var = __shfl_sync(0xffffffff, ss, 0) * inv32;

        const int c = c0 + lane;
        float y = d * rsqrtf(var + eps);
        y = y * gamma[c] + beta[c];

        // swish, multiply by per-channel weight, swish again
        float a = y / (1.0f + expf(-y));
        float z = a * mul[c];
        float out = z / (1.0f + expf(-z));
        x[idx] = out;
    }
}

torch::Tensor post_forward_cuda(torch::Tensor x,
                                torch::Tensor gamma,
                                torch::Tensor beta,
                                torch::Tensor mul) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int G = C / 32;

    const dim3 block(256);
    const dim3 grid(N);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    post_gn_swish_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(),
        gamma.data_ptr<float>(),
        beta.data_ptr<float>(),
        mul.data_ptr<float>(),
        C, G);
    return x;
}
'''

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


class ModelNew(nn.Module):
    """
    Linear followed by fused GroupNorm + Swish + channel multiply + Swish.
    """
    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):
        y = F.linear(x, self.gemm.weight, self.gemm.bias)
        return _ext.post_forward(y, self.group_norm.weight, self.group_norm.bias, self.multiply_weight)