import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, adds a bias term, and applies ReLU.
    """
    def __init__(self, in_features, out_features, bias_shape):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=False)
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor with shape (batch_size, in_features).
        Returns:
            torch.Tensor: Output tensor with shape (batch_size, out_features).
        """
        x = self.gemm(x)
        x = x + self.bias
        x = torch.relu(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
bias_shape = (out_features,)

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

def get_init_inputs():
    return [in_features, out_features, bias_shape]

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_bias_relu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
"""

_cuda_src = r"""
#include <torch/types.h>
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>

#define CUBLAS_CHECK(stmt) do {                                      \
    cublasStatus_t _status = (stmt);                                 \
    TORCH_CHECK(_status == CUBLAS_STATUS_SUCCESS,                    \
                "cuBLAS error: ", static_cast<int>(_status));        \
} while (0)

__device__ __forceinline__ float relu_f(float x) {
    return x > 0.0f ? x : 0.0f;
}

__global__ void add_bias_relu_vec4_kernel(float4* __restrict__ out4,
                                          const float4* __restrict__ bias4,
                                          int rows, int n4) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    int r = blockIdx.y;
    if (c < n4) {
        long long idx = static_cast<long long>(r) * n4 + c;
        float4 v = out4[idx];
        float4 b = bias4[c];
        v.x = relu_f(v.x + b.x);
        v.y = relu_f(v.y + b.y);
        v.z = relu_f(v.z + b.z);
        v.w = relu_f(v.w + b.w);
        out4[idx] = v;
    }
}

__global__ void add_bias_relu_scalar_kernel(float* __restrict__ out,
                                            const float* __restrict__ bias,
                                            long long total, int n) {
    long long idx = static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    long long stride = static_cast<long long>(blockDim.x) * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = static_cast<int>(idx % n);
        float v = out[idx] + bias[c];
        out[idx] = v > 0.0f ? v : 0.0f;
    }
}

torch::Tensor linear_bias_relu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), "all tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat && bias.scalar_type() == at::kFloat,
                "only float32 tensors are supported");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(), "tensors must be contiguous");
    TORCH_CHECK(x.dim() == 2 && weight.dim() == 2, "x and weight must be 2D");

    const int M = static_cast<int>(x.size(0));
    const int K = static_cast<int>(x.size(1));
    const int N = static_cast<int>(weight.size(0));
    TORCH_CHECK(static_cast<int>(weight.size(1)) == K, "incompatible matrix dimensions");
    TORCH_CHECK(bias.numel() == N, "bias has wrong size");

    auto out = torch::empty({M, N}, x.options());

    const float alpha = 1.0f;
    const float beta = 0.0f;
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    CUBLAS_CHECK(cublasSetStream(handle, stream));

    // Row-major C = x @ weight.T.  Interpreting row-major buffers as column-major,
    // compute C^T (N x M) = weight (N x K) * x^T (K x M).
    CUBLAS_CHECK(cublasSgemm(handle,
                             CUBLAS_OP_T, CUBLAS_OP_N,
                             N, M, K,
                             &alpha,
                             weight.data_ptr<float>(), K,
                             x.data_ptr<float>(), K,
                             &beta,
                             out.data_ptr<float>(), N));

    const int threads = 256;
    if ((N % 4) == 0) {
        int n4 = N >> 2;
        dim3 block(threads);
        dim3 grid((n4 + threads - 1) / threads, M);
        add_bias_relu_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias.data_ptr<float>()),
            M, n4);
    } else {
        long long total = static_cast<long long>(M) * N;
        int blocks = static_cast<int>((total + threads - 1) / threads);
        if (blocks > 4096) blocks = 4096;
        add_bias_relu_scalar_kernel<<<blocks, threads, 0, stream>>>(out.data_ptr<float>(), bias.data_ptr<float>(), total, N);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="linear_bias_relu_cublas_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["linear_bias_relu_forward"],
    with_cuda=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Drop-in replacement for the reference module.  Parameters are kept with the
    same names/layout so external initialization or state_dict copying works.
    """
    def __init__(self, in_features, out_features, bias_shape):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=False)
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        return _ext.linear_bias_relu_forward(x, self.gemm.weight, self.bias)