import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a GEMM, applies Group Normalization, and then HardTanh.
    """
    def __init__(self, in_features, out_features, num_groups, hardtanh_min, hardtanh_max):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.hardtanh = nn.Hardtanh(min_val=hardtanh_min, max_val=hardtanh_max)

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_features).
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_features).
        """
        x = self.gemm(x)
        x = self.group_norm(x)
        x = self.hardtanh(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
num_groups = 16
hardtanh_min = -2.0
hardtanh_max = 2.0

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

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

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 x, torch::Tensor weight, torch::Tensor lin_bias,
                            torch::Tensor gn_weight, torch::Tensor gn_bias,
                            int64_t groups, double ht_min, double ht_max);
'''

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

__inline__ __device__ float warp_sum(float v) {
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
    return v;
}

__inline__ __device__ float block_sum(float v) {
    static __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_sum(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    v = (threadIdx.x < ((blockDim.x + 31) >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) v = warp_sum(v);
    return v;
}

__global__ void gn_ht_kernel_512(const float* __restrict__ inp,
                                 const float* __restrict__ gamma,
                                 const float* __restrict__ beta,
                                 float* __restrict__ out,
                                 int C, int groups,
                                 float ht_min, float ht_max) {
    int bg = blockIdx.x;
    int b = bg / groups;
    int g = bg - b * groups;
    int start = b * C + g * 512;
    int cstart = g * 512;

    int t = threadIdx.x;
    float v0 = inp[start + t];
    float v1 = inp[start + t + 256];
    float sum = v0 + v1;
    float sumsq = v0 * v0 + v1 * v1;

    float s = block_sum(sum);
    float ss = block_sum(sumsq);

    __shared__ float sh_mean, sh_inv;
    if (t == 0) {
        float mean = s * (1.0f / 512.0f);
        float var = ss * (1.0f / 512.0f) - mean * mean;
        var = var > 0.0f ? var : 0.0f;
        sh_mean = mean;
        sh_inv = rsqrtf(var + 1.0e-5f);
    }
    __syncthreads();

    float mean = sh_mean;
    float inv = sh_inv;
    int c = cstart + t;
    float y = (v0 - mean) * inv;
    y = y * gamma[c] + beta[c];
    y = fminf(fmaxf(y, ht_min), ht_max);
    out[start + t] = y;

    c += 256;
    y = (v1 - mean) * inv;
    y = y * gamma[c] + beta[c];
    y = fminf(fmaxf(y, ht_min), ht_max);
    out[start + t + 256] = y;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor lin_bias,
                            torch::Tensor gn_weight, torch::Tensor gn_bias,
                            int64_t groups, double ht_min, double ht_max) {
    auto y = at::addmm(lin_bias, x, weight.t());
    auto out = torch::empty_like(y);
    int B = (int)y.size(0);
    int C = (int)y.size(1);
    gn_ht_kernel_512<<<B * (int)groups, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), gn_weight.data_ptr<float>(), gn_bias.data_ptr<float>(),
        out.data_ptr<float>(), C, (int)groups, (float)ht_min, (float)ht_max);
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, num_groups, hardtanh_min, hardtanh_max):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.num_groups = int(num_groups)
        self.hardtanh_min = float(hardtanh_min)
        self.hardtanh_max = float(hardtanh_max)

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