import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a GEMM, scaling, hardtanh, and GELU activation.
    """
    def __init__(self, in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.scaling_factor = scaling_factor
        self.hardtanh = nn.Hardtanh(min_val=hardtanh_min, max_val=hardtanh_max)
        self.gelu = nn.GELU()

    def forward(self, x):
        x = self.gemm(x)
        x = x * self.scaling_factor
        x = self.hardtanh(x)
        x = self.gelu(x)
        return x

batch_size = 2048
in_features = 8192
out_features = 8192
scaling_factor = 0.5
hardtanh_min = -2
hardtanh_max = 2

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

def get_init_inputs():
    return [in_features, out_features, scaling_factor, 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>

// CUDA launcher defined in cuda source
void fused_act_launcher(torch::Tensor y, double scaling_factor, double hardtanh_min, double hardtanh_max);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                            double scaling_factor, double hardtanh_min, double hardtanh_max) {
    // Compute the GEMM using the same highly-optimized backend as torch.nn.Linear:
    // y = x @ weight.T + bias.  The following fused CUDA kernel performs all
    // pointwise post-processing in a single pass, eliminating three large
    // intermediate reads/writes and kernel launches from the reference.
    auto y = torch::addmm(bias, x, weight.t());
    fused_act_launcher(y, scaling_factor, hardtanh_min, hardtanh_max);
    return y;
}
'''

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

__global__ void fused_act_kernel(float* __restrict__ y, int64_t n,
                                 float scale, float hmin, float hmax) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    const float inv_sqrt2 = 0.70710678118654752440f;
    for (int64_t i = idx; i < n; i += stride) {
        float v = y[i] * scale;
        v = fminf(fmaxf(v, hmin), hmax);
        // Exact GELU, matching torch.nn.GELU() default closely in fp32.
        float r = 0.5f * v * (1.0f + erff(v * inv_sqrt2));
        y[i] = r;
    }
}

void fused_act_launcher(torch::Tensor y, double scaling_factor, double hardtanh_min, double hardtanh_max) {
    const int threads = 256;
    int64_t n = y.numel();
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_act_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), n,
        (float)scaling_factor, (float)hardtanh_min, (float)hardtanh_max);
}
'''

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

class ModelNew(nn.Module):
    """
    Drop-in replacement for the reference model.  Keeps the same public module
    structure/parameters (self.gemm) so benchmark harnesses can initialize or
    copy weights exactly as with the original model.
    """
    def __init__(self, in_features, out_features, scaling_factor, hardtanh_min, hardtanh_max):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.scaling_factor = float(scaling_factor)
        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.scaling_factor, self.hardtanh_min, self.hardtanh_max)