#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>

// Custom implementation using cuBLAS batched GEMM for 3D tensor-matrix multiplication.
// We want to compute, for each batch element n:
//      C[n] = A[n] * B,
// where A[n] is a matrix of shape (M, K) stored in row-major order,
// B is a matrix of shape (K, L) stored in row-major order (common to all batches),
// and the output C[n] should be of shape (M, L) in row-major order.
//
// To leverage cuBLAS which expects column-major data, we use the identity:
//     (A_r * B_r)^T = B_r^T * A_r^T,
// where A_r and B_r are the row-major matrices.
// We can then reinterpret the memory layout as follows:
//  - Interpret A (M×K row-major) as A_c = A^T (K×M column-major). Thus, ldb = K and stride for A = M*K.
//  - Interpret B (K×L row-major) as B_c = B^T (L×K column-major). Thus, lda = L and stride for B = 0.
//  - The resulting C will be computed as a column-major matrix C_c of dimensions (L, M) (with ldc = L and stride = M*L),
//    which, when reinterpreted in row-major, gives (M, L).
//
// GEMM call parameters:
//   transa = CUBLAS_OP_N and transb = CUBLAS_OP_N,
//   m = L, n = M, k = K,
// with batch count = N.
//
// Additionally, we set the cuBLAS math mode to CUBLAS_DEFAULT_MATH to obtain results that exactly match torch.matmul.
torch::Tensor tensor_matmul(torch::Tensor A, torch::Tensor B) {
    // Enforce correct CUDA device.
    c10::cuda::CUDAGuard device_guard(A.device());
    auto stream = c10::cuda::getCurrentCUDAStream();

    // Dimensions: A: (N, M, K), B: (K, L). Desired output C: (N, M, L).
    int N = A.size(0);
    int M = A.size(1);
    int K = A.size(2);
    TORCH_CHECK(B.dim() == 2, "B must be a 2D tensor");
    TORCH_CHECK(B.size(0) == K, "Dimension mismatch: the last dimension of A must equal the first dimension of B");
    int L = B.size(1);

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

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

    // GEMM parameters via row-major to column-major conversion:
    // For B: interpreted as B_c of dimensions (L, K), so lda = L.
    // For A: interpreted as A_c of dimensions (K, M), so ldb = K.
    // For C: computed as C_c of dimensions (L, M), so ldc = L.
    int lda = L;   // Leading dimension for B interpreted as B_c.
    int ldb = K;   // Leading dimension for A interpreted as A_c.
    int ldc = L;   // Leading dimension for output C interpreted as C_c.

    // Strides:
    long long strideA = 0;              // B is common for all batches.
    long long strideB = (long long) M * K;  // Each batch of A occupies M*K elements.
    long long strideC = (long long) M * L;  // Each batch of C occupies M*L elements.

    // Get or create a static cuBLAS handle.
    static cublasHandle_t handle = nullptr;
    if (handle == nullptr) {
        cublasCreate(&handle);
        // Set math mode to default to ensure numerical correctness.
        cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH);
    } else {
        cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH);
    }
    cublasSetStream(handle, stream);

    // Call cuBLAS batched GEMM:
    // Compute C_c = B_c * A_c which yields a column-major matrix of dimensions (L, M),
    // which when reinterpreted in row-major produces the desired (M, L) matrix.
    cublasStatus_t status = cublasSgemmStridedBatched(
        handle,
        CUBLAS_OP_N, // B is not transposed; interpreted as B_c.
        CUBLAS_OP_N, // A is not transposed; interpreted as A_c.
        /* m = */ L,  // Number of rows of B_c.
        /* n = */ M,  // Number of columns of A_c.
        /* k = */ K,  // Shared dimension.
        &alpha,
        /* A: */ B.data_ptr<float>(), lda, strideA,  // B is common for all batches.
        /* B: */ A.data_ptr<float>(), ldb, strideB,  // A has a batch stride of M*K.
        &beta,
        /* C: */ C.data_ptr<float>(), ldc, strideC,
        N  // Batch count.
    );
    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, "cublasSgemmStridedBatched failed");

    return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("tensor_matmul", &tensor_matmul, "3D tensor-matrix multiplication using cuBLAS batched GEMM");
}