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.T, B.T)

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

def get_inputs():
    A = torch.rand(K, M)
    B = torch.rand(N, K)
    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>

#define CUBLAS_CHECK(cmd) do {                                      \
    cublasStatus_t status = (cmd);                                  \
    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,                    \
                "cuBLAS error, status=", (int)status);             \
} while (0)

torch::Tensor matmul_at_bt(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor");
    TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor");
    TORCH_CHECK(A.scalar_type() == at::kFloat, "A must be float32");
    TORCH_CHECK(B.scalar_type() == at::kFloat, "B must be float32");
    TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
    TORCH_CHECK(B.is_contiguous(), "B must be contiguous");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "inputs must be 2D");
    TORCH_CHECK(A.size(0) == B.size(1), "inner dimension mismatch");

    const int64_t K64 = A.size(0);
    const int64_t M64 = A.size(1);
    const int64_t N64 = B.size(0);
    TORCH_CHECK(K64 <= INT_MAX && M64 <= INT_MAX && N64 <= INT_MAX, "matrix too large");

    const int K = static_cast<int>(K64);
    const int M = static_cast<int>(M64);
    const int N = static_cast<int>(N64);

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

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    CUBLAS_CHECK(cublasSetStream(handle, at::cuda::getCurrentCUDAStream()));
    // Match fp32 PyTorch matmul numerics (no TF32 tensor-core approximation).
    CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));

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

    // Row-major A is KxM and row-major B is NxK.  Interpreted as column-major,
    // A stores an MxK matrix A_col with A_col[m,k] = A[k,m], and B stores a
    // KxN matrix B_col with B_col[k,n] = B[n,k].  The row-major output C (MxN)
    // is the same memory as a column-major NxM matrix C_col = C^T, so compute:
    //   C_col = B_col^T * A_col^T.
    CUBLAS_CHECK(cublasSgemm(handle,
                             CUBLAS_OP_T, CUBLAS_OP_T,
                             N, M, K,
                             &alpha,
                             B.data_ptr<float>(), K,
                             A.data_ptr<float>(), M,
                             &beta,
                             C.data_ptr<float>(), N));
    return C;
}
'''

_ext = load_inline(
    name='matmul_at_bt_cublas_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources='',
    functions=['matmul_at_bt'],
    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_at_bt(A, B)