gpt-o3 / cuda7a2145

#include "kernel.h"
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

/*
 *  Python interface
 *
 *      C = run(A, B)
 *
 *  A : [M, 7168]   torch.float16 (CUDA)  – row-major
 *  B : [256,7168]  torch.float16 (CUDA)  – row-major
 *  C : [M, 256 ]   torch.float16 (CUDA)  – row-major
 */

torch::Tensor run(torch::Tensor A, torch::Tensor B)
{
    /*  sanity checks -------------------------------------------------------- */
    TORCH_CHECK(A.is_cuda(), "A must reside on CUDA device");
    TORCH_CHECK(B.is_cuda(), "B must reside on CUDA device");
    TORCH_CHECK(A.scalar_type() == at::kHalf, "A must be float16");
    TORCH_CHECK(B.scalar_type() == at::kHalf, "B must be float16");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be rank-2 tensors");
    TORCH_CHECK(A.size(1) == 7168,
                "A has wrong second dimension (expected 7168)");
    TORCH_CHECK(B.size(0) == 256 && B.size(1) == 7168,
                "B must have shape [256, 7168]");

    /*  make the inputs contiguous (no-op if already) ----------------------- */
    auto A_c = A.contiguous();
    auto B_c = B.contiguous();

    const int64_t M = A_c.size(0);

    /*  allocate output ------------------------------------------------------ */
    auto C = torch::empty({M, 256},
                          torch::TensorOptions()
                              .dtype(at::kHalf)
                              .device(A.device()));

    /*  raw pointers --------------------------------------------------------- */
    const __half* d_A = reinterpret_cast<const __half*>(A_c.data_ptr<at::Half>());
    const __half* d_B = reinterpret_cast<const __half*>(B_c.data_ptr<at::Half>());
    __half*       d_C = reinterpret_cast<__half*>(C.data_ptr<at::Half>());

    /*  current CUDA stream -------------------------------------------------- */
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    /*  launch specialised kernel ------------------------------------------- */
    launch_gemm_n256_k7168(d_A, d_B, d_C, static_cast<int>(M), stream);

    /*  ensure completion ---------------------------------------------------- */
    cudaError_t err = cudaStreamSynchronize(stream);
    TORCH_CHECK(err == cudaSuccess,
                "CUDA kernel failed : ",
                cudaGetErrorString(err));

    return C;
}

/* -------------------------- PyBind registration --------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run,
          "Optimised GEMM  C = A·B^T  (A[M,7168]  ·  B[256,7168]^T)",
          py::arg("A"),
          py::arg("B"));
}