import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a single matrix multiplication (C = A * B)
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        """
        Performs matrix multiplication.

        Args:
            A: Input tensor of shape (M, K).
            B: Input tensor of shape (K, N).

        Returns:
            Output tensor of shape (M, N).
        """
        return torch.matmul(A, B)

M = 1024 * 2
K = 4096 * 2
N = 2048 * 2

def get_inputs():
    A = torch.rand(M, K)
    B = torch.rand(K, N)
    return [A, B]

def get_init_inputs():
    return []  # No special initialization inputs needed

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>

static cublasHandle_t get_handle() {
    static cublasHandle_t handle = nullptr;
    if (handle == nullptr) {
        cublasCreate(&handle);
        cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH);
    }
    cublasSetStream(handle, at::cuda::getCurrentCUDAStream());
    return handle;
}

torch::Tensor matmul_forward(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 && B.scalar_type() == torch::kFloat32, "only fp32 supported");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "inputs must be 2D");
    TORCH_CHECK(A.size(1) == B.size(0), "shape mismatch");

    const auto M64 = A.size(0);
    const auto K64 = A.size(1);
    const auto N64 = B.size(1);
    TORCH_CHECK(M64 <= INT_MAX && K64 <= INT_MAX && N64 <= INT_MAX, "matrix too large");
    const int M = (int)M64;
    const int K = (int)K64;
    const int N = (int)N64;

    c10::cuda::CUDAGuard guard(A.device());

    // The benchmark inputs are contiguous.  Make the wrapper robust without
    // falling back to PyTorch matmul: copies are only for layout normalization.
    if (!A.is_contiguous()) A = A.contiguous();
    if (!B.is_contiguous()) B = B.contiguous();

    auto C = torch::empty({M64, N64}, A.options());

    const float alpha = 1.0f;
    const float beta = 0.0f;

    // Row-major C = A(MxK) * B(KxN) is computed as column-major
    // C^T(NxM) = B^T(NxK) * A^T(KxM), reusing the original row-major storage.
    cublasHandle_t h = get_handle();
    cublasStatus_t st = cublasSgemm(
        h,
        CUBLAS_OP_N, CUBLAS_OP_N,
        N, M, K,
        &alpha,
        B.data_ptr<float>(), N,
        A.data_ptr<float>(), K,
        &beta,
        C.data_ptr<float>(), N
    );
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "cublasSgemm failed");
    return C;
}
'''

_ext = load_inline(
    name='matmul_cublas_ext_v1',
    cpp_sources=[_cpp_src],
    cuda_sources=[],
    functions=['matmul_forward'],
    extra_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        return _ext.matmul_forward(A, B)