gpt-o3_cuda_270394

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

#include "kernel.h"

/******************************************************************************
 *  Python-visible entry point
 *****************************************************************************/
torch::Tensor run(torch::Tensor A, torch::Tensor B)
{
    /* ---------------- argument checking ---------------- */
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2,
                "A and B must be 2-D tensors");
    TORCH_CHECK(A.size(1) == 2048 &&
                B.size(0) == 128  && B.size(1) == 2048,
                "Shapes must be  A[M,2048]  and  B[128,2048]");
    TORCH_CHECK(A.scalar_type() == at::kHalf &&
                B.scalar_type() == at::kHalf,
                "Tensors must be float16");
    TORCH_CHECK(A.is_cuda() && B.is_cuda(),
                "Tensors have to live on CUDA");

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

    const int64_t M = A_c.size(0);

    /* output tensor */
    auto options = torch::TensorOptions()
                     .dtype(at::kHalf)
                     .device(A.device());
    auto C = torch::empty({M, 128}, options);

    /* temporary buffer for transposed B */
    auto B_col = torch::empty({2048, 128}, options);

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

    /* 1. transpose B  -------------------------------------------------------- */
    launch_transpose_B(
        reinterpret_cast<const __half *>(B_c.data_ptr<at::Half>()),
        reinterpret_cast<      __half *>(B_col.data_ptr<at::Half>()),
        stream);

    /* 2. GEMM ---------------------------------------------------------------- */
    launch_gemm_n128_k2048(
        reinterpret_cast<const __half *>(A_c.data_ptr<at::Half>()),
        reinterpret_cast<const __half *>(B_col.data_ptr<at::Half>()),
        reinterpret_cast<      __half *>(C.data_ptr<at::Half>()),
        static_cast<int>(M),
        stream);

    /* make sure the kernel finished before returning to Python */
    CUDA_CHECK(cudaStreamSynchronize(stream));

    return C;
}

/******************************************************************************
 *  PyBind11 module definition
 *****************************************************************************/
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("run", &run,
          "gemm_n128_k2048 (CUDA, FP16)  –  C = A @ B.T");
}