import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, applies ReLU, and divides by a constant.
    """
    def __init__(self, in_features, out_features, divisor):
        super(Model, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.divisor = divisor

    def forward(self, x):
        x = self.linear(x)
        x = torch.relu(x)
        x = x / self.divisor
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
divisor = 2.0

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

def get_init_inputs():
    return [in_features, out_features, 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 <c10/cuda/CUDAGuard.h>
#include <cublas_v2.h>

void relu_div_bias_launcher(torch::Tensor y, torch::Tensor bias, double divisor);

static inline void check_cublas(cublasStatus_t status) {
    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, "cuBLAS call failed with status ", (int)status);
}

torch::Tensor linear_relu_div(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, "x must be float32");
    TORCH_CHECK(weight.scalar_type() == torch::kFloat32, "weight must be float32");
    TORCH_CHECK(bias.scalar_type() == torch::kFloat32, "bias must be float32");
    TORCH_CHECK(x.dim() == 2 && weight.dim() == 2 && bias.dim() == 1, "bad tensor dimensions");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(), "tensors must be contiguous");

    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 && bias.size(0) == N64, "incompatible shapes");
    TORCH_CHECK(M64 <= INT_MAX && N64 <= INT_MAX && K64 <= INT_MAX, "problem too large");
    const int M = (int)M64;
    const int N = (int)N64;
    const int K = (int)K64;

    c10::cuda::CUDAGuard device_guard(x.device());
    auto y = torch::empty({M64, N64}, x.options());

    // Row-major y = x * weight^T.  cuBLAS is column-major; viewing row-major
    // buffers as transposed column-major matrices gives C^T = weight * x^T.
    const float alpha = 1.0f;
    const float beta = 0.0f;
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    check_cublas(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));

    relu_div_bias_launcher(y, bias, divisor);
    return y;
}
'''

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

__global__ void relu_div_bias_kernel(float* __restrict__ y,
                                     const float* __restrict__ bias,
                                     int M, int N, float inv_div) {
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    if (row < M && col < N) {
        const int idx = row * N + col;
        float v = y[idx] + bias[col];
        v = v > 0.0f ? v * inv_div : 0.0f;
        y[idx] = v;
    }
}

void relu_div_bias_launcher(torch::Tensor y, torch::Tensor bias, double divisor) {
    const int M = (int)y.size(0);
    const int N = (int)y.size(1);
    const float inv_div = 1.0f / (float)divisor;
    dim3 block(32, 8);
    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
    relu_div_bias_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), bias.data_ptr<float>(), M, N, inv_div);
}
'''

_ext = load_inline(
    name='linear_relu_div_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['linear_relu_div'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    verbose=False,
)

class ModelNew(nn.Module):
    """
    Fused replacement for Linear -> ReLU -> divide.  Parameters and public
    interface match the reference Model.
    """
    def __init__(self, in_features, out_features, divisor):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.divisor = divisor

    def forward(self, x):
        return _ext.linear_relu_div(x, self.linear.weight, self.linear.bias, float(self.divisor))