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 os
os.makedirs('/scratch/abaner/tmp', exist_ok=True)
os.environ['TMPDIR'] = '/scratch/abaner/tmp'
os.environ['MAX_JOBS'] = '1'

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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor linear_relu_div_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor scaled_bias, double divisor);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <mutex>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

static cublasLtHandle_t ltHandle = nullptr;
static std::once_flag handle_once;

static void check_cublas(cublasStatus_t s, const char* m) {
    TORCH_CHECK(s == CUBLAS_STATUS_SUCCESS, m, " cublas status ", (int)s);
}

__global__ void scale_bias_kernel(const float* __restrict__ b, float* __restrict__ sb, int n, float inv) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) sb[i] = b[i] * inv;
}

torch::Tensor linear_relu_div_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor scaled_bias, double divisor) {
    CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(bias); CHECK_CUDA(scaled_bias);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight); CHECK_CONTIGUOUS(bias); CHECK_CONTIGUOUS(scaled_bias);
    CHECK_FLOAT(x); CHECK_FLOAT(weight); CHECK_FLOAT(bias); CHECK_FLOAT(scaled_bias);

    int64_t M64 = x.size(0), K64 = x.size(1), N64 = weight.size(0);
    TORCH_CHECK(x.dim() == 2 && weight.dim() == 2 && bias.dim() == 1, "bad dims");
    TORCH_CHECK(weight.size(1) == K64 && bias.size(0) == N64 && scaled_bias.numel() >= N64, "shape mismatch");
    int M = (int)M64, N = (int)N64, K = (int)K64;

    auto out = torch::empty({M64, N64}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    float inv = 1.0f / (float)divisor;

    scale_bias_kernel<<<(N + 255) / 256, 256, 0, stream>>>(bias.data_ptr<float>(), scaled_bias.data_ptr<float>(), N, inv);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    std::call_once(handle_once, [](){ check_cublas(cublasLtCreate(&ltHandle), "cublasLtCreate failed"); });

    cublasLtMatmulDesc_t opDesc;
    cublasLtMatrixLayout_t aDesc, bDesc, cDesc;
    check_cublas(cublasLtMatmulDescCreate(&opDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F), "op desc create failed");

    // Treat row-major output [M,N] as a column-major [N,M] matrix:
    // out^T = weight [N,K] * x^T [K,M].  The row-major storage of weight is
    // a column-major [K,N] view, so use transa=T.
    cublasOperation_t transa = CUBLAS_OP_T;
    cublasOperation_t transb = CUBLAS_OP_N;
    check_cublas(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)), "set transa failed");
    check_cublas(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb)), "set transb failed");

    cublasLtEpilogue_t epi = CUBLASLT_EPILOGUE_RELU_BIAS;
    check_cublas(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epi, sizeof(epi)), "set epilogue failed");
    void* bias_ptr = (void*)scaled_bias.data_ptr<float>();
    check_cublas(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)), "set bias failed");

    check_cublas(cublasLtMatrixLayoutCreate(&aDesc, CUDA_R_32F, K, N, K), "a desc create failed");
    check_cublas(cublasLtMatrixLayoutCreate(&bDesc, CUDA_R_32F, K, M, K), "b desc create failed");
    check_cublas(cublasLtMatrixLayoutCreate(&cDesc, CUDA_R_32F, N, M, N), "c desc create failed");

    cublasLtMatmulPreference_t pref;
    check_cublas(cublasLtMatmulPreferenceCreate(&pref), "pref create failed");
    size_t maxWorkspace = 0;
    check_cublas(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &maxWorkspace, sizeof(maxWorkspace)), "pref set failed");
    cublasLtMatmulHeuristicResult_t heuristic;
    int returned = 0;
    check_cublas(cublasLtMatmulAlgoGetHeuristic(ltHandle, opDesc, aDesc, bDesc, cDesc, cDesc, pref, 1, &heuristic, &returned), "heuristic failed");
    TORCH_CHECK(returned > 0, "no cublasLt algo found");

    float alpha = inv;
    float beta = 0.0f;
    check_cublas(cublasLtMatmul(ltHandle, opDesc, &alpha,
                                 weight.data_ptr<float>(), aDesc,
                                 x.data_ptr<float>(), bDesc,
                                 &beta,
                                 out.data_ptr<float>(), cDesc,
                                 out.data_ptr<float>(), cDesc,
                                 &heuristic.algo, nullptr, 0, stream), "matmul failed");

    cublasLtMatmulPreferenceDestroy(pref);
    cublasLtMatrixLayoutDestroy(cDesc);
    cublasLtMatrixLayoutDestroy(bDesc);
    cublasLtMatrixLayoutDestroy(aDesc);
    cublasLtMatmulDescDestroy(opDesc);
    return out;
}
'''

_ext = load_inline(
    name='linear_relu_div_lt_ext_final',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['linear_relu_div_cuda'],
    extra_ldflags=['-lcublasLt', '-lcublas'],
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, divisor):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.divisor = divisor
        self.register_buffer('scaled_bias', torch.empty(out_features, dtype=torch.float32), persistent=False)

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