submission 855334

#!POPCORN leaderboard eigh
#!POPCORN gpu B200

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

_CPP_SRC = r"""
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> syevj_small_batched(torch::Tensor input);
std::vector<torch::Tensor> xsyev_batched(torch::Tensor input);
"""

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cusolverDn.h>
#include <climits>
#include <vector>

static cusolverDnHandle_t eigh_handle() {
    static cusolverDnHandle_t handle = nullptr;
    if (handle == nullptr) {
        cusolverStatus_t status = cusolverDnCreate(&handle);
        TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnCreate failed");
    }
    return handle;
}

static cusolverDnParams_t eigh_params() {
    static cusolverDnParams_t params = nullptr;
    if (params == nullptr) {
        cusolverStatus_t status = cusolverDnCreateParams(&params);
        TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnCreateParams failed");
    }
    return params;
}

static syevjInfo_t jacobi_params() {
    static syevjInfo_t params = nullptr;
    if (params == nullptr) {
        cusolverStatus_t status = cusolverDnCreateSyevjInfo(&params);
        TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnCreateSyevjInfo failed");
        status = cusolverDnXsyevjSetTolerance(params, 1.0e-5);
        TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnXsyevjSetTolerance failed");
        status = cusolverDnXsyevjSetMaxSweeps(params, 10);
        TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnXsyevjSetMaxSweeps failed");
    }
    return params;
}

static void check_cuda_tail(const char* tag) {
    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, tag, ": ", cudaGetErrorString(err));
}

static void check_square_cuda_fp32(torch::Tensor input) {
    TORCH_CHECK(input.is_cuda(), "input must be CUDA");
    TORCH_CHECK(input.dtype() == torch::kFloat32, "input must be float32");
    TORCH_CHECK(input.dim() == 3 && input.size(1) == input.size(2),
                "input must have shape (batch, n, n)");
    TORCH_CHECK(input.size(0) <= INT_MAX && input.size(1) <= INT_MAX,
                "dimensions exceed int range");
}

std::vector<torch::Tensor> syevj_small_batched(torch::Tensor input) {
    check_square_cuda_fp32(input);
    const int64_t batch64 = input.size(0);
    const int64_t n64 = input.size(1);
    TORCH_CHECK(n64 >= 1 && n64 <= 32, "syevj batched supports 1 <= n <= 32 here");
    const int batch = static_cast<int>(batch64);
    const int n = static_cast<int>(n64);

    auto q = torch::empty_strided({batch64, n64, n64}, {n64 * n64, 1, n64}, input.options());
    q.copy_(input);
    auto values = torch::empty({batch64, n64}, input.options());
    auto info = torch::empty({batch64}, torch::TensorOptions().device(input.device()).dtype(torch::kInt32));

    int lwork = 0;
    cusolverStatus_t status = cusolverDnSsyevjBatched_bufferSize(
        eigh_handle(), CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, n,
        q.data_ptr<float>(), n, values.data_ptr<float>(), &lwork, jacobi_params(), batch);
    TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnSsyevjBatched_bufferSize failed");

    auto work = torch::empty({static_cast<int64_t>(lwork)}, input.options());
    status = cusolverDnSsyevjBatched(
        eigh_handle(), CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, n,
        q.data_ptr<float>(), n, values.data_ptr<float>(), work.data_ptr<float>(), lwork,
        info.data_ptr<int>(), jacobi_params(), batch);
    TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnSsyevjBatched failed");
    check_cuda_tail("syevj_small_batched cuda tail failed");
    return {q, values};
}

std::vector<torch::Tensor> xsyev_batched(torch::Tensor input) {
    check_square_cuda_fp32(input);
    const int64_t batch64 = input.size(0);
    const int64_t n64 = input.size(1);
    TORCH_CHECK(n64 >= 1, "xsyev batched requires n >= 1");
    const int batch = static_cast<int>(batch64);
    const int n = static_cast<int>(n64);

    auto q = torch::empty_strided({batch64, n64, n64}, {n64 * n64, 1, n64}, input.options());
    q.copy_(input);
    auto values = torch::empty({batch64, n64}, input.options());
    auto info = torch::empty({batch64}, torch::TensorOptions().device(input.device()).dtype(torch::kInt32));

    size_t work_device_size = 0;
    size_t work_host_size = 0;
    cusolverStatus_t status = cusolverDnXsyevBatched_bufferSize(
        eigh_handle(), eigh_params(), CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, n,
        CUDA_R_32F, q.data_ptr<float>(), n, CUDA_R_32F, values.data_ptr<float>(), CUDA_R_32F,
        &work_device_size, &work_host_size, batch);
    TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnXsyevBatched_bufferSize failed");

    auto work_device = torch::empty(
        {static_cast<int64_t>(work_device_size)},
        torch::TensorOptions().device(input.device()).dtype(torch::kUInt8));
    std::vector<unsigned char> work_host(work_host_size);

    status = cusolverDnXsyevBatched(
        eigh_handle(), eigh_params(), CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, n,
        CUDA_R_32F, q.data_ptr<float>(), n, CUDA_R_32F, values.data_ptr<float>(), CUDA_R_32F,
        work_device.data_ptr(), work_device_size, work_host_size ? work_host.data() : nullptr,
        work_host_size, info.data_ptr<int>(), batch);
    TORCH_CHECK(status == CUSOLVER_STATUS_SUCCESS, "cusolverDnXsyevBatched failed");
    check_cuda_tail("xsyev_batched cuda tail failed");
    return {q, values};
}
"""

_NATIVE = load_inline(
    name="eigh_smalln_eligible_loose10_v1",
    cpp_sources=[_CPP_SRC],
    cuda_sources=[_CUDA_SRC],
    functions=["syevj_small_batched", "xsyev_batched"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcusolver"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    if data.is_cuda and data.dtype == torch.float32 and data.dim() == 3 and data.shape[-2] == data.shape[-1]:
        n = data.shape[-1]
        if n <= 32:
            q, values = _NATIVE.syevj_small_batched(data)
            return q, values
        if n >= 64:
            q, values = _NATIVE.xsyev_batched(data)
            return q, values
    values, vectors = torch.linalg.eigh(data)
    return vectors, values