import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a GEMM, Group Normalization, Minimum operation, and Bias addition.
    """
    def __init__(self, in_features, out_features, num_groups, bias_shape):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        x = self.gemm(x)
        x = self.group_norm(x)
        x = torch.min(x, dim=1, keepdim=True)[0] 
        x = x + self.bias
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
num_groups = 512
bias_shape = (1, out_features, 1, 1)

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

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

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

cpp_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/ops/linear.h>

torch::Tensor gn_min_bias_launcher(torch::Tensor y, torch::Tensor gn_w, torch::Tensor gn_b, torch::Tensor out_bias, int64_t num_groups);

torch::Tensor model_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor lin_bias, torch::Tensor gn_w, torch::Tensor gn_b, torch::Tensor out_bias, int64_t num_groups) {
    auto y = at::linear(x, weight, lin_bias);
    return gn_min_bias_launcher(y, gn_w, gn_b, out_bias, num_groups);
}
'''

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

__global__ __launch_bounds__(256) void gn_min_kernel16(const float* __restrict__ y,
                              const float* __restrict__ gamma,
                              const float* __restrict__ beta,
                              float* __restrict__ mins,
                              int C, int G) {
    int n = blockIdx.x;
    int tid = threadIdx.x;
    int lane16 = tid & 15;
    int halfwarp = tid >> 4;
    int groups_per_iter = blockDim.x >> 4;
    unsigned mask = ((tid & 16) == 0) ? 0x0000ffffu : 0xffff0000u;
    float local_min = FLT_MAX;

    for (int g = halfwarp; g < G; g += groups_per_iter) {
        int c = (g << 4) + lane16;
        float val = y[(long long)n * C + c];
        float sum = val;
        #pragma unroll
        for (int off = 8; off > 0; off >>= 1) sum += __shfl_down_sync(mask, sum, off, 16);
        float mean = __shfl_sync(mask, sum, 0, 16) * 0.0625f;
        float d = val - mean;
        float ss = d * d;
        #pragma unroll
        for (int off = 8; off > 0; off >>= 1) ss += __shfl_down_sync(mask, ss, off, 16);
        float var = __shfl_sync(mask, ss, 0, 16) * 0.0625f;
        float v = d * rsqrtf(var + 1.0e-5f);
        v = v * gamma[c] + beta[c];
        local_min = fminf(local_min, v);
    }

    __shared__ float smem[256];
    smem[tid] = local_min;
    __syncthreads();
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (tid < stride) smem[tid] = fminf(smem[tid], smem[tid + stride]);
        __syncthreads();
    }
    if (tid == 0) mins[n] = smem[0];
}

__global__ __launch_bounds__(256) void add_bias_kernel(const float* __restrict__ mins,
                                const float* __restrict__ bias,
                                float* __restrict__ out,
                                int N, int C) {
    int n = blockIdx.x * 32 + threadIdx.x;
    int c = blockIdx.y * 8 + threadIdx.y;
    if (n < N && c < C) {
        out[(long long)c * N + n] = mins[n] + bias[c];
    }
}

torch::Tensor gn_min_bias_launcher(torch::Tensor y, torch::Tensor gn_w, torch::Tensor gn_b, torch::Tensor out_bias, int64_t num_groups) {
    const int N = (int)y.size(0);
    const int C = (int)y.size(1);
    auto mins = torch::empty({N}, y.options());
    auto out = torch::empty({1, C, N, 1}, y.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    gn_min_kernel16<<<N, 256, 0, stream>>>(y.data_ptr<float>(), gn_w.data_ptr<float>(), gn_b.data_ptr<float>(), mins.data_ptr<float>(), C, (int)num_groups);
    dim3 block(32, 8);
    dim3 grid((N + 31) / 32, (C + 7) / 8);
    add_bias_kernel<<<grid, block, 0, stream>>>(mins.data_ptr<float>(), out_bias.data_ptr<float>(), out.data_ptr<float>(), N, C);
    return out;
}
'''

_ext = load_inline(name='gn_min_bias_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['model_forward'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, num_groups, bias_shape):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.num_groups = num_groups
    def forward(self, x):
        return _ext.model_forward(x, self.gemm.weight, self.gemm.bias, self.group_norm.weight, self.group_norm.bias, self.bias, self.num_groups)