import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a matrix multiplication, divides by a scalar, and applies GELU activation.
    """
    def __init__(self, input_size, output_size, divisor):
        super(Model, self).__init__()
        self.linear = nn.Linear(input_size, output_size)
        self.divisor = divisor

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, input_size).
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, output_size).
        """
        x = self.linear(x)
        x = x / self.divisor
        x = torch.nn.functional.gelu(x)
        return x

batch_size = 1024
input_size = 8192
output_size = 8192
divisor = 10.0

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

def get_init_inputs():
    return [input_size, output_size, divisor]

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

_cpp_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/Context.h>
#include <cublas_v2.h>

void gelu_bias_div_launcher(torch::Tensor y, torch::Tensor bias, double inv_divisor);

torch::Tensor linear_div_gelu(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double divisor) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
    TORCH_CHECK(bias.is_cuda(), "bias must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 is supported");
    TORCH_CHECK(weight.scalar_type() == torch::kFloat32, "only float32 is supported");
    TORCH_CHECK(bias.scalar_type() == torch::kFloat32, "only float32 is supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");
    TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
    TORCH_CHECK(x.dim() == 2 && weight.dim() == 2 && bias.dim() == 1, "bad tensor ranks");

    const int64_t M64 = x.size(0);
    const int64_t K64 = x.size(1);
    const int64_t N64 = weight.size(0);
    TORCH_CHECK(weight.size(1) == K64, "incompatible dimensions");
    TORCH_CHECK(bias.size(0) == N64, "incompatible bias");
    TORCH_CHECK(M64 <= INT_MAX && N64 <= INT_MAX && K64 <= INT_MAX, "dimensions too large");
    const int M = static_cast<int>(M64);
    const int N = static_cast<int>(N64);
    const int K = static_cast<int>(K64);

    auto y = torch::empty({M64, N64}, x.options());

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

    // Match PyTorch's TF32 setting for fp32 matmul as closely as possible.
#if defined(CUBLAS_TF32_TENSOR_OP_MATH) && defined(CUBLAS_PEDANTIC_MATH)
    cublasSetMathMode(handle, at::globalContext().allowTF32CuBLAS() ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_PEDANTIC_MATH);
#elif defined(CUBLAS_PEDANTIC_MATH)
    cublasSetMathMode(handle, CUBLAS_PEDANTIC_MATH);
#else
    cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH);
#endif

    // Row-major y[M,N] = (x[M,K] * weight[N,K]^T) / divisor.
    // Interpreting row-major buffers as column-major gives y^T[N,M] = weight * x^T.
    const float alpha = static_cast<float>(1.0 / divisor);
    const float beta = 0.0f;
    cublasStatus_t st = cublasSgemm(
        handle,
        CUBLAS_OP_T, CUBLAS_OP_N,
        N, M, K,
        &alpha,
        weight.data_ptr<float>(), K,
        x.data_ptr<float>(), K,
        &beta,
        y.data_ptr<float>(), N);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "cublasSgemm failed with status ", (int)st);

    gelu_bias_div_launcher(y, bias, 1.0 / divisor);
    return y;
}
'''

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

__global__ void gelu_bias_div_kernel(float* __restrict__ y,
                                     const float* __restrict__ bias,
                                     int M, int N, float inv_divisor) {
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    const int row = blockIdx.y;
    if (col < N && row < M) {
        const int idx = row * N + col;
        float v = y[idx] + bias[col] * inv_divisor;
        // Exact GELU: 0.5*x*(1 + erf(x/sqrt(2))).
        float r = 0.5f * v * (1.0f + erff(v * 0.70710678118654752440f));
        y[idx] = r;
    }
}

void gelu_bias_div_launcher(torch::Tensor y, torch::Tensor bias, double inv_divisor) {
    const int M = static_cast<int>(y.size(0));
    const int N = static_cast<int>(y.size(1));
    const dim3 block(256);
    const dim3 grid((N + block.x - 1) / block.x, M);
    gelu_bias_div_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), bias.data_ptr<float>(), M, N, static_cast<float>(inv_divisor));
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}
'''

_ext = load_inline(
    name='linear_div_gelu_fused_ext_v1',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['linear_div_gelu'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, input_size, output_size, divisor):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(input_size, output_size)
        self.divisor = divisor

    def forward(self, x):
        if not x.is_contiguous():
            x = x.contiguous()
        return _ext.linear_div_gelu(x, self.linear.weight, self.linear.bias, float(self.divisor))