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

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

torch::Tensor gn_ht_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, int64_t num_groups, double ht_min, double ht_max, double eps);
"""

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

__global__ void gn_ht_kernel(const float* __restrict__ x,
                             const float* __restrict__ gamma,
                             const float* __restrict__ beta,
                             float* __restrict__ y,
                             int rows, int cols, int group_size,
                             float ht_min, float ht_max, float eps) {
    int bg = blockIdx.x;
    int row = bg / gridDim.y;
    int g = bg - row * gridDim.y;
    int base = row * cols + g * group_size;
    int tid = threadIdx.x;

    float sum = 0.0f;
    float sumsq = 0.0f;
    // Designed for group_size=512, blockDim=256: two elements per thread.
    for (int i = tid; i < group_size; i += blockDim.x) {
        float v = x[base + i];
        sum += v;
        sumsq += v * v;
    }

    // warp reductions
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffff, sum, offset);
        sumsq += __shfl_down_sync(0xffffffff, sumsq, offset);
    }

    __shared__ float sh_sum[8];
    __shared__ float sh_sumsq[8];
    int lane = tid & 31;
    int warp = tid >> 5;
    if (lane == 0) {
        sh_sum[warp] = sum;
        sh_sumsq[warp] = sumsq;
    }
    __syncthreads();

    if (warp == 0) {
        sum = (tid < (blockDim.x >> 5)) ? sh_sum[lane] : 0.0f;
        sumsq = (tid < (blockDim.x >> 5)) ? sh_sumsq[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum += __shfl_down_sync(0xffffffff, sum, offset);
            sumsq += __shfl_down_sync(0xffffffff, sumsq, offset);
        }
        if (tid == 0) {
            float mean = sum / (float)group_size;
            float var = sumsq / (float)group_size - mean * mean;
            if (var < 0.0f) var = 0.0f;
            sh_sum[0] = mean;
            sh_sumsq[0] = rsqrtf(var + eps);
        }
    }
    __syncthreads();

    float mean = sh_sum[0];
    float invstd = sh_sumsq[0];
    for (int i = tid; i < group_size; i += blockDim.x) {
        int c = g * group_size + i;
        float v = (x[base + i] - mean) * invstd;
        v = v * gamma[c] + beta[c];
        v = fminf(fmaxf(v, ht_min), ht_max);
        y[base + i] = v;
    }
}

torch::Tensor gn_ht_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, int64_t num_groups, double ht_min, double ht_max, double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    auto y = torch::empty_like(x);
    int rows = (int)x.size(0);
    int cols = (int)x.size(1);
    int groups = (int)num_groups;
    int group_size = cols / groups;
    dim3 block(256);
    dim3 grid(rows * groups);
    // recover row/group from 1D grid for maximum compatibility
    // kernel uses gridDim.y, so launch equivalent 2D when supported
    dim3 grid2(groups, rows);
    gn_ht_kernel<<<grid2, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), y.data_ptr<float>(),
        rows, cols, group_size, (float)ht_min, (float)ht_max, (float)eps);
    return y;
}
"""

_ext = load_inline(name="gn_ht_ext_v1",
                   cpp_sources=_cpp_src,
                   cuda_sources=_cuda_src,
                   functions=["gn_ht_cuda"],
                   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, 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.hardtanh = nn.Hardtanh(min_val=hardtanh_min, max_val=hardtanh_max)
        self.num_groups = num_groups
        self.hardtanh_min = float(hardtanh_min)
        self.hardtanh_max = float(hardtanh_max)

    def forward(self, x):
        z = F.linear(x, self.gemm.weight, self.gemm.bias)
        return _ext.gn_ht_cuda(z.contiguous(), self.group_norm.weight, self.group_norm.bias,
                               self.num_groups, self.hardtanh_min, self.hardtanh_max, self.group_norm.eps)