claude-opus-4-1-20250805 / cuda896e50

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

// Main entry point function
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // Validate input tensors
    TORCH_CHECK(A.dim() == 2, "A must be 2D tensor");
    TORCH_CHECK(B.dim() == 2, "B must be 2D tensor");
    TORCH_CHECK(A.dtype() == torch::kFloat16, "A must be float16");
    TORCH_CHECK(B.dtype() == torch::kFloat16, "B must be float16");
    TORCH_CHECK(A.is_cuda(), "A must be on CUDA device");
    TORCH_CHECK(B.is_cuda(), "B must be on CUDA device");
    TORCH_CHECK(A.device() == B.device(), "A and B must be on the same device");
    
    // Check dimensions
    int M = A.size(0);
    int K_A = A.size(1);
    int N = B.size(0);
    int K_B = B.size(1);
    
    TORCH_CHECK(K_A == K_SIZE, "A must have K dimension = ", K_SIZE, ", got ", K_A);
    TORCH_CHECK(N == N_SIZE, "B must have N dimension = ", N_SIZE, ", got ", N);
    TORCH_CHECK(K_B == K_SIZE, "B must have K dimension = ", K_SIZE, ", got ", K_B);
    
    // Ensure tensors are contiguous
    A = A.contiguous();
    B = B.contiguous();
    
    // Create output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::kFloat16)
        .device(A.device());
    torch::Tensor C = torch::empty({M, N_SIZE}, options);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Get pointers to tensor data
    const half* A_ptr = reinterpret_cast<const half*>(A.data_ptr<at::Half>());
    const half* B_ptr = reinterpret_cast<const half*>(B.data_ptr<at::Half>());
    half* C_ptr = reinterpret_cast<half*>(C.data_ptr<at::Half>());
    
    // Launch kernel
    launch_gemm_kernel(A_ptr, B_ptr, C_ptr, M, stream);
    
    // PyTorch handles synchronization automatically when the tensor is accessed
    
    return C;
}

// Python binding
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "GEMM operation with N=5120, K=2048",
          py::arg("A"), py::arg("B"));
}