gemini-2.5-pro / cuda1d80a9

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

#include "kernel.h"

#include <stdexcept>
#include <string>

// Helper macros for concise tensor validation
#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_HALF(x) TORCH_CHECK(x.scalar_type() == torch::kFloat16, #x " must be a float16 tensor")

/**
 * @brief PyTorch extension entry point for the GEMM operation.
 *
 * This function validates input tensors and calls the CUDA kernel launcher
 * to perform the computation C = A * B.T on the GPU.
 *
 * @param A A torch::Tensor of shape [M, 4096] and dtype float16.
 * @param B A torch::Tensor of shape [4096, 4096] and dtype float16.
 * @return A torch::Tensor C of shape [M, 4096] and dtype float16 containing the result.
 */
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // --- Input Validation ---
    CHECK_CUDA(A);
    CHECK_CUDA(B);
    CHECK_CONTIGUOUS(A);
    CHECK_CONTIGUOUS(B);
    CHECK_HALF(A);
    CHECK_HALF(B);

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

    // --- Dimension Checks ---
    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);

    const int N_spec = 4096;
    const int K_spec = 4096;

    TORCH_CHECK(K_A == K_spec, "A must have shape [M, 4096], but K is ", K_A);
    TORCH_CHECK(N_B == N_spec, "B must have shape [4096, 4096], but N is ", N_B);
    TORCH_CHECK(K_B == K_spec, "B must have shape [4096, 4096], but K is ", K_B);
    TORCH_CHECK(A.device() == B.device(), "Tensors must be on the same CUDA device");

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

    // --- Kernel Execution ---
    try {
        // Get the current CUDA stream from PyTorch's context to ensure proper synchronization
        cudaStream_t stream = at::cuda::getCurrentCUDAStream();

        // Get raw data pointers. at::Half is compatible with cuda_fp16.h::half
        const half* A_ptr = reinterpret_cast<const half*>(A.data_ptr<at::Half>());
        const half* B_ptr = reinterpret_cast<const half*>(B.data_ptr<at::Half>());
        half* C_ptr = reinterpret_cast<half*>(C.data_ptr<at::Half>());

        // Launch the cuBLAS-based kernel
        gemm_n4096_k4096_launcher(M, A_ptr, B_ptr, C_ptr, stream);

    } catch (const std::exception& e) {
        // Propagate exceptions from the CUDA/cuBLAS calls to Python
        throw std::runtime_error(std::string("CUDA kernel execution failed: ") + e.what());
    }

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

    return C;
}

// Pybind11 module definition to expose the 'run' function to Python
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "GEMM (A * B.T) for N=4096, K=4096 using a B200-optimized cuBLAS kernel");
}