claude-opus-4-1 / cuda1970e7

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include "kernel.h"
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // Input validation
    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 a CUDA tensor");
    TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor");
    TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
    TORCH_CHECK(B.is_contiguous(), "B must be contiguous");
    
    // Dimension validation
    TORCH_CHECK(A.dim() == 2, "A must be 2D");
    TORCH_CHECK(B.dim() == 2, "B must be 2D");
    
    const int64_t M = A.size(0);
    const int64_t K_A = A.size(1);
    const int64_t N = B.size(0);
    const int64_t K_B = B.size(1);
    
    TORCH_CHECK(K_A == 4096, "A's K dimension must be 4096");
    TORCH_CHECK(N == 4096, "B's N dimension must be 4096");
    TORCH_CHECK(K_B == 4096, "B's K dimension must be 4096");
    
    // Set the CUDA device
    c10::cuda::CUDAGuard device_guard(A.device());
    
    // Create output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::kFloat16)
        .device(A.device())
        .requires_grad(false);
    torch::Tensor C = torch::empty({M, N}, options);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Get raw pointers
    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, static_cast<int>(M), stream);
    
    // Check for errors
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        TORCH_CHECK(false, "CUDA kernel launch failed: ", cudaGetErrorString(err));
    }
    
    return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Optimized GEMM for N=4096, K=4096",
          py::arg("A"), py::arg("B"));
}