gemini-2.5-pro / cudacda2a1

#include <torch/extension.h>
#include <vector>
#include "kernel.h"

// C++ type from PyTorch for FP16
using at::Half;

// --- PyTorch Binding ---

/**
 * @brief Python-callable function to execute the GEMM operation.
 *
 * This function acts as the interface between Python (PyTorch) and the custom CUDA kernel.
 * It handles tensor validation, memory management, and kernel launching.
 *
 * @param A A PyTorch tensor of shape [M, 4096] and dtype float16, located on a CUDA device.
 * @param B A PyTorch tensor of shape [2048, 4096] and dtype float16, located on the same CUDA device.
 * @return A new PyTorch tensor C of shape [M, 2048] and dtype float16, containing the result of A * B.T.
 */
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // --- Input Validation ---
    TORCH_CHECK(A.dim() == 2, "Input tensor A must be 2-dimensional");
    TORCH_CHECK(B.dim() == 2, "Input tensor B must be 2-dimensional");

    TORCH_CHECK(A.is_cuda(), "Input tensor A must be a CUDA tensor");
    TORCH_CHECK(B.is_cuda(), "Input tensor B must be a CUDA tensor");

    TORCH_CHECK(A.device() == B.device(), "Input tensors A and B must be on the same device");

    TORCH_CHECK(A.scalar_type() == torch::kFloat16, "Input tensor A must have dtype float16");
    TORCH_CHECK(B.scalar_type() == torch::kFloat16, "Input tensor B must have dtype float16");

    // Ensure tensors are contiguous in memory, as the kernel assumes a packed layout.
    A = A.contiguous();
    B = B.contiguous();

    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);

    // Check fixed dimensions
    const int N_spec = 2048;
    const int K_spec = 4096;
    TORCH_CHECK(N_B == N_spec, "Input tensor B must have N=2048 rows, but got ", N_B);
    TORCH_CHECK(K_A == K_spec, "Input tensor A must have K=4096 columns, but got ", K_A);
    TORCH_CHECK(K_B == K_spec, "Input tensor B must have K=4096 columns, but got ", K_B);

    // --- Output Tensor Allocation ---
    // Create the output tensor C with the correct shape and options (device, dtype)
    auto C = torch::empty({M, N_spec}, A.options());

    // --- Kernel Execution ---
    // Get raw data pointers from the PyTorch tensors
    half* C_ptr = reinterpret_cast<half*>(C.data_ptr<Half>());
    const half* A_ptr = reinterpret_cast<const half*>(A.data_ptr<Half>());
    const half* B_ptr = reinterpret_cast<const half*>(B.data_ptr<Half>());

    // Call the host launcher function from the CUDA file
    gemm_n2048_k4096_cuda(C_ptr, A_ptr, B_ptr, M);

    return C;
}

// Binds the C++ `run` function to a Python module.
// This allows calling `gemm_n2048_k4096.run(A, B)` from Python.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "GEMM C=A*B.T (N=2048, K=4096) implementation on B200");
}