submission 863598

import os

import torch
import torch.utils.cpp_extension as cpp_ext
from task import input_t, output_t

# ---------------------------------------------------------------------------
# Two-sided tournament Jacobi for n=32.
# 16 warps per matrix, 6 sweeps, PAD=33 conflict-free layout.
# ---------------------------------------------------------------------------
_CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <torch/extension.h>

#define N32    32
#define PAD    33      // col stride: 33 ≡ 1 mod 32 → conflict-free both dims
#define NWARPS 16      // warps per CTA (= pairs per round)
#define NSWEEP  6      // 2-sided converges in 5-7 sweeps for cond=1

__global__ void __launch_bounds__(NWARPS * N32)
jacobi_16w_v2(
    const float* __restrict__ A,
    float*       __restrict__ Q,
    float*       __restrict__ L,
    int batch
) {
    int warp = threadIdx.x >> 5;     // warp index = pair index pi (0..15)
    int lane = threadIdx.x & 31;     // lane (0..31)
    int mat  = blockIdx.x;
    if (mat >= batch) return;

    __shared__ float sc[N32][PAD];   // A, column-major
    __shared__ float sv[N32][PAD];   // V, column-major
    __shared__ float lam[N32];
    __shared__ int   ord[N32];

    // Each warp loads 2 columns (16 × 2 = 32 columns total)
    const float* mA = A + (long)mat * N32 * N32;
    {
        int c0 = warp * 2;
        sc[c0    ][lane] = mA[(long)lane * N32 + c0    ];
        sc[c0 + 1][lane] = mA[(long)lane * N32 + c0 + 1];
        sv[c0    ][lane] = (c0     == lane) ? 1.0f : 0.0f;
        sv[c0 + 1][lane] = (c0 + 1 == lane) ? 1.0f : 0.0f;
    }
    __syncthreads();

    const unsigned WMASK = 0xffffffffu;

    for (int sweep = 0; sweep < NSWEEP; sweep++) {
        for (int r = 0; r < N32 - 1; r++) {
            int p, q;
            if (warp == 0) {
                p = r; q = N32 - 1;
            } else {
                int a = (r >= warp) ? r - warp : r - warp + (N32 - 1);
                int b = (r + warp < N32 - 1) ? r + warp : r + warp - (N32 - 1);
                p = (a < b) ? a : b;
                q = (a < b) ? b : a;
            }

            float ap = sc[lane][p];
            float aq = sc[lane][q];

            float a_pp = __shfl_sync(WMASK, ap, p);
            float a_qq = __shfl_sync(WMASK, aq, q);
            float a_pq = __shfl_sync(WMASK, aq, p);

            float c = 1.0f, s = 0.0f;
            if (a_pq != 0.0f) {
                float tau = (a_qq - a_pp) * (0.5f / a_pq);
                float t   = copysignf(1.0f, tau) / (fabsf(tau) + sqrtf(1.0f + tau * tau));
                c = rsqrtf(1.0f + t * t);
                s = c * t;
            }

            sc[lane][p] = c * ap - s * aq;
            sc[lane][q] = s * ap + c * aq;
            __syncthreads();

            float bp = sc[p][lane];
            float bq = sc[q][lane];
            sc[p][lane] = c * bp - s * bq;
            sc[q][lane] = s * bp + c * bq;

            float vp = sv[p][lane];
            float vq = sv[q][lane];
            sv[p][lane] = c * vp - s * vq;
            sv[q][lane] = s * vp + c * vq;
            __syncthreads();
        }
    }

    if (lane == 0) {
        int c0 = warp * 2;
        lam[c0    ] = sc[c0    ][c0    ];
        lam[c0 + 1] = sc[c0 + 1][c0 + 1];
        ord[c0    ] = c0;
        ord[c0 + 1] = c0 + 1;
    }
    __syncthreads();

    if (threadIdx.x == 0) {
        for (int i = 1; i < N32; i++) {
            int   ki = ord[i];
            float vi = lam[ki];
            int   j  = i;
            while (j > 0 && lam[ord[j - 1]] > vi) {
                ord[j] = ord[j - 1]; j--;
            }
            ord[j] = ki;
        }
        float tmp[N32];
        for (int i = 0; i < N32; i++) tmp[i] = lam[ord[i]];
        for (int i = 0; i < N32; i++) lam[i]  = tmp[i];
    }
    __syncthreads();

    float* mL = L + (long)mat * N32;
    float* mQ = Q + (long)mat * N32 * N32;
    {
        int c0 = warp * 2;
        if (lane == 0) {
            mL[c0    ] = lam[c0    ];
            mL[c0 + 1] = lam[c0 + 1];
        }
        mQ[(long)lane * N32 + c0    ] = sv[ord[c0    ]][lane];
        mQ[(long)lane * N32 + c0 + 1] = sv[ord[c0 + 1]][lane];
    }
}

void eigh_n32_v2(torch::Tensor A, torch::Tensor Q, torch::Tensor L) {
    int batch = A.size(0);
    dim3 block(NWARPS * N32);
    dim3 grid(batch);
    jacobi_16w_v2<<<grid, block>>>(
        A.data_ptr<float>(), Q.data_ptr<float>(), L.data_ptr<float>(), batch
    );
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("eigh_n32_v2", &eigh_n32_v2, "2-sided tournament Jacobi n=32 (16w, no-div, 6sw)");
}
"""

# ---------------------------------------------------------------------------
# cuSOLVER XsyevBatched with per-(n,batch) cached workspace.
#
# Per-call allocation: Ac = A.clone() (contiguous copy of A that cuSOLVER
# overwrites). After the first call per (n,batch), PyTorch's CUDA caching
# allocator returns the same memory block, so cudaMalloc is free. Caching
# Ac explicitly would require an explicit copy anyway (to preserve A_in),
# and returning a view of a static buffer is not safe across calls.
#
# Output convention: cuSOLVER writes Q in Fortran column-major; reading that
# memory as C row-major gives Q^T. We return Ac.transpose(-1,-2) — a zero-
# copy view — so the caller sees eigenvectors as columns (linalg.eigh convention).
#
# Fill mode: CUBLAS_FILL_MODE_UPPER is faster at n>=2048 on B200; LOWER wins
# for smaller n. Include fill mode in the cache key since buffer sizes differ.
# ---------------------------------------------------------------------------
_CPP_SRC = r"""
#include <torch/extension.h>
#include <cusolverDn.h>
#include <unordered_map>
#include <string>
#include <vector>

#define CUSOLVER_CHECK(x) TORCH_CHECK((x)==CUSOLVER_STATUS_SUCCESS, "cusolver error ", (int)(x))

struct WsEntry {
    torch::Tensor dbuf;
    std::vector<uint8_t> hbuf;
    cusolverDnParams_t params;
    size_t dsz;
    size_t hsz;
};

static cusolverDnHandle_t g_handle = nullptr;
static std::unordered_map<std::string, WsEntry> g_ws_cache;

static cusolverDnHandle_t get_handle() {
    if (!g_handle) {
        CUSOLVER_CHECK(cusolverDnCreate(&g_handle));
    }
    return g_handle;
}

// Returns {Ac [batch,n,n], vals [batch,n]}.
// Ac contains Q in Fortran column-major (= Q^T in C row-major).
// Caller does .transpose(-1,-2) for eigenvectors-as-columns convention.
// upper=true uses CUBLAS_FILL_MODE_UPPER, which is ~4% faster at n>=2048.
std::vector<torch::Tensor> eigh_batched(torch::Tensor A, bool upper) {
    TORCH_CHECK(A.is_cuda() && A.dtype() == torch::kFloat32 && A.dim() == 3);
    int64_t batch = A.size(0);
    int64_t n = A.size(1);
    auto handle = get_handle();

    // cuSOLVER overwrites A in place; work on a contiguous clone.
    // PyTorch's CUDA caching allocator reuses the same block after the first call.
    auto Ac = A.is_contiguous() ? A.clone() : A.contiguous();
    auto W = torch::empty({batch, n}, Ac.options());
    auto info = torch::empty({batch},
        torch::TensorOptions().dtype(torch::kInt32).device(A.device()));

    auto uplo = upper ? CUBLAS_FILL_MODE_UPPER : CUBLAS_FILL_MODE_LOWER;

    // Cache key includes fill mode since buffer sizes differ.
    std::string key = std::to_string(n) + "," + std::to_string(batch)
                      + "," + (upper ? "U" : "L");
    auto it = g_ws_cache.find(key);

    if (it == g_ws_cache.end()) {
        cusolverDnParams_t params;
        CUSOLVER_CHECK(cusolverDnCreateParams(&params));

        size_t dsz = 0, hsz = 0;
        CUSOLVER_CHECK(cusolverDnXsyevBatched_bufferSize(
            handle, params,
            CUSOLVER_EIG_MODE_VECTOR, uplo,
            n, CUDA_R_32F, Ac.data_ptr(), n,
            CUDA_R_32F, W.data_ptr(),
            CUDA_R_32F, &dsz, &hsz, batch));

        WsEntry entry;
        entry.params = params;
        entry.dsz = dsz;
        entry.hsz = hsz;
        entry.dbuf = torch::empty(
            {(int64_t)(dsz > 0 ? dsz : 1)},
            torch::TensorOptions().dtype(torch::kUInt8).device(A.device()));
        entry.hbuf.resize(hsz > 0 ? hsz : 1);
        g_ws_cache[key] = std::move(entry);
        it = g_ws_cache.find(key);
    }

    auto& ws = it->second;
    CUSOLVER_CHECK(cusolverDnXsyevBatched(
        handle, ws.params,
        CUSOLVER_EIG_MODE_VECTOR, uplo,
        n, CUDA_R_32F, Ac.data_ptr(), n,
        CUDA_R_32F, W.data_ptr(),
        CUDA_R_32F,
        ws.dsz > 0 ? ws.dbuf.data_ptr() : nullptr, ws.dsz,
        ws.hsz > 0 ? ws.hbuf.data() : nullptr, ws.hsz,
        info.data_ptr<int>(), batch));
    return {Ac, W};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("eigh_batched", &eigh_batched, "cuSOLVER XsyevBatched with cached workspace");
}
"""

_ext_n32 = None
_ext_cusolver = None


def _get_ext_n32():
    global _ext_n32
    if _ext_n32 is None:
        _ext_n32 = cpp_ext.load_inline(
            name="jacobi_16w_v2",
            cpp_sources="",
            cuda_sources=_CUDA_SRC,
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=False,
        )
    return _ext_n32


def _get_ext_cusolver():
    global _ext_cusolver
    if _ext_cusolver is None:
        cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda")
        _ext_cusolver = cpp_ext.load_inline(
            name="eigh_batched_cached",
            cpp_sources=_CPP_SRC,
            extra_include_paths=[os.path.join(cuda_home, "include")],
            extra_ldflags=[
                f"-L{os.path.join(cuda_home, 'lib64')}",
                "-lcusolver",
            ],
            verbose=False,
        )
    return _ext_cusolver


def _eigh_n32(A: torch.Tensor):
    batch, n, _ = A.shape
    assert n == 32
    A_c = A.contiguous()
    Q = torch.empty_like(A_c)
    L = torch.empty(batch, n, device=A.device, dtype=A.dtype)
    _get_ext_n32().eigh_n32_v2(A_c, Q, L)
    return Q, L


def custom_kernel(data: input_t) -> output_t:
    n = data.shape[-1]
    if n == 32:
        return _eigh_n32(data)
    # CUBLAS_FILL_MODE_UPPER is ~4% faster for the batched reduction at n>=2048.
    upper = n >= 2048
    vecs_t, values = _get_ext_cusolver().eigh_batched(data, upper)
    # cuSOLVER writes column-major eigvecs; read as row-major gives V^T.
    # Transpose back to get Q with eigenvectors as columns (linalg.eigh convention).
    return vecs_t.transpose(-1, -2), values