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>
torch::Tensor fused_linear_act(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale, double minv, double maxv);
'''

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

__global__ void act2d(float* __restrict__ y, const float* __restrict__ bias,
                      int M, int N, float scale, float minv, float maxv) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y;
    if (col < N) {
        long idx = (long)row * N + col;
        const float inv_sqrt2 = 0.70710678118654752440f;
        float v = (y[idx] + bias[col]) * scale;
        v = fminf(fmaxf(v, minv), maxv);
        y[idx] = 0.5f * v * (1.0f + erff(v * inv_sqrt2));
    }
}

torch::Tensor fused_linear_act(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                               double scale, double minv, double maxv) {
    int64_t M64 = x.size(0), K64 = x.size(1), N64 = weight.size(0);
    int M = (int)M64, N = (int)N64, K = (int)K64;
    auto y = torch::empty({M64, N64}, x.options());

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasSetStream(handle, stream);

    float alpha = 1.0f, beta = 0.0f;
    // Treat row-major tensors as column-major transposes: y^T = weight * x^T.
    cublasStatus_t stat = cublasGemmEx(handle,
        CUBLAS_OP_T, CUBLAS_OP_N,
        N, M, K,
        &alpha,
        weight.data_ptr<float>(), CUDA_R_32F, K,
        x.data_ptr<float>(), CUDA_R_32F, K,
        &beta,
        y.data_ptr<float>(), CUDA_R_32F, N,
        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
    TORCH_CHECK(stat == CUBLAS_STATUS_SUCCESS, "cublasGemmEx failed");

    dim3 block(256);
    dim3 grid((N + 255) / 256, M);
    act2d<<<grid, block, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(),
                                      M, N, (float)scale, (float)minv, (float)maxv);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

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

class ModelNew(nn.Module):
    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 = scaling_factor
        self.hardtanh_min = hardtanh_min
        self.hardtanh_max = hardtanh_max

    def forward(self, x):
        return _ext.fused_linear_act(
            x, self.gemm.weight, self.gemm.bias,
            float(self.scaling_factor), float(self.hardtanh_min), float(self.hardtanh_max)
        )