import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs 3D tensor-matrix multiplication.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A, B):
        """
        Performs 3D tensor-matrix multiplication.

        Args:
            A (torch.Tensor): Input 3D tensor of shape (N, M, K).
            B (torch.Tensor): Input matrix of shape (K, L).

        Returns:
            torch.Tensor: Output tensor of shape (N, M, L), resulting from the multiplication of A and B along the last dimension of A.
        """
        return torch.matmul(A, B)

N = 16
M = 1024
K = 2048
L = 768

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

torch::Tensor matmul3d_forward(torch::Tensor A, torch::Tensor B);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cublas_v2.h>

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

torch::Tensor matmul3d_forward(torch::Tensor A, torch::Tensor B) {
    c10::cuda::CUDAGuard device_guard(A.device());

    const int64_t N = A.size(0);
    const int64_t M = A.size(1);
    const int64_t K = A.size(2);
    const int64_t L = B.size(1);
    const int64_t rows = N * M;

    auto C = torch::empty({N, M, L}, A.options());

    // Row-major C[rows,L] = A[rows,K] @ B[K,L].
    // Interpreting the same buffers as column-major gives:
    //   C_col[L,rows] = B_col[L,K] * A_col[K,rows]
    // which is exactly the transpose view of the desired row-major result.
    const float alpha = 1.0f;
    const float beta  = 0.0f;

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    auto stream = at::cuda::getCurrentCUDAStream();
    check_cublas_status(cublasSetStream(handle, stream.stream()));

    check_cublas_status(cublasSgemm(
        handle,
        CUBLAS_OP_N, CUBLAS_OP_N,
        static_cast<int>(L), static_cast<int>(rows), static_cast<int>(K),
        &alpha,
        B.data_ptr<float>(), static_cast<int>(L),
        A.data_ptr<float>(), static_cast<int>(K),
        &beta,
        C.data_ptr<float>(), static_cast<int>(L)
    ));

    return C;
}
"""

_ext = load_inline(
    name="matmul3d_cublas_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["matmul3d_forward"],
    extra_cflags=["-O3"],
    extra_cuda_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, B):
        return _ext.matmul3d_forward(A, B)