gemini-2.5-pro / cudae94bd9

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include "kernel.h"
#include <stdexcept>

// Macro to check tensor properties, raising an error if a check fails.
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FP16(x) TORCH_CHECK(x.scalar_type() == torch::kFloat16, #x " must be a float16 tensor")

/**
 * @brief Main function exposed to Python for running the GEMM operation.
 *
 * This function takes two PyTorch tensors A and B, validates their properties,
 * allocates an output tensor C, and calls the custom CUDA WMMA kernel
 * to perform the computation C = A * B.T.
 *
 * @param A A torch::Tensor with shape [M, 2048] and dtype float16, on a CUDA device.
 * @param B A torch::Tensor with shape [5120, 2048] and dtype float16, on a CUDA device.
 * @return A torch::Tensor containing the result of the matrix multiplication.
 */
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // --- 1. Input Validation ---
    CHECK_CUDA(A);
    CHECK_CUDA(B);

    CHECK_FP16(A);
    CHECK_FP16(B);

    CHECK_CONTIGUOUS(A);
    CHECK_CONTIGUOUS(B);

    TORCH_CHECK(A.dim() == 2, "A must be a 2D tensor");
    TORCH_CHECK(B.dim() == 2, "B must be a 2D tensor");

    // --- 2. Shape Verification ---
    const int M = A.size(0);
    const int K_A = A.size(1);
    const int N_B = B.size(0);
    const int K_B = B.size(1);

    // Fixed dimensions from specification
    const int N_spec = 5120;
    const int K_spec = 2048;

    TORCH_CHECK(K_A == K_spec, "A tensor has incorrect K dimension: got ", K_A, ", expected ", K_spec);
    TORCH_CHECK(N_B == N_spec, "B tensor has incorrect N dimension: got ", N_B, ", expected ", N_spec);
    TORCH_CHECK(K_B == K_spec, "B tensor has incorrect K dimension: got ", K_B, ", expected ", K_spec);

    // --- 3. Output Tensor Allocation ---
    auto options = torch::TensorOptions()
        .device(A.device())
        .dtype(torch::kFloat16);
    torch::Tensor C = torch::empty({M, N_spec}, options);

    // --- 4. Kernel Launch ---
    // Get the current CUDA stream from PyTorch to ensure proper synchronization.
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Get raw data pointers from PyTorch tensors.
    const half* ptr_A = reinterpret_cast<const half*>(A.data_ptr<at::Half>());
    const half* ptr_B = reinterpret_cast<const half*>(B.data_ptr<at::Half>());
    half* ptr_C = reinterpret_cast<half*>(C.data_ptr<at::Half>());

    // Call the CUDA kernel launcher.
    gemm_n5120_k2048_launcher(ptr_A, ptr_B, ptr_C, M, stream);

    // Check for any asynchronous errors from the kernel launch.
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(std::string("CUDA error after kernel launch: ") + cudaGetErrorString(err));
    }
    
    return C;
}

// --- 5. Pybind11 Module Definition ---
// This creates the Python module and exposes the `run` function.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "GEMM C[M, N] = A[M, K] @ B.T[K, N] where N=5120, K=2048 (CUDA/WMMA implementation)");
}