import torch
import torch.nn as nn

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

        Args:
            A (torch.Tensor): Input matrix A of shape (N, N).
            B (torch.Tensor): Input matrix B of shape (N, N).

        Returns:
            torch.Tensor: Output matrix C of shape (N, N).
        """
        return torch.matmul(A, B)

N = 2048 * 2

def get_inputs():
    A = torch.rand(N, N)
    B = torch.rand(N, 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 <cuda_runtime.h>
#include <cublas_v2.h>

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

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, "inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "inputs must be matrices");
    TORCH_CHECK(A.size(0) == A.size(1) && B.size(0) == B.size(1) && A.size(1) == B.size(0), "bad matrix sizes");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(), "inputs must be contiguous");

    const int64_t n64 = A.size(0);
    TORCH_CHECK(n64 <= INT_MAX, "matrix too large");
    const int n = (int)n64;

    c10::cuda::CUDAGuard device_guard(A.device());
    auto C = torch::empty({n64, n64}, A.options());

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    const float alpha = 1.0f;
    const float beta = 0.0f;

    // Tensors are row-major.  cuBLAS is column-major, so computing
    // C_col = B_col * A_col produces the row-major C = A @ B in memory.
#if defined(CUBLAS_COMPUTE_32F_FAST_TF32)
    check_cublas(cublasGemmEx(
        handle,
        CUBLAS_OP_N, CUBLAS_OP_N,
        n, n, n,
        &alpha,
        (const void*)B.data_ptr<float>(), CUDA_R_32F, n,
        (const void*)A.data_ptr<float>(), CUDA_R_32F, n,
        &beta,
        (void*)C.data_ptr<float>(), CUDA_R_32F, n,
        CUBLAS_COMPUTE_32F_FAST_TF32,
        CUBLAS_GEMM_DEFAULT_TENSOR_OP));
#else
    check_cublas(cublasSgemm(
        handle,
        CUBLAS_OP_N, CUBLAS_OP_N,
        n, n, n,
        &alpha,
        B.data_ptr<float>(), n,
        A.data_ptr<float>(), n,
        &beta,
        C.data_ptr<float>(), n));
#endif
    return C;
}
'''

_ext = load_inline(
    name='modelnew_cublas_tf32_matmul_ext',
    cpp_sources=_cpp_src,
    functions=['matmul_forward'],
    with_cuda=True,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    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)