claude-opus-4-1-20250805 / cuda6fecc6

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

// Helper function to check tensor properties
void check_cuda_tensor(const torch::Tensor& tensor, const std::string& name) {
    if (!tensor.is_cuda()) {
        throw std::runtime_error(name + " must be a CUDA tensor");
    }
    if (!tensor.is_contiguous()) {
        throw std::runtime_error(name + " must be contiguous");
    }
}

// Main run function
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // Input validation
    check_cuda_tensor(A, "A");
    check_cuda_tensor(B, "B");
    
    // Check dtypes - handle both Half and Float16
    if (A.scalar_type() != torch::ScalarType::Half) {
        throw std::runtime_error("A must be float16");
    }
    if (B.scalar_type() != torch::ScalarType::Half) {
        throw std::runtime_error("B must be float16");
    }
    
    // Check dimensions
    if (A.dim() != 2) {
        throw std::runtime_error("A must be 2-dimensional");
    }
    if (B.dim() != 2) {
        throw std::runtime_error("B must be 2-dimensional");
    }
    
    int M = A.size(0);
    int K_A = A.size(1);
    int N = B.size(0);
    int K_B = B.size(1);
    
    // Verify dimensions match specification
    if (N != N_FIXED) {
        throw std::runtime_error("B dimension 0 must be " + std::to_string(N_FIXED) + 
                                ", got " + std::to_string(N));
    }
    if (K_A != K_FIXED) {
        throw std::runtime_error("A dimension 1 must be " + std::to_string(K_FIXED) + 
                                ", got " + std::to_string(K_A));
    }
    if (K_B != K_FIXED) {
        throw std::runtime_error("B dimension 1 must be " + std::to_string(K_FIXED) + 
                                ", got " + std::to_string(K_B));
    }
    
    // Allocate output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::ScalarType::Half)
        .device(A.device())
        .requires_grad(false);
    torch::Tensor C = torch::empty({M, N_FIXED}, options);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Get raw pointers - cast to half*
    const half* A_ptr = reinterpret_cast<const half*>(A.data_ptr<torch::Half>());
    const half* B_ptr = reinterpret_cast<const half*>(B.data_ptr<torch::Half>());
    half* C_ptr = reinterpret_cast<half*>(C.data_ptr<torch::Half>());
    
    // Launch kernel
    launch_gemm_kernel(A_ptr, B_ptr, C_ptr, M, stream);
    
    // Ensure kernel completion for correctness
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(std::string("CUDA kernel launch failed: ") + cudaGetErrorString(err));
    }
    
    return C;
}

// Python binding
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Optimized GEMM kernel for N=256, K=7168 (C = A @ B.T)",
          py::arg("A"), py::arg("B"));
}