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

cpp_source = """
#include <torch/extension.h>
torch::Tensor einsum_forward(torch::Tensor A, torch::Tensor B);
"""

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

__global__ void matmul4d_tiny_kernel(const float* __restrict__ A,
                                     const float* __restrict__ B,
                                     float* __restrict__ C,
                                     long long M, int L, int K) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = M * (long long)K;
    if (idx >= total) return;
    int k = (int)(idx % K);
    long long m = idx / K;
    float acc = 0.0f;
    for (int l = 0; l < L; ++l) acc += A[m * (long long)L + l] * B[l * (long long)K + k];
    C[idx] = acc;
}

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

torch::Tensor einsum_forward(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 && B.scalar_type() == torch::kFloat32, "float32 only");
    TORCH_CHECK(A.dim() == 4 && B.dim() == 2, "bad dims");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(), "contiguous required");

    const auto b = A.size(0);
    const auto ii = A.size(1);
    const auto j = A.size(2);
    const auto L64 = A.size(3);
    TORCH_CHECK(B.size(0) == L64, "contracting dimension mismatch");
    const auto K64 = B.size(1);

    auto C = torch::empty({b, ii, j, K64}, A.options());
    const int64_t M64 = b * ii * j;
    if (M64 == 0 || K64 == 0 || L64 == 0) return C;
    TORCH_CHECK(K64 <= INT_MAX && L64 <= INT_MAX, "dims too large");
    const int K = (int)K64;
    const int L = (int)L64;

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

    // A minimal CUDA kernel is used for tiny cases where cuBLAS launch/heuristic overhead dominates.
    // The benchmark dimensions take the SGEMM path below.
    if (M64 * K64 < 4096) {
        const int threads = 256;
        const int blocks = (int)((M64 * K64 + threads - 1) / threads);
        matmul4d_tiny_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M64, L, K);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return C;
    }

    TORCH_CHECK(M64 <= INT_MAX, "M too large for cublas sgemm");
    const int M = (int)M64;

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    check_cublas(cublasSetStream(handle, at::cuda::getCurrentCUDAStream()));
    // Use full fp32 SGEMM math so results match torch.einsum at fp32 tolerance.
    check_cublas(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));

    const float alpha = 1.0f;
    const float beta = 0.0f;
    // Row-major C[M,K] = A[M,L] * B[L,K] is computed as column-major
    // C^T[K,M] = B^T[K,L] * A^T[L,M], using the original contiguous buffers.
    check_cublas(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                             K, M, L,
                             &alpha,
                             B.data_ptr<float>(), K,
                             A.data_ptr<float>(), L,
                             &beta,
                             C.data_ptr<float>(), K));
    return C;
}
'''

_ext = load_inline(
    name='einsum4d_cublas_kernel_ext_final',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['einsum_forward'],
    verbose=False,
    with_cuda=True,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
)

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

    def forward(self, A, B):
        return _ext.einsum_forward(A, B)