submission 873092

#!POPCORN leaderboard eigh
#!POPCORN gpu B200

"""
GPUMode B200 batched symmetric eigendecomposition scaffold, v6.

n == 32:
    Custom cyclic Jacobi CUDA kernel, one block per matrix.

Other sizes:
    torch.linalg.eigh fallback.

Important build detail:
    The CUDA translation unit contains no PyTorch headers. PyTorch's Python
    3.13 / CUDA 13 runner currently exposes a template-parsing incompatibility
    when NVCC directly parses torch/extension.h, so Tensor handling stays in
    the C++ translation unit and the CUDA unit exposes a plain C launcher.
"""

from __future__ import annotations

import os
import tempfile
from pathlib import Path

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


_CPP_SRC = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>

#include <cstdint>
#include <vector>


extern "C" void launch_jacobi_eigh_32(
    const float* input,
    float* vectors,
    float* values,
    int64_t batch
);


std::vector<torch::Tensor> eigh32_cuda(
    torch::Tensor input
) {
    constexpr int64_t N = 32;

    TORCH_CHECK(input.is_cuda(), "input must be CUDA");
    TORCH_CHECK(
        input.scalar_type() == torch::kFloat32,
        "input must be float32"
    );
    TORCH_CHECK(
        input.dim() == 3 &&
        input.size(1) == N &&
        input.size(2) == N,
        "eigh32_cuda expects [batch, 32, 32]"
    );
    TORCH_CHECK(
        input.is_contiguous(),
        "input must be contiguous"
    );

    c10::cuda::CUDAGuard device_guard(input.device());

    const int64_t batch = input.size(0);

    torch::Tensor vectors = torch::empty_like(input);
    torch::Tensor values = torch::empty(
        {batch, N},
        input.options()
    );

    launch_jacobi_eigh_32(
        input.data_ptr<float>(),
        vectors.data_ptr<float>(),
        values.data_ptr<float>(),
        batch
    );

    return {vectors, values};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
    module.def(
        "eigh32_cuda",
        &eigh32_cuda,
        "Batched 32x32 symmetric Jacobi eigendecomposition"
    );
}

"""


_CUDA_SRC = r"""
#include <cuda_runtime.h>

#include <cmath>
#include <cstdint>


namespace {

constexpr int N = 32;
constexpr int NN = N * N;
constexpr int PAIRS = N / 2;
constexpr int ROUNDS = N - 1;
constexpr int SWEEPS = 14;


__global__ void jacobi_eigh_32_kernel(
    const float* __restrict__ input,
    float* __restrict__ vectors,
    float* __restrict__ values
) {
    const int matrix = static_cast<int>(blockIdx.x);
    const int tid = static_cast<int>(threadIdx.x);

    __shared__ float a0[NN];
    __shared__ float a1[NN];
    __shared__ float q0[NN];
    __shared__ float q1[NN];

    __shared__ int circle[N];
    __shared__ int partner[N];
    __shared__ int order[N];

    // For output column j:
    // J[j, j] = alpha[j]
    // J[partner[j], j] = beta[j]
    __shared__ float alpha[N];
    __shared__ float beta[N];

    const long long base =
        static_cast<long long>(matrix) * NN;

    if (tid < NN) {
        a0[tid] = input[base + tid];

        const int row = tid / N;
        const int col = tid - row * N;

        q0[tid] = row == col ? 1.0f : 0.0f;
    }

    if (tid < N) {
        circle[tid] = tid;
        partner[tid] = tid;
        alpha[tid] = 1.0f;
        beta[tid] = 0.0f;
        order[tid] = tid;
    }

    __syncthreads();

    bool parity = false;

    #pragma unroll 1
    for (int sweep = 0; sweep < SWEEPS; ++sweep) {
        if (tid < N) {
            circle[tid] = tid;
        }

        __syncthreads();

        #pragma unroll 1
        for (int round = 0; round < ROUNDS; ++round) {
            float* cur_a = parity ? a1 : a0;
            float* next_a = parity ? a0 : a1;

            float* cur_q = parity ? q1 : q0;
            float* next_q = parity ? q0 : q1;

            // Build 16 disjoint plane rotations.
            if (tid < PAIRS) {
                const int p = circle[tid];
                const int q = circle[N - 1 - tid];

                const float app = cur_a[p * N + p];
                const float aqq = cur_a[q * N + q];
                const float apq = cur_a[p * N + q];

                float c = 1.0f;
                float s = 0.0f;

                const float scale =
                    fabsf(app) + fabsf(aqq) + 1.0f;

                if (fabsf(apq) > 2.0e-7f * scale) {
                    const float tau =
                        (aqq - app) / (2.0f * apq);

                    const float t =
                        copysignf(1.0f, tau) /
                        (
                            fabsf(tau) +
                            hypotf(1.0f, tau)
                        );

                    c = rsqrtf(1.0f + t * t);
                    s = t * c;
                }

                partner[p] = q;
                partner[q] = p;

                alpha[p] = c;
                alpha[q] = c;

                // Q[:, p] = c*Q[:, p] - s*Q[:, q]
                // Q[:, q] = s*Q[:, p] + c*Q[:, q]
                beta[p] = -s;
                beta[q] = s;
            }

            __syncthreads();

            if (tid < NN) {
                const int i = tid / N;
                const int j = tid - i * N;

                const int pi = partner[i];
                const int pj = partner[j];

                const float ai = alpha[i];
                const float bi = beta[i];
                const float aj = alpha[j];
                const float bj = beta[j];

                // next_a = J^T * cur_a * J
                next_a[i * N + j] =
                    ai * aj * cur_a[i * N + j] +
                    bi * aj * cur_a[pi * N + j] +
                    ai * bj * cur_a[i * N + pj] +
                    bi * bj * cur_a[pi * N + pj];

                // next_q = cur_q * J
                next_q[i * N + j] =
                    aj * cur_q[i * N + j] +
                    bj * cur_q[i * N + pj];
            }

            __syncthreads();

            parity = !parity;

            // Circle-method schedule:
            // keep position 0 fixed and rotate positions 1..31.
            if (tid == 0) {
                const int last = circle[N - 1];

                for (int k = N - 1; k > 1; --k) {
                    circle[k] = circle[k - 1];
                }

                circle[1] = last;
            }

            __syncthreads();
        }
    }

    float* final_a = parity ? a1 : a0;
    float* final_q = parity ? q1 : q0;

    if (tid == 0) {
        for (int i = 0; i < N; ++i) {
            order[i] = i;
        }

        // Sort eigenvalues ascending.
        for (int i = 0; i < N - 1; ++i) {
            int best = i;

            float best_value =
                final_a[order[i] * N + order[i]];

            for (int j = i + 1; j < N; ++j) {
                const float candidate =
                    final_a[order[j] * N + order[j]];

                if (candidate < best_value) {
                    best = j;
                    best_value = candidate;
                }
            }

            const int tmp = order[i];
            order[i] = order[best];
            order[best] = tmp;
        }

        for (int j = 0; j < N; ++j) {
            const int src = order[j];

            values[
                static_cast<long long>(matrix) * N + j
            ] = final_a[src * N + src];
        }
    }

    __syncthreads();

    if (tid < NN) {
        const int row = tid / N;
        const int out_col = tid - row * N;
        const int src_col = order[out_col];

        vectors[base + tid] =
            final_q[row * N + src_col];
    }
}

} // namespace


extern "C" void launch_jacobi_eigh_32(
    const float* input,
    float* vectors,
    float* values,
    int64_t batch
) {
    jacobi_eigh_32_kernel<<<
        static_cast<unsigned int>(batch),
        NN
    >>>(
        input,
        vectors,
        values
    );
}
"""


def _write_source_file(path: Path, source: str) -> None:
    try:
        if path.exists() and path.read_text(encoding="utf-8") == source:
            return
    except OSError:
        pass

    temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")
    temporary.write_text(source, encoding="utf-8")
    os.replace(temporary, path)


_SOURCE_DIR = (
    Path(tempfile.gettempdir()) /
    "gpumode_eigh_jacobi32_v6_sources"
)
_SOURCE_DIR.mkdir(parents=True, exist_ok=True)

_CPP_PATH = _SOURCE_DIR / "eigh_wrapper.cpp"
_CUDA_PATH = _SOURCE_DIR / "jacobi32.cu"

_write_source_file(_CPP_PATH, _CPP_SRC)
_write_source_file(_CUDA_PATH, _CUDA_SRC)


_mod = load(
    name="gpumode_eigh_jacobi32_v6",
    sources=[
        str(_CPP_PATH),
        str(_CUDA_PATH),
    ],
    extra_cflags=[
        "-O3",
        "-std=c++17",
    ],
    extra_cuda_cflags=[
        "-O3",
        "-std=c++17",
        "--use_fast_math",
    ],
    with_cuda=True,
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    if data.shape[-1] == 32:
        q, eigenvalues = _mod.eigh32_cuda(data)
        return q, eigenvalues

    eigenvalues, q = torch.linalg.eigh(data)
    return q, eigenvalues