gpt-5-2025-08-07 / cuda351c51

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <vector>
#include <stdexcept>
#include <sstream>
#include "kernel.h"

namespace py = pybind11;

#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_DTYPE_HALF(x) TORCH_CHECK(x.dtype() == torch::kFloat16, #x " must be torch.float16")
#define CUDA_CHECK(err) do { \
    cudaError_t err__ = (err); \
    TORCH_CHECK(err__ == cudaSuccess, "CUDA error: ", cudaGetErrorString(err__)); \
} while (0)

static inline std::string shape_str(const torch::Tensor& t) {
    std::ostringstream oss;
    oss << "[";
    for (int i = 0; i < t.dim(); ++i) {
        oss << t.size(i);
        if (i + 1 < t.dim()) oss << ", ";
    }
    oss << "]";
    return oss.str();
}

torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // Validate inputs
    TORCH_CHECK(A.dim() == 2, "A must be 2D, got ", A.dim(), "D with shape ", shape_str(A));
    TORCH_CHECK(B.dim() == 2, "B must be 2D, got ", B.dim(), "D with shape ", shape_str(B));

    // Move to CUDA if needed and make contiguous + correct dtype
    if (!A.is_cuda()) A = A.to(torch::kCUDA);
    if (!B.is_cuda()) B = B.to(torch::kCUDA);
    if (!A.is_contiguous()) A = A.contiguous();
    if (!B.is_contiguous()) B = B.contiguous();
    if (A.dtype() != torch::kFloat16) A = A.to(torch::kFloat16);
    if (B.dtype() != torch::kFloat16) B = B.to(torch::kFloat16);

    CHECK_CUDA(A);
    CHECK_CUDA(B);
    CHECK_CONTIGUOUS(A);
    CHECK_CONTIGUOUS(B);
    CHECK_DTYPE_HALF(A);
    CHECK_DTYPE_HALF(B);

    const int64_t M64 = A.size(0);
    const int64_t K64 = A.size(1);
    TORCH_CHECK(K64 == GEMM_K_CONST, "A.shape[1] must be ", GEMM_K_CONST, ", got ", K64);
    TORCH_CHECK(B.size(0) == GEMM_N_CONST && B.size(1) == GEMM_K_CONST,
                "B must have shape [", GEMM_N_CONST, ", ", GEMM_K_CONST, "], got ", shape_str(B));

    // Output tensor C [M, N=256]
    auto options = A.options().dtype(torch::kFloat16);
    torch::Tensor C = torch::empty({M64, GEMM_N_CONST}, options);

    if (M64 == 0) {
        return C;
    }

    const int M = static_cast<int>(M64);
    const int lda = GEMM_K_CONST; // row-major A
    const int ldb = GEMM_K_CONST; // row-major B
    const int ldc = GEMM_N_CONST; // row-major C

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Launch kernel via launcher in .cu (avoids <<<>>> in this .cpp TU)
    gemm_n256_k7168_launch(
        reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
        reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
        reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
        M,
        lda, ldb, ldc,
        stream
    );

    CUDA_CHECK(cudaGetLastError());
    // No explicit sync; PyTorch stream semantics handle dependencies

    return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "gemm_n256_k7168 CUDA kernel",
          py::arg("A"), py::arg("B"));
}