claude-opus-4-1 / cudad26d88

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

// Macro for checking CUDA errors
#define CUDA_CHECK(call) \
    do { \
        cudaError_t error = call; \
        if (error != cudaSuccess) { \
            AT_ERROR("CUDA error at ", __FILE__, ":", __LINE__, \
                     " code=", error, "(", cudaGetErrorString(error), ")"); \
        } \
    } while(0)

// Check tensor properties
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FP16(x) TORCH_CHECK(x.scalar_type() == torch::kFloat16, #x " must be float16")

torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // Input validation
    CHECK_CUDA(A);
    CHECK_CUDA(B);
    CHECK_CONTIGUOUS(A);
    CHECK_CONTIGUOUS(B);
    
    // Check dimensions
    TORCH_CHECK(A.dim() == 2, "A must be 2D, got ", A.dim(), "D");
    TORCH_CHECK(B.dim() == 2, "B must be 2D, got ", B.dim(), "D");
    
    // Get dimensions
    const int64_t M = A.size(0);
    const int64_t K_A = A.size(1);
    const int64_t N_B = B.size(0);
    const int64_t K_B = B.size(1);
    
    // Validate dimensions
    TORCH_CHECK(K_A == K_FIXED, "A must have K dimension = ", K_FIXED, ", got ", K_A);
    TORCH_CHECK(N_B == N_FIXED, "B must have N dimension = ", N_FIXED, ", got ", N_B);
    TORCH_CHECK(K_B == K_FIXED, "B must have K dimension = ", K_FIXED, ", got ", K_B);
    
    // Convert to fp16 if necessary
    torch::Tensor A_fp16 = A;
    torch::Tensor B_fp16 = B;
    
    if (A.scalar_type() != torch::kFloat16) {
        A_fp16 = A.to(torch::kFloat16);
    }
    if (B.scalar_type() != torch::kFloat16) {
        B_fp16 = B.to(torch::kFloat16);
    }
    
    // Ensure contiguous
    A_fp16 = A_fp16.contiguous();
    B_fp16 = B_fp16.contiguous();
    
    // Create output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::kFloat16)
        .device(A.device())
        .requires_grad(false);
    
    torch::Tensor C = torch::zeros({M, N_FIXED}, options);
    
    // Get raw pointers
    const half* A_ptr = reinterpret_cast<const half*>(A_fp16.data_ptr<at::Half>());
    const half* B_ptr = reinterpret_cast<const half*>(B_fp16.data_ptr<at::Half>());
    half* C_ptr = reinterpret_cast<half*>(C.data_ptr<at::Half>());
    
    // Get current CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_gemm_kernel(A_ptr, B_ptr, C_ptr, static_cast<int>(M), stream);
    
    // Ensure kernel completion for correctness
    CUDA_CHECK(cudaStreamSynchronize(stream));
    
    return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Optimized GEMM kernel for (M x 4096) @ (6144 x 4096)^T",
          py::arg("A"), py::arg("B"));
}