gemini-2.5-pro / cudaadc04b

#include <torch/extension.h>
#include <stdexcept>
#include <string>
#include <vector>

#include "kernel.h"

// CUDA API error checking macro
#define CUDA_CHECK(status)                                                     \
  do {                                                                         \
    cudaError_t error = status;                                                \
    if (error != cudaSuccess) {                                                \
      throw std::runtime_error(std::string("CUDA error in " __FILE__ ":" +     \
                                           std::to_string(__LINE__)) +         \
                               ": " + cudaGetErrorString(error));              \
    }                                                                          \
  } while (0)

/**
 * @brief Python-bindable function that serves as the entry point.
 *
 * This function validates input tensors from PyTorch, prepares memory,
 * and calls the CUDA kernel launcher.
 *
 * @param A A torch::Tensor of shape [M, 14336] and dtype float16.
 * @param B A torch::Tensor of shape [4096, 14336] and dtype float16.
 * @return A torch::Tensor of shape [M, 4096] and dtype float16 containing the result.
 */
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
  // --- Input Validation ---
  TORCH_CHECK(A.is_cuda(), "Input tensor A must be on a CUDA device");
  TORCH_CHECK(B.is_cuda(), "Input tensor B must be on a CUDA device");
  TORCH_CHECK(A.device() == B.device(),
              "Input tensors A and B must be on the same CUDA device");

  TORCH_CHECK(A.scalar_type() == torch::kFloat16,
              "Input tensor A must be of type float16");
  TORCH_CHECK(B.scalar_type() == torch::kFloat16,
              "Input tensor B must be of type float16");

  TORCH_CHECK(A.dim() == 2, "Input tensor A must be 2-dimensional");
  TORCH_CHECK(B.dim() == 2, "Input tensor B must be 2-dimensional");

  // Check fixed dimensions as per specification
  const int N_fixed = 4096;
  const int K_fixed = 14336;
  TORCH_CHECK(B.size(0) == N_fixed, "Input tensor B must have N=", N_fixed,
              " rows, but got ", B.size(0));
  TORCH_CHECK(A.size(1) == K_fixed, "Input tensor A must have K=", K_fixed,
              " columns, but got ", A.size(1));
  TORCH_CHECK(B.size(1) == K_fixed, "Input tensor B must have K=", K_fixed,
              " columns, but got ", B.size(1));

  // Ensure tensors are contiguous for predictable memory layout
  A = A.contiguous();
  B = B.contiguous();

  // Get problem dimensions from input tensors
  const int M = A.size(0);
  const int N = B.size(0);

  // --- Output Tensor Creation ---
  auto C_options =
      torch::TensorOptions().device(A.device()).dtype(torch::kFloat16);
  torch::Tensor C = torch::empty({M, N}, C_options);
  if (M == 0) {
    return C;
  }

  // --- Kernel Execution ---
  // Get raw data pointers. We must reinterpret_cast because at::Half and __half
  // are distinct types.
  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>());

  // Get the current CUDA stream from PyTorch's context
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  // Launch the CUDA kernel
  gemm_n4096_k14336_cuda(M, A_ptr, B_ptr, C_ptr, stream);

  // Check for any asynchronous errors from the kernel launch
  CUDA_CHECK(cudaGetLastError());

  return C;
}

// Pybind11 module definition to expose the 'run' function to Python
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run,
        "GEMM N=4096, K=14336 (A[M,K] @ B[N,K].T -> C[M,N]) implementation for "
        "B200 using WMMA",
        py::arg("A"), py::arg("B"));
}