import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs 4D tensor-matrix multiplication: 
        C[b, i, j, k] = sum_l A[b, i, j, l] * B[l, k]

    Args:
        A (torch.Tensor): Input 4D tensor of shape (b, i, j, l)
        B (torch.Tensor): Input matrix of shape (l, k)

    Returns:
        torch.Tensor: Output 4D tensor of shape (b, i, j, k)
    """
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, A, B):
        """
        Performs the 4D tensor-matrix multiplication.

        Args:
            A (torch.Tensor): Input 4D tensor of shape (b, i, j, l)
            B (torch.Tensor): Input matrix of shape (l, k)

        Returns:
            torch.Tensor: Output 4D tensor of shape (b, i, j, k)
        """
        return torch.einsum("bijl,lk->bijk", A, B)

# Test code
b = 8
i = 256
j = 512
l = 256
k = 768

def get_inputs():
    A = torch.rand(b, i, j, l)
    B = torch.rand(l, 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

_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(cublasStatus_t status) {
    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, "cuBLAS error: ", (int)status);
}

torch::Tensor tensor4_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() == at::kFloat && B.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(A.dim() == 4 && B.dim() == 2, "expected A 4D and B 2D");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(), "inputs must be contiguous");
    TORCH_CHECK(A.size(3) == B.size(0), "contracting dimensions differ");

    const auto b = A.size(0);
    const auto i = A.size(1);
    const auto j = A.size(2);
    const auto l = A.size(3);
    const auto k = B.size(1);
    const int64_t M64 = b * i * j;
    TORCH_CHECK(M64 <= INT_MAX && k <= INT_MAX && l <= INT_MAX, "dimensions too large for cuBLAS sgemm");

    auto C = torch::empty({b, i, j, k}, A.options());
    if (M64 == 0 || k == 0) return C;

    c10::cuda::CUDAGuard device_guard(A.device());
    auto handle = at::cuda::getCurrentCUDABlasHandle();

    const int M = (int)M64;
    const int N = (int)k;
    const int L = (int)l;
    const float alpha = 1.0f;
    const float beta = 0.0f;

    // Row-major C[M,N] = A[M,L] * B[L,N].  cuBLAS is column-major, so compute
    // C_col[N,M] = B_col[N,L] * A_col[L,M] over the same underlying storage.
    check_cublas(cublasSgemm(handle,
                             CUBLAS_OP_N, CUBLAS_OP_N,
                             N, M, L,
                             &alpha,
                             B.data_ptr<float>(), N,
                             A.data_ptr<float>(), L,
                             &beta,
                             C.data_ptr<float>(), N));
    return C;
}
'''

_ext = load_inline(
    name='tensor4_matmul_cublas_ext',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['tensor4_matmul_forward'],
    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.tensor4_matmul_forward(A, B)