submission 843903

#!POPCORN leaderboard eigh
#!POPCORN gpu B200

import os
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

DBG = "DBG" in os.environ

CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cusolverDn.h>
#include <vector>
#include <mutex>

#define CUDA_CHECK(call) do {                                                      \
    cudaError_t err = call;                                                        \
    if (err != cudaSuccess) {                                                      \
        throw std::runtime_error(std::string("CUDA error: ") +                    \
                                 cudaGetErrorString(err));                         \
    }                                                                              \
} while(0)

#define CUSOLVER_CHECK(call) do {                                                  \
    cusolverStatus_t err = call;                                                   \
    if (err != CUSOLVER_STATUS_SUCCESS) {                                          \
        throw std::runtime_error("cuSOLVER error: " +                             \
                                 std::to_string((int)err));                        \
    }                                                                              \
} while(0)

extern "C" {
cusolverStatus_t CUSOLVERAPI cusolverDnXsyevBatched_bufferSize(
    cusolverDnHandle_t, cusolverDnParams_t, cusolverEigMode_t, cublasFillMode_t,
    int64_t, cudaDataType, const void*, int64_t, cudaDataType, const void*,
    cudaDataType, size_t*, size_t*, int64_t);

cusolverStatus_t CUSOLVERAPI cusolverDnXsyevBatched(
    cusolverDnHandle_t, cusolverDnParams_t, cusolverEigMode_t, cublasFillMode_t,
    int64_t, cudaDataType, void*, int64_t, cudaDataType, void*,
    cudaDataType, void*, size_t, void*, size_t, int*, int64_t);
}

static cusolverDnHandle_t g_handle = nullptr;
static cusolverDnParams_t g_params = nullptr;
static std::mutex g_mutex;

static void ensure_handle() {
    std::lock_guard<std::mutex> lock(g_mutex);
    if (!g_handle) {
        CUSOLVER_CHECK(cusolverDnCreate(&g_handle));
        CUSOLVER_CHECK(cusolverDnCreateParams(&g_params));
    }
}

template <typename scalar_t>
cudaDataType cuda_data_type();

template <>
cudaDataType cuda_data_type<double>() { return CUDA_R_64F; }

template <>
cudaDataType cuda_data_type<float>() { return CUDA_R_32F; }

static void* g_dev_work = nullptr;
static size_t g_dev_work_bytes = 0;
static void* g_host_work = nullptr;
static size_t g_host_work_bytes = 0;
static int* g_dev_info = nullptr;
static int g_dev_info_cap = 0;

static void ensure_workspace(size_t dev_bytes, size_t host_bytes, int batch) {
    if (dev_bytes > g_dev_work_bytes) {
        if (g_dev_work) cudaFree(g_dev_work);
        CUDA_CHECK(cudaMalloc(&g_dev_work, dev_bytes));
        g_dev_work_bytes = dev_bytes;
    }
    if (host_bytes > g_host_work_bytes) {
        free(g_host_work);
        g_host_work = malloc(host_bytes);
        g_host_work_bytes = host_bytes;
    }
    if (batch > g_dev_info_cap) {
        if (g_dev_info) cudaFree(g_dev_info);
        CUDA_CHECK(cudaMalloc(&g_dev_info, sizeof(int) * batch));
        g_dev_info_cap = batch;
    }
}

template <typename scalar_t>
void batched_syevd_kernel(cusolverDnHandle_t handle, cusolverDnParams_t params,
                           cusolverEigMode_t jobz, cublasFillMode_t uplo,
                           int64_t n, scalar_t* d_A, int64_t lda,
                           scalar_t* d_W, int64_t batchSize) {
    cudaDataType dt = cuda_data_type<scalar_t>();

    size_t dev_bytes = 0, host_bytes = 0;
    CUSOLVER_CHECK(cusolverDnXsyevBatched_bufferSize(
        handle, params, jobz, uplo, n, dt, d_A, lda, dt, d_W,
        dt, &dev_bytes, &host_bytes, batchSize));

    ensure_workspace(dev_bytes, host_bytes, batchSize);

    CUSOLVER_CHECK(cusolverDnXsyevBatched(
        handle, params, jobz, uplo, n, dt, d_A, lda, dt, d_W,
        dt, g_dev_work, dev_bytes, g_host_work, host_bytes,
        g_dev_info, batchSize));

    std::vector<int> h_info(batchSize);
    CUDA_CHECK(cudaMemcpy(h_info.data(), g_dev_info, sizeof(int) * batchSize,
                          cudaMemcpyDeviceToHost));
    for (int64_t i = 0; i < batchSize; i++) {
        if (h_info[i] != 0) {
            throw std::runtime_error("syevd batched failed at batch " +
                                     std::to_string(i) +
                                     " with info=" + std::to_string(h_info[i]));
        }
    }
}

std::vector<torch::Tensor> syevd_cuda(torch::Tensor A, bool compute_eigenvectors) {
    TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor");
    TORCH_CHECK(A.dim() == 2 || A.dim() == 3, "A must be 2 or 3-dimensional");
    TORCH_CHECK(A.size(-1) == A.size(-2), "A must be square");
    TORCH_CHECK(A.scalar_type() == torch::kFloat64 ||
                A.scalar_type() == torch::kFloat32,
                "A must be float32 or float64");

    auto A_contig = A.contiguous();
    int64_t n = A_contig.size(-1);
    int64_t lda = n;
    int64_t batch = A_contig.dim() == 2 ? 1 : A_contig.size(0);

    ensure_handle();

    torch::Tensor eigenvectors = torch::empty_like(A_contig);
    eigenvectors.copy_(A_contig);
    torch::Tensor eigenvalues = torch::empty({batch, n}, A_contig.options());

    cusolverEigMode_t jobz = compute_eigenvectors
        ? CUSOLVER_EIG_MODE_VECTOR
        : CUSOLVER_EIG_MODE_NOVECTOR;
    cublasFillMode_t uplo = CUBLAS_FILL_MODE_LOWER;

    auto dtype = A_contig.scalar_type();

    AT_DISPATCH_FLOATING_TYPES(dtype, "syevd_cuda", [&] {
        batched_syevd_kernel<scalar_t>(
            g_handle, g_params, jobz, uplo, n,
            eigenvectors.data_ptr<scalar_t>(), lda,
            eigenvalues.data_ptr<scalar_t>(), batch);
    });

    return {eigenvectors.transpose(1, 2).contiguous(), eigenvalues};
}
"""

CPP_SRC = """
std::vector<torch::Tensor> syevd_cuda(torch::Tensor A, bool compute_eigenvectors);
"""

_syevd_module = load_inline(
    name="syevd_module",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=["syevd_cuda"],
    verbose=True,
    extra_ldflags=["-lcusolver", "-lcublas"],
)


def custom_kernel(data: input_t) -> output_t:
    result = _syevd_module.syevd_cuda(data, True)
    return (result[0], result[1])

if DBG:
    def custom_kernel(data: input_t) -> output_t:
        return torch.linalg.eigh(data)[::-1]


if DBG and __name__ == "__main__":
    from test_eigh import run_test_suite, solve_ref

    torch.set_default_dtype(torch.float32)
    print("\n=== cuSOLVER syevd kernel tests ===")

    def test_wrapper(unbatched_data):
        return custom_kernel(unbatched_data)

    run_test_suite(solver=test_wrapper, sizes=[64, 128], conds=[5, 10], rtol=5e-4)
    print("\nAll tests completed.")