claude-opus-4-1_cuda_b43068

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

// Helper function to check CUDA errors
#define CHECK_CUDA(x) do { \
    cudaError_t err = x; \
    if (err != cudaSuccess) { \
        throw std::runtime_error(std::string("CUDA error: ") + cudaGetErrorString(err)); \
    } \
} while(0)

// Helper macros for tensor checks
#define CHECK_INPUT(x) do { \
    TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor"); \
    TORCH_CHECK(x.is_contiguous(), #x " must be contiguous"); \
    TORCH_CHECK(x.dtype() == torch::kFloat16, #x " must be float16"); \
} while(0)

// Main entry point function
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // Validate inputs
    CHECK_INPUT(A);
    CHECK_INPUT(B);
    
    // Get dimensions
    const int M = A.size(0);
    const int K_A = A.size(1);
    const int N = B.size(0);
    const int K_B = B.size(1);
    
    // Validate dimensions
    TORCH_CHECK(K_A == K_SIZE, "A dimension K must be " + std::to_string(K_SIZE) + ", got " + std::to_string(K_A));
    TORCH_CHECK(K_B == K_SIZE, "B dimension K must be " + std::to_string(K_SIZE) + ", got " + std::to_string(K_B));
    TORCH_CHECK(N == N_SIZE, "B dimension N must be " + std::to_string(N_SIZE) + ", got " + std::to_string(N));
    
    // Create output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::kFloat16)
        .device(A.device())
        .requires_grad(false);
    torch::Tensor C = torch::empty({M, N_SIZE}, options);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch custom kernel
    launch_gemm_kernel(
        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,
        stream
    );
    
    // Ensure kernel completes successfully
    CHECK_CUDA(cudaGetLastError());
    
    return C;
}

// Python binding
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Optimized GEMM kernel for M x 14336 * 4096 x 14336 -> M x 4096",
          py::arg("A"), py::arg("B"));
}